diff --git a/icechunk-python/python/icechunk/_icechunk_python.pyi b/icechunk-python/python/icechunk/_icechunk_python.pyi index 74cefe5fe..d7cba78e2 100644 --- a/icechunk-python/python/icechunk/_icechunk_python.pyi +++ b/icechunk-python/python/icechunk/_icechunk_python.pyi @@ -2424,6 +2424,23 @@ class IcechunkError(Exception): class ConflictError(Exception): """An error that occurs when a conflict is detected""" + def __init__( + self, + expected_parent: str | None = None, + actual_parent: str | None = None, + ) -> None: + """ + Create a new ConflictError. + + Parameters + ---------- + expected_parent: str | None + The expected parent snapshot ID. + actual_parent: str | None + The actual parent snapshot ID of the branch. + """ + ... + @property def expected_parent(self) -> str: """The expected parent snapshot ID. @@ -2485,6 +2502,26 @@ class ConflictType(Enum): class Conflict: """A conflict detected between snapshots""" + def __init__( + self, + conflict_type: ConflictType, + path: str, + conflicted_chunks: list[list[int]] | None = None, + ) -> None: + """ + Create a new Conflict. + + Parameters + ---------- + conflict_type: ConflictType + The type of conflict. + path: str + The path of the node that caused the conflict. + conflicted_chunks: list[list[int]] | None + If the conflict is a chunk conflict, the list of chunk indices in conflict. + """ + ... + @property def conflict_type(self) -> ConflictType: """The type of conflict detected @@ -2515,6 +2552,19 @@ class Conflict: class RebaseFailedError(IcechunkError): """An error that occurs when a rebase operation fails""" + def __init__(self, snapshot: str, conflicts: list[Conflict]) -> None: + """ + Create a new RebaseFailedError. + + Parameters + ---------- + snapshot: str + The snapshot ID that the session was rebased to. + conflicts: list[Conflict] + The conflicts that occurred during the rebase operation. + """ + ... + @property def snapshot(self) -> str: """The snapshot ID that the session was rebased to""" diff --git a/icechunk-python/src/conflicts.rs b/icechunk-python/src/conflicts.rs index 5ebabe437..e8b63dd2e 100644 --- a/icechunk-python/src/conflicts.rs +++ b/icechunk-python/src/conflicts.rs @@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize}; use crate::impl_pickle; -#[pyclass(name = "ConflictType", eq)] +#[pyclass(name = "ConflictType", module = "icechunk", eq)] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) enum PyConflictType { NewNodeConflictsWithExistingNode = 1, @@ -59,6 +59,27 @@ impl Display for PyConflictType { #[pymethods] impl PyConflictType { + #[new] + fn new(value: i32) -> PyResult { + match value { + 1 => Ok(PyConflictType::NewNodeConflictsWithExistingNode), + 2 => Ok(PyConflictType::NewNodeInInvalidGroup), + 3 => Ok(PyConflictType::ZarrMetadataDoubleUpdate), + 4 => Ok(PyConflictType::ZarrMetadataUpdateOfDeletedArray), + 5 => Ok(PyConflictType::ZarrMetadataUpdateOfDeletedGroup), + 6 => Ok(PyConflictType::ChunkDoubleUpdate), + 7 => Ok(PyConflictType::ChunksUpdatedInDeletedArray), + 8 => Ok(PyConflictType::ChunksUpdatedInUpdatedArray), + 9 => Ok(PyConflictType::DeleteOfUpdatedArray), + 10 => Ok(PyConflictType::DeleteOfUpdatedGroup), + 11 => Ok(PyConflictType::MoveOperationCannotBeRebased), + _ => Err(PyErr::new::(format!( + "Invalid ConflictType value: {}", + value + ))), + } + } + fn __repr__(&self) -> String { format!("{self:?}") } @@ -66,11 +87,31 @@ impl PyConflictType { fn __str__(&self) -> String { format!("{self}") } + + fn __reduce__(&self, py: Python<'_>) -> PyResult<(Py, Py)> { + use pyo3::IntoPyObjectExt; + let cls = py.get_type::().into_py_any(py)?; + let value: i32 = match self { + PyConflictType::NewNodeConflictsWithExistingNode => 1, + PyConflictType::NewNodeInInvalidGroup => 2, + PyConflictType::ZarrMetadataDoubleUpdate => 3, + PyConflictType::ZarrMetadataUpdateOfDeletedArray => 4, + PyConflictType::ZarrMetadataUpdateOfDeletedGroup => 5, + PyConflictType::ChunkDoubleUpdate => 6, + PyConflictType::ChunksUpdatedInDeletedArray => 7, + PyConflictType::ChunksUpdatedInUpdatedArray => 8, + PyConflictType::DeleteOfUpdatedArray => 9, + PyConflictType::DeleteOfUpdatedGroup => 10, + PyConflictType::MoveOperationCannotBeRebased => 11, + }; + let args = (value,).into_py_any(py)?; + Ok((cls, args)) + } } impl_pickle!(PyConflictType); -#[pyclass(name = "Conflict")] +#[pyclass(name = "Conflict", module = "icechunk")] #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct PyConflict { #[pyo3(get)] @@ -83,6 +124,16 @@ pub(crate) struct PyConflict { #[pymethods] impl PyConflict { + #[new] + #[pyo3(signature = (conflict_type, path, conflicted_chunks=None))] + fn new( + conflict_type: PyConflictType, + path: String, + conflicted_chunks: Option>>, + ) -> Self { + Self { conflict_type, path, conflicted_chunks } + } + fn __repr__(&self) -> String { format!("Conflict({:?}, path={})", self.conflict_type, self.path) } @@ -90,6 +141,18 @@ impl PyConflict { fn __str__(&self) -> String { format!("{}: {}", self.path, self.conflict_type) } + + fn __reduce__(&self, py: Python<'_>) -> PyResult<(Py, Py)> { + use pyo3::IntoPyObjectExt; + let cls = py.get_type::().into_py_any(py)?; + let args = ( + self.conflict_type.clone(), + self.path.clone(), + self.conflicted_chunks.clone(), + ) + .into_py_any(py)?; + Ok((cls, args)) + } } impl_pickle!(PyConflict); diff --git a/icechunk-python/tests/test_conflicts.py b/icechunk-python/tests/test_conflicts.py index d66b9cf29..2a9829f55 100644 --- a/icechunk-python/tests/test_conflicts.py +++ b/icechunk-python/tests/test_conflicts.py @@ -260,3 +260,26 @@ async def test_rebase_async(any_spec_version: int | None) -> None: # Should have session_b's values due to UseOurs selection assert array_final[0, 0] == 2 assert array_final.attrs["repo"] == 1 + + +def test_conflict_repr_and_str() -> None: + """Test that Conflict types have correct string representations.""" + conflict = icechunk.Conflict( + icechunk.ConflictType.ChunkDoubleUpdate, + "/my/array", + [[0, 0]], + ) + + # Test __repr__ + repr_str = repr(conflict) + assert "ChunkDoubleUpdate" in repr_str + assert "/my/array" in repr_str + + # Test __str__ + str_str = str(conflict) + assert "/my/array" in str_str + + # Test ConflictType repr and str + conflict_type = icechunk.ConflictType.ZarrMetadataDoubleUpdate + assert "ZarrMetadataDoubleUpdate" in repr(conflict_type) + assert "metadata" in str(conflict_type).lower() diff --git a/icechunk-python/tests/test_pickle.py b/icechunk-python/tests/test_pickle.py index 8462baaf8..4646e095b 100644 --- a/icechunk-python/tests/test_pickle.py +++ b/icechunk-python/tests/test_pickle.py @@ -6,7 +6,11 @@ import zarr from icechunk import ( + Conflict, + ConflictError, + ConflictType, IcechunkError, + RebaseFailedError, Repository, RepositoryConfig, S3StaticCredentials, @@ -123,3 +127,55 @@ def test_pickle_error() -> None: pickled = pickle.dumps(error) roundtripped = pickle.loads(pickled) assert error.message == roundtripped.message + + +def test_pickle_conflict_types() -> None: + """Test that Conflict types can be constructed and pickled correctly.""" + # Test Conflict with conflicted_chunks + conflict = Conflict( + ConflictType.ChunkDoubleUpdate, + "/some/array", + [[0, 0], [1, 1]], + ) + assert conflict.conflict_type == ConflictType.ChunkDoubleUpdate + assert conflict.path == "/some/array" + assert conflict.conflicted_chunks == [[0, 0], [1, 1]] + + # Test pickle roundtrip + roundtripped = pickle.loads(pickle.dumps(conflict)) + assert roundtripped.conflict_type == conflict.conflict_type + assert roundtripped.path == conflict.path + assert roundtripped.conflicted_chunks == conflict.conflicted_chunks + + # Test Conflict without conflicted_chunks + conflict_no_chunks = Conflict( + ConflictType.ZarrMetadataDoubleUpdate, + "/another/array", + ) + assert conflict_no_chunks.conflicted_chunks is None + + roundtripped_no_chunks = pickle.loads(pickle.dumps(conflict_no_chunks)) + assert roundtripped_no_chunks.conflict_type == conflict_no_chunks.conflict_type + assert roundtripped_no_chunks.path == conflict_no_chunks.path + assert roundtripped_no_chunks.conflicted_chunks is None + + # Test RebaseFailedError construction and pickle + error = RebaseFailedError("snapshot_123", [conflict, conflict_no_chunks]) + assert error.snapshot == "snapshot_123" + assert len(error.conflicts) == 2 + assert error.conflicts[0].path == "/some/array" + assert error.conflicts[1].path == "/another/array" + + roundtripped_error = pickle.loads(pickle.dumps(error)) + assert roundtripped_error.snapshot == error.snapshot + assert len(roundtripped_error.conflicts) == len(error.conflicts) + assert roundtripped_error.conflicts[0].path == error.conflicts[0].path + + # Test ConflictError construction and pickle + conflict_error = ConflictError("expected_snap", "actual_snap") + assert conflict_error.expected_parent == "expected_snap" + assert conflict_error.actual_parent == "actual_snap" + + roundtripped_conflict_error = pickle.loads(pickle.dumps(conflict_error)) + assert roundtripped_conflict_error.expected_parent == conflict_error.expected_parent + assert roundtripped_conflict_error.actual_parent == conflict_error.actual_parent