gh-74876: Support pickling weak references and weak containers

Doc/library/pickle.rst

@@ -517,7 +517,11 @@ The following types can be pickled:
 * classes accessible from the top level of a module;
 
 * instances of such classes for which the result of calling :meth:`~object.__getstate__`
-  is picklable  (see section :ref:`pickle-inst` for details).
+  is picklable  (see section :ref:`pickle-inst` for details);
+
+* weak references and the weak containers from the :mod:`weakref` module,
+  provided that the weakly referenced objects are strongly referenced
+  elsewhere in the same pickle (see :class:`weakref.ref`).
 
 Attempts to pickle unpicklable objects will raise the :exc:`PicklingError`
 exception; when this happens, an unspecified number of bytes may have already

Doc/library/weakref.rst

@@ -126,6 +126,17 @@ See :ref:`__slots__ documentation <slots>` for details.
    Weak references are :ref:`generic <generics>` over the type of the object they
    reference.
 
+   A weak reference can be pickled
+   if the referent is strongly referenced elsewhere in the same pickle;
+   otherwise pickling raises :exc:`pickle.PicklingError`.
+   The unpickled weak reference refers to the unpickled referent.
+   A reference to a dead object unpickles as a reference to a dead object.
+   The *callback* is not preserved.
+
+   Weak references are atomic for the :mod:`copy` module:
+   :func:`copy.copy` and :func:`copy.deepcopy` return
+   the original weak reference without copying the referent.
+
    .. attribute:: __callback__
 
       This read-only attribute returns the callback currently associated to the
@@ -138,6 +149,9 @@ See :ref:`__slots__ documentation <slots>` for details.
    .. versionchanged:: next
       Raise :exc:`!TypeError` if *callback* is not callable or ``None``.
 
+   .. versionchanged:: next
+      Weak references can now be pickled.
+
 
 .. function:: proxy(object[, callback])
 
@@ -203,9 +217,16 @@ See :ref:`__slots__ documentation <slots>` for details.
       >>> d[k2] = 2   # d = {k2: 2}
       >>> del k1      # d = {k2: 2}
 
+   A :class:`!WeakKeyDictionary` can be copied,
+   and pickled if the keys are strongly referenced
+   elsewhere in the same pickle (see :class:`ref`).
+
    .. versionchanged:: 3.9
       Added support for ``|`` and ``|=`` operators, as specified in :pep:`584`.
 
+   .. versionchanged:: next
+      :class:`!WeakKeyDictionary` can now be pickled.
+
 :class:`WeakKeyDictionary` objects have an additional method that
 exposes the internal references directly.  The references are not guaranteed to
 be "live" at the time they are used, so the result of calling the references
@@ -224,9 +245,16 @@ than needed.
    Mapping class that references values weakly.  Entries in the dictionary will be
    discarded when no strong reference to the value exists any more.
 
+   A :class:`!WeakValueDictionary` can be copied,
+   and pickled if the values are strongly referenced
+   elsewhere in the same pickle (see :class:`ref`).
+
    .. versionchanged:: 3.9
       Added support for ``|`` and ``|=`` operators, as specified in :pep:`584`.
 
+   .. versionchanged:: next
+      :class:`!WeakValueDictionary` can now be pickled.
+
 :class:`WeakValueDictionary` objects have an additional method that has the
 same issues as the :meth:`WeakKeyDictionary.keyrefs` method.
 
@@ -241,6 +269,14 @@ same issues as the :meth:`WeakKeyDictionary.keyrefs` method.
    Set class that keeps weak references to its elements.  An element will be
    discarded when no strong reference to it exists any more.
 
+   A :class:`!WeakSet` can be pickled if the elements are strongly
+   referenced elsewhere in the same pickle (see :class:`ref`).
+   A copy or a deep copy shares the elements with the original;
+   user-defined attributes are copied or deep-copied.
+
+   .. versionchanged:: next
+      :class:`!WeakSet` can now be pickled and copied.
+
 
 .. class:: WeakMethod(method[, callback])
 
@@ -270,8 +306,16 @@ same issues as the :meth:`WeakKeyDictionary.keyrefs` method.
 
    *callback* is the same as the parameter of the same name to the :func:`ref` function.
 
+   A :class:`!WeakMethod` can be copied,
+   and pickled if the object the method is bound to is strongly referenced
+   elsewhere in the same pickle (see :class:`ref`).
+   A copy or a deep copy is an equal weak method to the same object.
+
    .. versionadded:: 3.4
 
+   .. versionchanged:: next
+      :class:`!WeakMethod` can now be pickled and copied.
+
 .. class:: finalize(obj, func, /, *args, **kwargs)
 
    Return a callable finalizer object which will be called when *obj*

Doc/whatsnew/3.16.rst

@@ -341,6 +341,17 @@ os
   (Contributed by Maurycy Pawłowski-Wieroński in :gh:`149464`.)
 
 
+pickle
+------
+
+* Weak references (:class:`weakref.ref`), :class:`weakref.WeakMethod`,
+  :class:`weakref.WeakSet`, :class:`weakref.WeakValueDictionary` and
+  :class:`weakref.WeakKeyDictionary` can now be pickled if the weakly
+  referenced objects are strongly referenced elsewhere in the same pickle;
+  otherwise pickling raises :exc:`pickle.PicklingError`.
+  :class:`!WeakMethod` and :class:`!WeakSet` can now also be copied.
+  (Contributed by Serhiy Storchaka in :gh:`74876`.)
+
 re
 --
 

Lib/_weakrefset.py

@@ -9,6 +9,10 @@
 
 
 class WeakSet:
+    # Slots give __getstate__() a uniform (dict, slots) shape for __reduce__.
+    # __dict__ carries user attributes; add() needs __weakref__.
+    __slots__ = '__dict__', '__weakref__', 'data', '_remove'
+
     def __init__(self, data=None):
         self.data = set()
 
@@ -40,7 +44,14 @@ def __contains__(self, item):
         return wr in self.data
 
     def __reduce__(self):
-        return self.__class__, (list(self),), self.__getstate__()
+        # Pickle the elements as weak references; data and _remove are
+        # rebuilt by the reconstructor.  Copying shares the elements.
+        dict, slots = self.__getstate__()
+        for name in ('data', '_remove'):
+            slots.pop(name, None)
+        state = (dict, slots) if dict or slots else None
+        return (_restore_weakset, (self.__class__, [ref(item) for item in self]),
+                state)
 
     def add(self, item):
         self.data.add(ref(item, self._remove))
@@ -145,3 +156,10 @@ def __repr__(self):
         return repr(self.data)
 
     __class_getitem__ = classmethod(GenericAlias)
+
+
+def _restore_weakset(cls, refs):
+    self = cls.__new__(cls)
+    WeakSet.__init__(self)
+    self.update(item for r in refs for item in [r()] if item is not None)
+    return self

Lib/pickle.py

@@ -32,6 +32,7 @@
 from struct import pack, unpack
 import io
 import codecs
+import weakref
 import _compat_pickle
 
 __all__ = ["PickleError", "PicklingError", "UnpicklingError", "Pickler",
@@ -485,6 +486,9 @@ def __init__(self, file, protocol=None, *, fix_imports=True,
         self.write = self.framer.write
         self._write_large_bytes = self.framer.write_large_bytes
         self.memo = {}
+        # Ids of objects reached so far only through weak references; any
+        # left at the end of dump() is an error.  See save_weakref.
+        self._weak_only = set()
         self.proto = int(protocol)
         self.bin = protocol >= 1
         self.fast = 0
@@ -499,6 +503,7 @@ def clear_memo(self):
         useful when re-using picklers.
         """
         self.memo.clear()
+        self._weak_only.clear()
 
     def dump(self, obj):
         """Write a pickled representation of obj to the open file."""
@@ -512,6 +517,14 @@ def dump(self, obj):
         if self.proto >= 4:
             self.framer.start_framing()
         self.save(obj)
+        if self._weak_only:
+            # These referents would be collected right after loading.  Report
+            # distinct type names: referents may be numerous and unhashable.
+            names = {_T(self.memo[i][1])
+                     for i in self._weak_only if i in self.memo}
+            raise PicklingError(
+                "cannot pickle weak reference(s) whose referent(s) have no "
+                "strong reference in the pickle: " + ", ".join(sorted(names)))
         self.write(STOP)
         self.framer.end_framing()
 
@@ -559,7 +572,7 @@ def get(self, i):
 
         return GET + repr(i).encode("ascii") + b'\n'
 
-    def save(self, obj, save_persistent_id=True):
+    def save(self, obj, save_persistent_id=True, weak=False):
         self.framer.commit_frame()
 
         # Check for persistent id (defined by a subclass)
@@ -572,6 +585,9 @@ def save(self, obj, save_persistent_id=True):
         # Check the memo
         x = self.memo.get(id(obj))
         if x is not None:
+            # A strong reference anchors the object: no longer weak-only.
+            if not weak:
+                self._weak_only.discard(id(obj))
             self.write(self.get(x[0]))
             return
 
@@ -633,6 +649,12 @@ def save(self, obj, save_persistent_id=True):
             exc.add_note(f'when serializing {_T(obj)} object')
             raise
 
+        # A completed non-weak save anchors obj.  Only save_reduce() memoizes
+        # an object after its arguments, so only here can a weak reference to
+        # a still-in-progress obj have tagged it.
+        if not weak:
+            self._weak_only.discard(id(obj))
+
     def persistent_id(self, obj):
         # This exists so a subclass can override it
         return None
@@ -1264,6 +1286,40 @@ def save_type(self, obj):
     dispatch[FunctionType] = save_global
     dispatch[type] = save_type
 
+    def save_weakref(self, obj):
+        # Reconstruct as weakref.ref(referent).  The referent is saved
+        # *weakly* (no anchoring, see save()) and tagged weak-only only after
+        # its whole subtree, so a strong cycle inside it cannot anchor it.
+        referent = obj()
+        if referent is None:
+            # All dead references are interchangeable: pickle by name as the
+            # weakref._dead_ref singleton.
+            self.write(GLOBAL + b'weakref\n_dead_ref\n')
+            return
+
+        save = self.save
+        write = self.write
+        save(weakref.ref)
+        # A referent already in the memo is anchored or already tracked;
+        # otherwise it is tagged below, until a strong save of it.
+        anchored = id(referent) in self.memo
+        # Build the 1-tuple (referent,); TUPLE1 needs protocol 2.
+        if self.proto < 2:
+            write(MARK)
+        save(referent, weak=True)
+        if not anchored and id(referent) in self.memo:
+            self._weak_only.add(id(referent))
+        write(TUPLE1 if self.proto >= 2 else TUPLE)
+        write(REDUCE)
+        # The referent's subtree may have pickled this very (cached) weakref
+        # already; reuse the memoized one.
+        if id(obj) in self.memo:
+            write(POP + self.get(self.memo[id(obj)][0]))
+        else:
+            self.memoize(obj)
+
+    dispatch[weakref.ReferenceType] = save_weakref
+
 
 # Unpickling machinery
 

Lib/test/pickletester.py

@@ -1755,9 +1755,61 @@ def t():
             [ToBeUnpickled] * 2)
 
 
+class WeakrefTarget:
+    # Weakly referenceable, picklable, and not hashable -- the weak-only
+    # error path must not assume hashable referents.
+    __hash__ = None
+
+
 class AbstractPicklingErrorTests:
     # Subclass must define self.dumps, self.pickler.
 
+    def test_unpicklable_weakref_referent(self):
+        # Pickling a weak reference pickles its referent too, so it fails when
+        # the referent is not picklable (a lambda).
+        f = lambda: 42
+        wr = weakref.ref(f)
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                self.assertRaises(pickle.PicklingError, self.dumps, wr, proto)
+
+    def test_weakref_without_strong_ref(self):
+        # The error names the distinct referent types: bounded however many
+        # referents there are, and not requiring them to be hashable.
+        referents = [WeakrefTarget() for _ in range(5)]
+        wrs = [weakref.ref(obj) for obj in referents]
+        typename = f'{WeakrefTarget.__module__}.{WeakrefTarget.__qualname__}'
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                with self.assertRaises(pickle.PicklingError) as cm:
+                    self.dumps(wrs, proto)
+                self.assertEqual(str(cm.exception),
+                    'cannot pickle weak reference(s) whose referent(s) have no '
+                    'strong reference in the pickle: ' + typename)
+
+    def test_weakref_reused_pickler(self):
+        # Both re-initializing the pickler and clearing its memo reset the
+        # weak-only bookkeeping left by a refused dump.
+        obj = WeakrefTarget()
+        wr = weakref.ref(obj)
+        p = self.pickler(io.BytesIO())
+        self.assertRaises(pickle.PicklingError, p.dump, wr)
+        f = io.BytesIO()
+        p.__init__(f)
+        p.dump([obj, wr])
+        obj2, wr2 = pickle.loads(f.getvalue())
+        self.assertIs(wr2(), obj2)
+
+        f = io.BytesIO()
+        p = self.pickler(f)
+        self.assertRaises(pickle.PicklingError, p.dump, wr)
+        p.clear_memo()
+        f.seek(0)
+        f.truncate()
+        p.dump([obj, wr])
+        obj2, wr2 = pickle.loads(f.getvalue())
+        self.assertIs(wr2(), obj2)
+
     def test_bad_reduce_result(self):
         obj = REX([print, ()])
         for proto in protocols:
@@ -3529,6 +3581,47 @@ def test_newobj_proxies(self):
                     self.assertEqual(B(x), B(y), detail)
                     self.assertEqual(x.__dict__, y.__dict__, detail)
 
+    def test_weakref(self):
+        obj = WeakrefTarget()
+        wr = weakref.ref(obj)
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                obj2, wr2 = self.loads(self.dumps([obj, wr], proto))
+                self.assertIs(type(wr2), weakref.ref)
+                self.assertIs(wr2(), obj2)
+                # Identity holds regardless of pickling order.
+                wr2, obj2 = self.loads(self.dumps([wr, obj], proto))
+                self.assertIs(wr2(), obj2)
+
+    def test_weakref_callback(self):
+        # The callback is not preserved: pickling it strongly would defeat
+        # the weak-only detection.
+        obj = WeakrefTarget()
+        wr = weakref.ref(obj, lambda r: None)
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                obj2, wr2 = self.loads(self.dumps([obj, wr], proto))
+                self.assertIs(wr2(), obj2)
+                self.assertIsNone(wr2.__callback__)
+
+    def test_dead_weakref(self):
+        # A dead reference stays dead, and all dead references unpickle as
+        # the same object.
+        obj = WeakrefTarget()
+        wr = weakref.ref(obj)
+        obj2 = WeakrefTarget()
+        wr2 = weakref.ref(obj2)
+        del obj, obj2
+        support.gc_collect()
+        self.assertIsNone(wr())
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                a, b = self.loads(self.dumps([wr, wr2], proto))
+                self.assertIs(type(a), weakref.ref)
+                self.assertIs(a, b)
+                support.gc_collect()
+                self.assertIsNone(a())
+
     def test_newobj_overridden_new(self):
         # Test that Python class with C implemented __new__ is pickleable
         for proto in protocols:

Lib/test/test_pickle.py

@@ -507,7 +507,7 @@ class SizeofTests(unittest.TestCase):
         check_sizeof = support.check_sizeof
 
         def test_pickler(self):
-            basesize = support.calcobjsize('7P2n3i2n4i2P')
+            basesize = support.calcobjsize('7P2n3i2n4i2P1n')
             p = _pickle.Pickler(io.BytesIO())
             self.assertEqual(object.__sizeof__(p), basesize)
             MT_size = struct.calcsize('3nP0n')