Merge pull request #192 from asmeurer/serialization-fix

Fix serialization of Cythonized objects

docs/changelog.md

@@ -1,5 +1,12 @@
 # ndindex Changelog
 
+## Version 1.9.2 (2024-09-25)
+
+## Minor Changes
+
+- Fixes an issue with pickle and deepcopy serialization introduced in ndindex
+  1.9.
+
 ## Version 1.9.1 (2024-09-23)
 
 This version is identical to 1.9, but includes some fixes to the release

ndindex/_slice.pyx

@@ -92,6 +92,15 @@ cdef class _Slice:
 
         self._reduced = _reduced
 
+    def __getnewargs__(self):
+        return self.args
+
+    def __getstate__(self):
+        return self._reduced
+
+    def __setstate__(self, state):
+        self._reduced = state
+
     @property
     def raw(self):
         return slice(*self.args)

ndindex/_tuple.pyx

@@ -111,7 +111,7 @@ cdef class _Tuple:
                     raise IndexError("an index can only have a single ellipsis ('...')")
             if isinstance(newarg, _Tuple):
                 if len(args) == 1:
-                    raise ValueError("tuples inside of tuple indices are not supported. Did you mean to call _Tuple(*args) instead of _Tuple(args)?")
+                    raise ValueError("tuples inside of tuple indices are not supported. Did you mean to call Tuple(*args) instead of Tuple(args)?")
                 raise ValueError("tuples inside of tuple indices are not supported. If you meant to use a fancy index, use a list or array instead.")
             newargs.append(newarg)
             if isinstance(newarg, _ArrayIndex):
@@ -152,6 +152,15 @@ cdef class _Tuple:
 
         self.args = tuple(newargs)
 
+    def __getnewargs__(self):
+        return self.args
+
+    def __setstate__(self, state):
+        pass
+
+    def __getstate__(self):
+        return ()
+
     @property
     def raw(self):
         return tuple(arg.raw for arg in self.args)

ndindex/tests/helpers.py

@@ -1,9 +1,9 @@
 import sys
 from itertools import chain
 import warnings
-from functools import wraps
+from functools import wraps, partial
 
-from numpy import intp, bool_, array, broadcast_shapes
+from numpy import ndarray, generic, intp, bool_, asarray, broadcast_shapes
 import numpy.testing
 
 from pytest import fail
@@ -373,16 +373,41 @@ def matmul_arrays_st(draw):
 
 reduce_kwargs = sampled_from([{}, {'negative_int': False}, {'negative_int': True}])
 
-def assert_equal(actual, desired, err_msg='', verbose=True):
+def assert_equal(actual, desired, allow_scalar_0d=False, err_msg='', verbose=True):
     """
-    Same as numpy.testing.assert_equal except it also requires the shapes and
-    dtypes to be equal.
+    Assert that two objects are equal.
+
+    - If the objects are ndarrays, this is the same as
+      numpy.testing.assert_equal except it also requires the shapes and dtypes
+      to be equal
+
+    - If the objects are tuples, recursively call assert_equal to support
+      tuples of arrays.
+
+    - If allow_scalar_0d=True, scalars will be considered equal to equivalent
+      0-D arrays.
+
+    - Require the types of actual and desired to be exactly the same
+      (excepting for scalars when allow_scalar_0d=True).
 
     """
-    numpy.testing.assert_equal(actual, desired, err_msg=err_msg,
-                               verbose=verbose)
-    assert actual.shape == desired.shape, err_msg or f"{actual.shape} != {desired.shape}"
-    assert actual.dtype == desired.dtype, err_msg or f"{actual.dtype} != {desired.dtype}"
+    if not (allow_scalar_0d and (isinstance(actual, generic)
+                                 or isinstance(desired, generic))):
+        assert type(actual) is type(desired), err_msg or f"{type(actual)} != {type(desired)}"
+
+    if isinstance(actual, (ndarray, generic)):
+        numpy.testing.assert_equal(actual, desired, err_msg=err_msg,
+                                   verbose=verbose)
+        assert actual.shape == desired.shape, err_msg or f"{actual.shape} != {desired.shape}"
+        assert actual.dtype == desired.dtype, err_msg or f"{actual.dtype} != {desired.dtype}"
+    elif isinstance(actual, tuple):
+        assert len(actual) == len(desired), err_msg
+        for i, j in zip(actual, desired):
+            assert_equal(i, j, err_msg=err_msg, verbose=verbose)
+    else:
+        assert actual == desired, err_msg
+
+assert_equal_allow_scalar_0d = partial(assert_equal, allow_scalar_0d=True)
 
 def warnings_are_errors(f):
     @wraps(f)
@@ -441,7 +466,7 @@ def assert_equal(x, y):
                 # deprecation was removed and lists are always interpreted as
                 # array indices.
                 if ("Using a non-tuple sequence for multidimensional indexing is deprecated" in w.args[0]): # pragma: no cover
-                    idx = array(idx, dtype=intp)
+                    idx = asarray(idx, dtype=intp)
                     a_raw = raw_func(a, idx)
                 elif "Out of bound index found. This was previously ignored when the indexing result contained no elements. In the future the index error will be raised. This error occurs either due to an empty slice, or if an array has zero elements even before indexing." in w.args[0]:
                     same_exception = False

ndindex/tests/test_ellipsis.py

@@ -4,7 +4,8 @@
 from hypothesis.strategies import one_of, integers
 
 from ..ndindex import ndindex
-from .helpers import check_same, prod, shapes, ellipses, reduce_kwargs
+from .helpers import (check_same, prod, shapes, ellipses, reduce_kwargs,
+                      assert_equal_allow_scalar_0d)
 
 def test_ellipsis_exhaustive():
     for n in range(10):
@@ -24,7 +25,9 @@ def test_ellipsis_reduce_exhaustive():
 @given(ellipses(), shapes, reduce_kwargs)
 def test_ellipsis_reduce_hypothesis(idx, shape, kwargs):
     a = arange(prod(shape)).reshape(shape)
-    check_same(a, idx, ndindex_func=lambda a, x: a[x.reduce(shape, **kwargs).raw])
+    check_same(a, idx,
+               ndindex_func=lambda a, x: a[x.reduce(shape, **kwargs).raw],
+               assert_equal=assert_equal_allow_scalar_0d)
 
 def test_ellipsis_reduce_no_shape_exhaustive():
     for n in range(10):
@@ -34,7 +37,8 @@ def test_ellipsis_reduce_no_shape_exhaustive():
 @given(ellipses(), shapes, reduce_kwargs)
 def test_ellipsis_reduce_no_shape_hypothesis(idx, shape, kwargs):
     a = arange(prod(shape)).reshape(shape)
-    check_same(a, idx, ndindex_func=lambda a, x: a[x.reduce(**kwargs).raw])
+    check_same(a, idx, ndindex_func=lambda a, x: a[x.reduce(**kwargs).raw],
+               assert_equal=assert_equal_allow_scalar_0d)
 
 @given(ellipses(), one_of(shapes, integers(0, 10)))
 def test_ellipsis_isempty_hypothesis(idx, shape):

ndindex/tests/test_expand.py

@@ -9,7 +9,8 @@
 from ..integerarray import IntegerArray
 from ..integer import Integer
 from ..tuple import Tuple
-from .helpers import ndindices, check_same, short_shapes, prod
+from .helpers import (ndindices, check_same, short_shapes, prod,
+                      assert_equal_allow_scalar_0d)
 
 @example(True, (1,))
 @example((Ellipsis, array([[ True,  True]])), (1, 2))
@@ -41,7 +42,7 @@ def test_expand_hypothesis(idx, shape):
     index = ndindex(idx)
 
     check_same(a, index.raw, ndindex_func=lambda a, x: a[x.expand(shape).raw],
-               same_exception=False)
+               same_exception=False, assert_equal=assert_equal_allow_scalar_0d)
 
     try:
         expanded = index.expand(shape)

ndindex/tests/test_ndindex.py

@@ -1,52 +1,53 @@
 import inspect
 import warnings
+import copy
+import pickle
 
 import numpy as np
 
 from hypothesis import given, example, settings
+from hypothesis.strategies import sampled_from
 
 from pytest import raises
 
+from ..array import ArrayIndex
 from ..ndindex import ndindex
 from ..booleanarray import BooleanArray
 from ..integer import Integer
 from ..ellipsis import ellipsis
 from ..integerarray import IntegerArray
+from ..slice import Slice
+from ..tuple import Tuple
 from .helpers import ndindices, check_same, assert_equal
 
-
 @example(None)
 @example([1, 2])
 @given(ndindices)
 def test_eq(idx):
     index = ndindex(idx)
     new = type(index)(*index.args)
 
-    if isinstance(new.raw, np.ndarray):
-        # trying to get a single value out of comparing two arrays requires all sorts of special handling, just let numpy do it
-        assert np.array_equal(new.raw, index.raw)
-    else:
-        assert new.raw == index.raw
+    assert_equal(new.raw, index.raw)
 
-    def assert_equal(a, b):
+    def check_eq(a, b):
         assert a == b
         assert b == a
         assert not (a != b)
         assert not (b != a)
 
-    def assert_not_equal(a, b):
+    def check_neq(a, b):
         assert a != b
         assert b != a
         assert not (a == b)
         assert not (b == a)
 
-    assert_equal(new, index)
-    assert_equal(new.raw, index)
-    assert_equal(new, index.raw)
+    check_eq(new, index)
+    check_eq(new.raw, index)
+    check_eq(new, index.raw)
 
-    assert_equal(index.raw, index)
+    check_eq(index.raw, index)
     assert hash(new) == hash(index)
-    assert_not_equal(index, 'a')
+    check_neq(index, 'a')
 
     try:
         h = hash(idx)
@@ -86,21 +87,24 @@ def test_ndindex(idx):
     assert index == idx
     assert ndindex[idx] == index
 
-    def test_raw_eq(idx, index):
-        if isinstance(idx, np.ndarray):
-            assert_equal(index.raw, idx)
-        elif isinstance(idx, list):
+    def assert_raw_eq(idx, index):
+        if isinstance(idx, (list, bool, np.bool_)):
+            assert isinstance(index, ArrayIndex)
             assert index.dtype in [np.intp, np.bool_]
             assert_equal(index.raw, np.asarray(idx, dtype=index.dtype))
+        elif isinstance(idx, np.integer):
+            assert type(index) is Integer
+            assert_equal(index.raw, int(idx))
         elif isinstance(idx, tuple):
-            assert type(index.raw) == type(idx)
-            assert len(index.raw) == len(idx)
+            assert type(index.raw) is tuple
+            assert len(idx) == len(index.raw)
             assert index.args == index.raw
-            for i, j in zip(idx, index.args):
-                test_raw_eq(i, j)
+            for i, j in zip(index.args, idx):
+                assert_raw_eq(j, i)
         else:
-            assert index.raw == idx
-    test_raw_eq(idx, index)
+            assert_equal(index.raw, idx)
+
+    assert_raw_eq(idx, index)
     assert ndindex(index.raw) == index
 
 def test_ndindex_invalid():
@@ -145,3 +149,40 @@ def test_repr_str(idx):
 
     # Str may not be re-creatable. Just test that it doesn't give an exception.
     str(index)
+
+# _Tuple does not serialize properly with protocols 0 and 1. Support could
+# probably be added if this is necessary.
+LOWEST_SUPPORTED_PROTOCOL = 2
+protocols = ["copy", "deepcopy"] + list(range(LOWEST_SUPPORTED_PROTOCOL, pickle.HIGHEST_PROTOCOL + 1))
+
+@given(ndindices, sampled_from(protocols))
+def test_serialization(idx, protocol):
+    index = ndindex(idx)
+
+    def serialize(index):
+        if protocol == "copy":
+            return copy.copy(index)
+        elif protocol == "deepcopy":
+            return copy.deepcopy(index)
+        else:
+            return pickle.loads(pickle.dumps(index, protocol=protocol))
+
+    roundtripped = serialize(index)
+    assert type(roundtripped) is type(index)
+    assert roundtripped == index
+    assert_equal(roundtripped.raw, index.raw)
+    assert_equal(roundtripped.args, index.args)
+
+    if isinstance(index, Slice):
+        assert index._reduced == roundtripped._reduced == False
+        s = index.reduce()
+        assert s._reduced == True
+        roundtripped_s = serialize(s)
+        assert roundtripped_s._reduced == True
+
+    if isinstance(index, Tuple):
+        assert all([i._reduced == False for i in index.args if isinstance(i, Slice)])
+        t = index.reduce()
+        assert all([i._reduced == True for i in t.args if isinstance(i, Slice)])
+        roundtripped_t = serialize(t)
+        assert all([i._reduced == True for i in roundtripped_t.args if isinstance(i, Slice)])

ndindex/tests/test_tuple.py

@@ -12,7 +12,8 @@
 from ..booleanarray import BooleanArray
 from ..integer import Integer
 from ..integerarray import IntegerArray
-from .helpers import check_same, Tuples, prod, short_shapes, iterslice, reduce_kwargs
+from .helpers import (check_same, Tuples, prod, short_shapes, iterslice,
+                      reduce_kwargs, assert_equal_allow_scalar_0d)
 
 def test_tuple_constructor():
     # Nested tuples are not allowed
@@ -96,7 +97,7 @@ def ndindex_func(a, index):
         return a[ndindex((*index.raw[:index.ellipsis_index], ...,
                           *index.raw[index.ellipsis_index+1:])).raw]
 
-    check_same(a, t, ndindex_func=ndindex_func)
+    check_same(a, t, ndindex_func=ndindex_func, assert_equal=assert_equal_allow_scalar_0d)
 
 @example((True, 0, False), 1, {})
 @example((..., None), (), {})
@@ -109,8 +110,9 @@ def test_tuple_reduce_no_shape_hypothesis(t, shape, kwargs):
 
     index = Tuple(*t)
 
-    check_same(a, index.raw, ndindex_func=lambda a, x: a[x.reduce(**kwargs).raw],
-               same_exception=False)
+    check_same(a, index.raw, ndindex_func=lambda a, x:
+               a[x.reduce(**kwargs).raw], same_exception=False,
+               assert_equal=assert_equal_allow_scalar_0d)
 
     reduced = index.reduce(**kwargs)
     if isinstance(reduced, Tuple):
@@ -143,8 +145,9 @@ def test_tuple_reduce_hypothesis(t, shape, kwargs):
 
     index = Tuple(*t)
 
-    check_same(a, index.raw, ndindex_func=lambda a, x: a[x.reduce(shape, **kwargs).raw],
-               same_exception=False)
+    check_same(a, index.raw, ndindex_func=lambda a, x: a[x.reduce(shape,
+                                                                  **kwargs).raw], same_exception=False,
+               assert_equal=assert_equal_allow_scalar_0d)
 
     negative_int = kwargs.get('negative_int', False)
 
@@ -197,7 +200,8 @@ def test_tuple_reduce_explicit():
 
         a = arange(prod(shape)).reshape(shape)
         check_same(a, before.raw, ndindex_func=lambda a, x:
-                   a[x.reduce(shape).raw])
+                   a[x.reduce(shape).raw],
+                   assert_equal=assert_equal_allow_scalar_0d)
 
         # Idempotency
         assert reduced.reduce() == reduced