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
@@ -73,10 +74,46 @@ def __hash__(self):
         else:
             raise RuntimeError("Cannot hash unfrozen list.")
 
+    def __deepcopy__(self, memo: dict[int, object]):
+        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
+
+        # 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 (
+            _reconstruct_pyfrozenlist,
+            (self._items, self._frozen),
+        )
+
+
+# Store a reference to the pure Python implementation before it's potentially replaced
 PyFrozenList = FrozenList
 
 
+def _reconstruct_pyfrozenlist(items: list[object], frozen: bool) -> PyFrozenList:
+    """Helper function to reconstruct the pure Python FrozenList during unpickling.
+    This function is needed since otherwise the class renaming confuses pickle."""
+    fl = PyFrozenList(items)
+    if frozen:
+        fl.freeze()
+    return fl
+
+
 if not NO_EXTENSIONS:
     try:
         from ._frozenlist import FrozenList as CFrozenList  # type: ignore

frozenlist/_frozenlist.pyx

@@ -144,5 +144,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,6 +1,7 @@
 # FIXME:
 # mypy: disable-error-code="misc"
 
+import pickle
 from collections.abc import MutableSequence
 from copy import deepcopy
 
@@ -273,6 +274,20 @@ 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 len(copied[2]) == 3
+        assert copied[2][0] == 1
+        assert copied[2][1] == 2
+        assert len(copied[2][2]) == 3
+        # ... and so on. Testing equality when a structure includes itself is tough.
+        assert orig.frozen
+
     def test_deepcopy_nested(self) -> None:
         inner = self.FrozenList([1, 2])
         orig = self.FrozenList([inner, 3])
@@ -372,6 +387,23 @@ 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)
+
+        assert unpickled.frozen == freeze
+
 
 class TestFrozenList(FrozenListMixin):
     FrozenList = FrozenList  # type: ignore[assignment]  # FIXME