Make sure that frozenlists are pickle-able

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`.

CONTRIBUTORS.txt

@@ -1,6 +1,7 @@
 - Contributors -
 ----------------
 Andrew Svetlov
+Charles Machalow
 Edgar Ramírez-Mondragón
 Marcin Konowalczyk
 Martijn Pieters

docs/spelling_wordlist.txt

@@ -302,6 +302,7 @@ unicode
 unittest
 Unittest
 unix
+unpickled
 unsets
 unstripped
 upstr

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:

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)

tests/test_frozenlist.py

@@ -1,13 +1,21 @@
 # FIXME:
 # mypy: disable-error-code="misc"
 
+import pickle
 from collections.abc import MutableSequence
 from copy import deepcopy
 
 import pytest
 
 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