Support pickle serialization (#53)

Author: NoureldinYosri

test/test_value.py

@@ -14,10 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
+import pickle
+
 import numpy as np
 import pytest
 
 from tunits import Value, UnitMismatchError
+import tunits as tu
 
 
 def test_construction() -> None:
@@ -263,3 +267,15 @@ def test_dimensionless() -> None:
         _ = A.dimensionless()
 
     assert B.dimensionless() == 1.2
+
+
+def test_pick_roundtrip() -> None:
+    units = itertools.product(
+        [tu.ns, tu.GHz, tu.dBm, tu.deg, 1, tu.GHz**0.5, tu.m**0.75, tu.s**0.25], repeat=3
+    )
+    for value in np.random.random(20):
+        for unit_list in units:
+            unit = np.prod(unit_list)  # type: ignore[arg-type]
+            x = value * unit
+            s = pickle.dumps(x)
+            assert x == pickle.loads(s)

test/test_value_array.py

@@ -12,11 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
+import pickle
+
 import numpy as np
 import pytest
 from tunits.core import raw_WithUnit, raw_UnitArray
 
 from tunits import ValueArray, UnitMismatchError, Value
+import tunits as tu
 
 
 def test_construction() -> None:
@@ -238,3 +242,15 @@ def test_dimensionless() -> None:
         _ = A.dimensionless()
 
     np.testing.assert_equal(B.dimensionless(), np.arange(5) / 1000)
+
+
+def test_pick_roundtrip() -> None:
+    units = itertools.product(
+        [tu.ns, tu.GHz, tu.dBm, tu.deg, 1, tu.GHz**0.5, tu.m**0.75, tu.s**0.25], repeat=3
+    )
+    for value in np.random.random((5, 20)):
+        for unit_list in units:
+            unit = np.prod(unit_list)  # type: ignore[arg-type]
+            x = value * unit
+            s = pickle.dumps(x)
+            assert all(x == pickle.loads(s))

tunits/core/cython/unit_array.pyx

@@ -237,4 +237,19 @@ cdef class UnitArray:
             raise ValueError("UnitArray power does not support third argument")
         return self.pow_frac(float_to_twelths_frac(exponent));
 
+    def __getstate__(self):
+        return {
+            'unit_count': self.unit_count,
+            'units': [*self],
+        }
+
+    def __setstate__(self, pickle_info: dict[str, Any]):
+        self.unit_count = pickle_info['unit_count']
+        self.units = <UnitTerm *>PyMem_Malloc(self.unit_count*sizeof(UnitTerm))
+        for i, (name, numer, denom) in enumerate(pickle_info['units']):
+            Py_INCREF(name)
+            self.units[i].name = <PyObject *>name
+            self.units[i].power.numer = numer
+            self.units[i].power.denom = denom
+
 _EmptyUnit = UnitArray()

tunits/core/cython/with_unit.pyx

@@ -685,6 +685,22 @@ cdef class WithUnit:
     def _from_json_dict_(cls, **kwargs):
         return cls(kwargs["value"], kwargs["unit"])
 
+    def __getstate__(self):
+        return {
+            'value': self.value,
+            'conv': self.conv,
+            'display_units': self.display_units.__getstate__(),
+            'base_units': self.base_units.__getstate__(),
+        }
+
+    def __setstate__(self, pickle_info):
+        self.value = pickle_info['value']
+        self.conv = pickle_info['conv']
+        self.display_units = UnitArray()
+        self.base_units = UnitArray()
+        self.display_units.__setstate__(pickle_info['display_units'])
+        self.base_units.__setstate__(pickle_info['base_units'])
+
 _try_interpret_as_with_unit = None
 _is_value_consistent_with_default_unit_database = None
 def init_base_unit_functions(