feat: add __reduce__ to Array for correct un/pickling

decent_array/_array.py

@@ -26,6 +26,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, Self
 
 from decent_array.interoperability._backend_manager import register_backend_listener
@@ -49,6 +50,11 @@ def _update_backend(backend: Backend | None) -> None:
 register_backend_listener(_update_backend)
 
 
+def _reconstruct_array(value: ArrayTypes) -> Array:
+    """Reconstruct array after pickling."""
+    return Array(value)
+
+
 class Array:
     """
     Wrapper around a single backend-native array.
@@ -391,6 +397,12 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Array:
         memo[id(self)] = copied
         return copied
 
+    # Pickle/unpickle --------------------------------------------------------
+
+    def __reduce__(self) -> tuple[Callable[..., Array], tuple[ArrayTypes]]:
+        """Handle unpickling of the array."""
+        return (_reconstruct_array, (self.value,))
+
     # Properties -----------------------------------------------------------
 
     @property

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "decent-array"
-version = "0.2.2"
+version = "0.2.3"
 authors = [{name = "Simon Granström"}, {name = "Nicola Bastianello"}]
 maintainers = [{name = "Team Decent"}]
 description = "A library of array operations and linear algebra primitives for interoperability across ML frameworks."

tests/test_array.py

@@ -575,9 +575,39 @@ def test_item_n_dim(backend: tuple) -> None:
 def test_deepcopy(backend: tuple) -> None:
     from copy import deepcopy
 
-    src = iop.from_numpy(np.array([1.0, 2.0, 3.0], dtype=np.float32))
+    src = _create_array([1.0, 2.0, 3.0])
     dst = deepcopy(src)
     np.testing.assert_allclose(_np(dst), [1.0, 2.0, 3.0])
     # Mutating the copy shouldn't affect the original.
     dst[0] = 99.0
     np.testing.assert_allclose(_np(src), [1.0, 2.0, 3.0])
+
+
+# Pickle/unpickle --------------------------------------------------------
+
+
+def test_pickle(backend: tuple) -> None:
+    import pickle
+
+    original = _create_array([1.0, 2.0, 3.0])
+
+    payload = pickle.dumps(
+        original,
+        protocol=pickle.HIGHEST_PROTOCOL,
+    )
+
+    restored = pickle.loads(payload)
+
+    np.testing.assert_array_equal(_np(restored.value), _np(original.value))
+
+
+def test_pickle_uses_reduce(backend: tuple) -> None:
+    import pickle
+    from unittest.mock import patch
+
+    array = _create_array([1.0, 2.0, 3.0])
+
+    with patch.object(Array, "__reduce__", wraps=array.__reduce__) as reduce:
+        payload = pickle.dumps(array)
+
+    reduce.assert_called_once()