diff --git a/CHANGES/718.bugfix.rst b/CHANGES/718.bugfix.rst new file mode 100644 index 00000000..b23860a2 --- /dev/null +++ b/CHANGES/718.bugfix.rst @@ -0,0 +1,5 @@ +Fixed a bug where :class:`~frozenlist.FrozenList` could not be pickled/unpickled. + +Pickling is a requirement for :class:`~frozenlist.FrozenList` to be able to be passed to :class:`multiprocessing.Process`\ es. Without this +fix, users receive a :exc:`TypeError` upon attempting to pickle a :class:`~frozenlist.FrozenList`. +-- by :user:`csm10495`. diff --git a/CONTRIBUTORS.txt b/CONTRIBUTORS.txt index 0abae2ea..b7caf28f 100644 --- a/CONTRIBUTORS.txt +++ b/CONTRIBUTORS.txt @@ -1,6 +1,7 @@ - Contributors - ---------------- Andrew Svetlov +Charles Machalow Edgar Ramírez-Mondragón Marcin Konowalczyk Martijn Pieters diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 1b7227a6..50ee5237 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -302,6 +302,7 @@ unicode unittest Unittest unix +unpickled unsets unstripped upstr diff --git a/frozenlist/__init__.py b/frozenlist/__init__.py index e83f9aa0..dc93dc6c 100644 --- a/frozenlist/__init__.py +++ b/frozenlist/__init__.py @@ -1,3 +1,4 @@ +import copy import os import types from collections.abc import MutableSequence @@ -12,7 +13,7 @@ @total_ordering -class FrozenList(MutableSequence): +class PyFrozenList(MutableSequence): __slots__ = ("_frozen", "_items") __class_getitem__ = classmethod(types.GenericAlias) @@ -73,8 +74,36 @@ def __hash__(self): else: raise RuntimeError("Cannot hash unfrozen list.") + def __deepcopy__(self, memo: dict[int, object]): + obj_id = id(self) -PyFrozenList = FrozenList + # Create new instance and register immediately + new_list = self.__class__([]) + memo[obj_id] = new_list + + # Deep copy items + new_list._items[:] = [copy.deepcopy(item, memo) for item in self._items] + + # Preserve frozen state + if self._frozen: + new_list.freeze() + + return new_list + + def __reduce__(self): + return ( + self.__class__, + (self._items,), + {"_frozen": self._frozen}, + ) + + def __setstate__(self, state: dict[str, object]): + self._frozen = state["_frozen"] + + +# Rename the pure Python implementation. The C extension (if available) will +# override this name. +FrozenList = PyFrozenList if not NO_EXTENSIONS: diff --git a/frozenlist/_frozenlist.pyx b/frozenlist/_frozenlist.pyx index a82d8c8f..6a3a3e8a 100644 --- a/frozenlist/_frozenlist.pyx +++ b/frozenlist/_frozenlist.pyx @@ -127,10 +127,6 @@ cdef class FrozenList: cdef FrozenList new_list obj_id = id(self) - # Return existing copy if already processed (circular reference) - if obj_id in memo: - return memo[obj_id] - # Create new instance and register immediately new_list = self.__class__([]) memo[obj_id] = new_list @@ -144,5 +140,15 @@ cdef class FrozenList: return new_list + def __reduce__(self): + return ( + self.__class__, + (self._items,), + {"_frozen": self._frozen.load()}, + ) + + def __setstate__(self, state): + self._frozen.store(state["_frozen"]) + MutableSequence.register(FrozenList) diff --git a/tests/test_frozenlist.py b/tests/test_frozenlist.py index c37f5c0d..3a57c81f 100644 --- a/tests/test_frozenlist.py +++ b/tests/test_frozenlist.py @@ -1,6 +1,7 @@ # FIXME: # mypy: disable-error-code="misc" +import pickle from collections.abc import MutableSequence from copy import deepcopy @@ -8,6 +9,13 @@ from frozenlist import FrozenList, PyFrozenList +try: + from frozenlist import CFrozenList + + HAS_CFROZENLIST = True +except ImportError: + HAS_CFROZENLIST = False + class FrozenListMixin: FrozenList = NotImplemented @@ -273,6 +281,17 @@ def test_deepcopy_frozen(self) -> None: with pytest.raises(RuntimeError): copied.append(4) + def test_deepcopy_frozen_circular(self) -> None: + orig = self.FrozenList([1, 2]) + orig.append(orig) # Create circular reference + orig.freeze() + copied = deepcopy(orig) + assert copied[0] == 1 + assert copied[1] == 2 + assert copied[2] is copied + assert copied is not orig + assert orig.frozen + def test_deepcopy_nested(self) -> None: inner = self.FrozenList([1, 2]) orig = self.FrozenList([inner, 3]) @@ -372,10 +391,76 @@ def test_deepcopy_multiple_references(self) -> None: assert len(copied[1]) == 3 # Should see the change assert len(shared) == 2 # Original unchanged + @pytest.mark.parametrize("freeze", [True, False], ids=["frozen", "not frozen"]) + def test_picklability(self, freeze: bool) -> None: + # Test that the list can be pickled and unpickled successfully + orig = self.FrozenList([1, 2, 3]) + if freeze: + orig.freeze() + + assert orig.frozen == freeze + + pickled = pickle.dumps(orig) + unpickled = pickle.loads(pickled) + assert unpickled == orig + assert unpickled is not orig + assert list(unpickled) == list(orig) -class TestFrozenList(FrozenListMixin): - FrozenList = FrozenList # type: ignore[assignment] # FIXME + assert unpickled.frozen == freeze + + @pytest.mark.parametrize("freeze", [True, False], ids=["frozen", "not frozen"]) + def test_picklability_forward_compatible(self, freeze: bool) -> None: + orig = self.FrozenList([1, 2]) + if freeze: + orig.freeze() + + assert orig.frozen == freeze + + # 0 is the original pickle protocol. It's compatible with all supported Python versions. + pickled = pickle.dumps(orig, protocol=0) + + # If this test breaks, we've changed the frozenlist data structure in an incompatible way + # with previous pickled binaries. + if self.FrozenList is PyFrozenList: + if freeze: + assert ( + pickled + == b"cfrozenlist\nPyFrozenList\np0\n((lp1\nI1\naI2\natp2\nRp3\n(dp4\nV_frozen\np5\nI01\nsb." + ) + else: + assert ( + pickled + == b"cfrozenlist\nPyFrozenList\np0\n((lp1\nI1\naI2\natp2\nRp3\n(dp4\nV_frozen\np5\nI00\nsb." + ) + elif self.FrozenList is FrozenList: + if freeze: + assert ( + pickled + == b"cfrozenlist._frozenlist\nFrozenList\np0\n((lp1\nI1\naI2\natp2\nRp3\n(dp4\nV_frozen\np5\nI01\nsb." + ) + else: + assert ( + pickled + == b"cfrozenlist._frozenlist\nFrozenList\np0\n((lp1\nI1\naI2\natp2\nRp3\n(dp4\nV_frozen\np5\nI00\nsb." + ) + else: + pytest.fail("Unknown FrozenList implementation.") + + +if HAS_CFROZENLIST: + # If we don't have CFrozenList, skip adding the test class specifically for it. + class TestFrozenList(FrozenListMixin): + FrozenList = CFrozenList # type: ignore[assignment] # FIXME class TestFrozenListPy(FrozenListMixin): + # All implementations will at least have the Python version. FrozenList = PyFrozenList # type: ignore[assignment] # FIXME + + +def test_frozenlist_aliasing() -> None: + """Test that FrozenList name points to the C extension if available, else to PyFrozenList.""" + if HAS_CFROZENLIST: + assert FrozenList is CFrozenList + else: + assert FrozenList is PyFrozenList