Implement __setitem__, __getitem__ and __delitem__

loro.pyi

@@ -185,6 +185,14 @@ class ImportStatus:
 class LoroCounter:
     id: ContainerID
     value: float
+    def __float__(self) -> float: ...
+    def __int__(self) -> int: ...
+    def __add__(self, other: typing.Any) -> float: ...
+    def __radd__(self, other: typing.Any) -> float: ...
+    def __sub__(self, other: typing.Any) -> float: ...
+    def __rsub__(self, other: typing.Any) -> float: ...
+    def __neg__(self) -> float: ...
+    def __abs__(self) -> float: ...
     def __new__(
         cls,
     ): ...
@@ -1051,6 +1059,24 @@ class LoroList:
         """
         ...
 
+    @typing.overload
+    def __getitem__(self, index: int) -> ValueOrContainer: ...
+
+    @typing.overload
+    def __getitem__(self, index: slice) -> list[ValueOrContainer]: ...
+
+    @typing.overload
+    def __setitem__(self, index: int, value: LoroValue) -> None: ...
+
+    @typing.overload
+    def __setitem__(self, index: slice, value: typing.Iterable[LoroValue]) -> None: ...
+
+    @typing.overload
+    def __delitem__(self, index: int) -> None: ...
+
+    @typing.overload
+    def __delitem__(self, index: slice) -> None: ...
+
     def insert_container(self, pos: int, child: Container) -> Container:
         r"""
         Insert a container with the given type at the given index.
@@ -1193,6 +1219,15 @@ class LoroMap:
         """
         ...
 
+    def __contains__(self, key: str) -> bool: ...
+
+    @typing.overload
+    def __getitem__(self, key: str) -> ValueOrContainer: ...
+
+    def __setitem__(self, key: str, value: LoroValue) -> None: ...
+
+    def __delitem__(self, key: str) -> None: ...
+
     def is_empty(self) -> bool:
         r"""
         Whether the map is empty.
@@ -1264,6 +1299,8 @@ class LoroMap:
         """
         ...
 
+    def items(self) -> list[tuple[str, ValueOrContainer]]: ...
+
     def get_last_editor(self, key: str) -> typing.Optional[int]:
         r"""
         Get the peer id of the last editor on the given entry
@@ -1468,6 +1505,18 @@ class LoroMovableList:
         """
         ...
 
+    @typing.overload
+    def __setitem__(self, index: int, value: LoroValue) -> None: ...
+
+    @typing.overload
+    def __setitem__(self, index: slice, value: typing.Iterable[LoroValue]) -> None: ...
+
+    @typing.overload
+    def __delitem__(self, index: int) -> None: ...
+
+    @typing.overload
+    def __delitem__(self, index: slice) -> None: ...
+
     def get_creator_at(self, pos: int) -> typing.Optional[int]:
         r"""
         Get the creator of the list item at the given position.
@@ -1514,6 +1563,16 @@ class LoroText:
     len_utf8: int
     len_unicode: int
     len_utf16: int
+    def __str__(self) -> str: ...
+    def __repr__(self) -> str: ...
+    def __len__(self) -> int: ...
+    @typing.overload
+    def __getitem__(self, index: int) -> str: ...
+    @typing.overload
+    def __getitem__(self, index: slice) -> str: ...
+    def __contains__(self, item: typing.Any) -> bool: ...
+    def __add__(self, other: typing.Union[str, "LoroText"]) -> str: ...
+    def __radd__(self, other: typing.Union[str, "LoroText"]) -> str: ...
     def __new__(
         cls,
     ): ...
@@ -1728,6 +1787,7 @@ class LoroTree:
     is_attached: bool
     roots: list[TreeID]
     id: ContainerID
+    def __contains__(self, target: TreeID) -> bool: ...
     def __new__(
         cls,
     ): ...

src/container.rs

@@ -7,6 +7,7 @@ mod movable_list;
 mod text;
 mod tree;
 mod unknown;
+pub mod utils;
 pub use counter::LoroCounter;
 pub use list::LoroList;
 pub use map::LoroMap;

src/container/counter.rs

@@ -7,7 +7,19 @@ use crate::{
     value::ContainerID,
 };
 use loro::{ContainerTrait, LoroCounter as LoroCounterInner};
-use pyo3::prelude::*;
+use pyo3::{exceptions::PyTypeError, prelude::*, Bound, PyRef};
+
+impl LoroCounter {
+    fn coerce_to_f64(other: &Bound<'_, PyAny>) -> PyLoroResult<f64> {
+        if let Ok(value) = other.extract::<f64>() {
+            Ok(value)
+        } else if let Ok(counter) = other.extract::<PyRef<LoroCounter>>() {
+            Ok(counter.get_value())
+        } else {
+            Err(PyTypeError::new_err("expected a number or LoroCounter").into())
+        }
+    }
+}
 
 #[pyclass(frozen)]
 #[derive(Debug, Clone, Default)]
@@ -46,6 +58,48 @@ impl LoroCounter {
         self.0.get_value()
     }
 
+    pub fn __float__(&self) -> f64 {
+        self.0.get_value()
+    }
+
+    pub fn __int__(&self) -> PyLoroResult<i64> {
+        let value = self.0.get_value();
+        if !value.is_finite() {
+            return Err(PyTypeError::new_err("cannot convert non-finite counter to int").into());
+        }
+        if value < i64::MIN as f64 || value > i64::MAX as f64 {
+            return Err(PyTypeError::new_err("counter value out of range for int").into());
+        }
+        Ok(value.trunc() as i64)
+    }
+
+    pub fn __add__(&self, other: Bound<'_, PyAny>) -> PyLoroResult<f64> {
+        let delta = Self::coerce_to_f64(&other)?;
+        Ok(self.0.get_value() + delta)
+    }
+
+    pub fn __radd__(&self, other: Bound<'_, PyAny>) -> PyLoroResult<f64> {
+        self.__add__(other)
+    }
+
+    pub fn __sub__(&self, other: Bound<'_, PyAny>) -> PyLoroResult<f64> {
+        let delta = Self::coerce_to_f64(&other)?;
+        Ok(self.0.get_value() - delta)
+    }
+
+    pub fn __rsub__(&self, other: Bound<'_, PyAny>) -> PyLoroResult<f64> {
+        let delta = Self::coerce_to_f64(&other)?;
+        Ok(delta - self.0.get_value())
+    }
+
+    pub fn __neg__(&self) -> f64 {
+        -self.0.get_value()
+    }
+
+    pub fn __abs__(&self) -> f64 {
+        self.0.get_value().abs()
+    }
+
     pub fn doc(&self) -> Option<LoroDoc> {
         self.0.doc().map(|doc| doc.into())
     }

src/container/list.rs

@@ -2,17 +2,22 @@ use std::sync::Arc;
 
 use loro::{ContainerTrait, LoroList as LoroListInner};
 use pyo3::prelude::*;
+use pyo3::{
+    exceptions::{PyIndexError, PyValueError},
+    BoundObject,
+};
 
+use crate::container::utils::{py_any_to_loro_values, slice_indices_positions, SliceOrInt};
 use crate::{
     doc::LoroDoc,
-    err::PyLoroResult,
+    err::{PyLoroError, PyLoroResult},
     event::{DiffEvent, Subscription},
     value::{ContainerID, LoroValue, ValueOrContainer, ID},
 };
 
 use super::{Container, Cursor, Side};
 
-#[pyclass(frozen)]
+#[pyclass(frozen, sequence)]
 #[derive(Debug, Clone, Default)]
 pub struct LoroList(pub LoroListInner);
 
@@ -106,6 +111,114 @@ impl LoroList {
         })
     }
 
+    pub fn __getitem__<'py>(
+        &self,
+        py: Python<'py>,
+        index: SliceOrInt<'py>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        match index {
+            SliceOrInt::Slice(slice) => {
+                let indices = slice.indices(self.0.len() as isize)?;
+                let mut i = indices.start;
+                let mut list: Vec<ValueOrContainer> = Vec::with_capacity(indices.slicelength);
+
+                for _ in 0..indices.slicelength {
+                    list.push(
+                        self.0
+                            .get(i as usize)
+                            .ok_or(PyIndexError::new_err("index out of range"))?
+                            .into(),
+                    );
+                    i += indices.step;
+                }
+                list.into_pyobject(py)
+            }
+            SliceOrInt::Int(idx) => {
+                let value: ValueOrContainer = self
+                    .0
+                    .get(usize::try_from(idx)?)
+                    .ok_or(PyIndexError::new_err("index out of range"))?
+                    .into();
+                Ok(value.into_pyobject(py)?.into_any().into_bound())
+            }
+        }
+    }
+
+    pub fn __setitem__<'py>(
+        &self,
+        index: SliceOrInt<'py>,
+        value: Bound<'py, PyAny>,
+    ) -> PyLoroResult<()> {
+        match index {
+            SliceOrInt::Int(idx) => {
+                let extracted: LoroValue = value.extract().map_err(PyLoroError::from)?;
+                self.0.delete(idx, 1).map_err(PyLoroError::from)?;
+                self.0.insert(idx, extracted.0).map_err(PyLoroError::from)?;
+                Ok(())
+            }
+            SliceOrInt::Slice(slice) => {
+                let len = self.__len__() as isize;
+                let indices = slice.indices(len).map_err(PyLoroError::from)?;
+
+                let values = py_any_to_loro_values(&value).map_err(PyLoroError::from)?;
+
+                if indices.step == 1 {
+                    self.0
+                        .delete(indices.start as usize, indices.slicelength)
+                        .map_err(PyLoroError::from)?;
+
+                    let mut pos = indices.start as usize;
+                    for v in values {
+                        self.0.insert(pos, v).map_err(PyLoroError::from)?;
+                        pos += 1;
+                    }
+                    Ok(())
+                } else {
+                    if values.len() != indices.slicelength {
+                        return Err(PyValueError::new_err(format!(
+                            "attempt to assign sequence of size {} to extended slice of size {}",
+                            values.len(),
+                            indices.slicelength
+                        ))
+                        .into());
+                    }
+
+                    let positions = slice_indices_positions(&indices);
+                    for (pos, v) in positions.into_iter().zip(values.into_iter()) {
+                        self.0.delete(pos, 1).map_err(PyLoroError::from)?;
+                        self.0.insert(pos, v).map_err(PyLoroError::from)?;
+                    }
+                    Ok(())
+                }
+            }
+        }
+    }
+
+    pub fn __delitem__<'py>(&self, index: SliceOrInt<'py>) -> PyLoroResult<()> {
+        match index {
+            SliceOrInt::Int(idx) => self.delete(idx, 1),
+            SliceOrInt::Slice(slice) => {
+                let len = self.__len__() as isize;
+                let indices = slice.indices(len).map_err(PyLoroError::from)?;
+
+                if indices.slicelength == 0 {
+                    return Ok(());
+                }
+
+                if indices.step == 1 {
+                    self.delete(indices.start as usize, indices.slicelength)
+                } else {
+                    let mut positions = slice_indices_positions(&indices);
+                    positions.sort_unstable();
+                    for pos in positions.into_iter().rev() {
+                        self.delete(pos, 1)?;
+                    }
+                    Ok(())
+                }
+            }
+        }
+    }
+
     /// Get the length of the list.
     #[inline]
     pub fn __len__(&self) -> usize {

src/container/map.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use loro::{ContainerTrait, LoroMap as LoroMapInner, PeerID};
-use pyo3::prelude::*;
+use pyo3::{exceptions::PyKeyError, prelude::*, PyErr};
 
 use crate::{
     doc::LoroDoc,
@@ -17,7 +17,7 @@ pub fn register_class(m: &Bound<'_, PyModule>) -> PyResult<()> {
     Ok(())
 }
 
-#[pyclass(frozen)]
+#[pyclass(frozen, mapping)]
 #[derive(Debug, Clone, Default)]
 pub struct LoroMap(pub LoroMapInner);
 
@@ -65,6 +65,34 @@ impl LoroMap {
         self.0.len()
     }
 
+    pub fn __contains__(&self, key: &str) -> bool {
+        self.0.get(key).is_some()
+    }
+
+    pub fn __getitem__(&self, key: &str) -> PyResult<ValueOrContainer> {
+        self.get(key)
+            .ok_or_else(|| PyKeyError::new_err(format!("Key {key} not found")))
+    }
+
+    pub fn __setitem__(&self, key: &str, value: LoroValue) -> PyLoroResult<()> {
+        self.insert(key, value)
+    }
+
+    pub fn __delitem__(&self, key: &str) -> PyResult<()> {
+        if !self.__contains__(key) {
+            return Err(PyKeyError::new_err(format!("Key {key} not found")));
+        }
+        self.delete(key).map_err(PyErr::from)?;
+        Ok(())
+    }
+
+    pub fn items(&self) -> Vec<(String, ValueOrContainer)> {
+        self.0
+            .keys()
+            .filter_map(|k| self.get(&k).map(|v| (k.to_string(), v)))
+            .collect()
+    }
+
     /// Get the ID of the map.
     #[getter]
     pub fn id(&self) -> ContainerID {