From f46a0909bb1d312e1263b600d04726dcbf298a98 Mon Sep 17 00:00:00 2001
From: nicola-bastianello <bastianello@pm.me>
Date: Mon, 7 Sep 2026 17:50:45 +0200
Subject: [PATCH] feat: add __reduce__ to Array for correct un/pickling

---
 decent_array/_array.py | 12 ++++++++++++
 pyproject.toml         |  2 +-
 tests/test_array.py    | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/decent_array/_array.py b/decent_array/_array.py
index 81f4f1d..45dbbe0 100644
--- a/decent_array/_array.py
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index 68b4542..937c05c 100644
--- a/pyproject.toml
+++ b/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."
diff --git a/tests/test_array.py b/tests/test_array.py
index 0e7faf9..de7078e 100644
--- a/tests/test_array.py
+++ b/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()