diff --git a/Doc/library/pickle.rst b/Doc/library/pickle.rst index 8eadc2cf2b1ef0..f45b92d2287d98 100644 --- a/Doc/library/pickle.rst +++ b/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 diff --git a/Doc/library/weakref.rst b/Doc/library/weakref.rst index 7cc0c33fda353c..fde37a56f8c0ed 100644 --- a/Doc/library/weakref.rst +++ b/Doc/library/weakref.rst @@ -126,6 +126,17 @@ See :ref:`__slots__ documentation ` for details. Weak references are :ref:`generic ` 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 ` 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 ` 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* diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index b6a5e4c7fbf1b8..92f0c4992b8b8d 100644 --- a/Doc/whatsnew/3.16.rst +++ b/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 -- diff --git a/Lib/_weakrefset.py b/Lib/_weakrefset.py index d1c7fcaeec9821..c462e63f430ff7 100644 --- a/Lib/_weakrefset.py +++ b/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 diff --git a/Lib/pickle.py b/Lib/pickle.py index f92b1fde768fc7..72a2959cdb99f9 100644 --- a/Lib/pickle.py +++ b/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 diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py index 9ba498ce8f575d..fd0e655deee1fa 100644 --- a/Lib/test/pickletester.py +++ b/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: diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py index 48375cf459ea0b..3e656acad026f4 100644 --- a/Lib/test/test_pickle.py +++ b/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') diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py index b9d1745592c09c..b642c0afb424b1 100644 --- a/Lib/test/test_weakref.py +++ b/Lib/test/test_weakref.py @@ -7,6 +7,7 @@ import operator import contextlib import copy +import pickle import threading import time import random @@ -46,6 +47,14 @@ def create_bound_method(): return C().method +class WeakValueDictionaryWithSlots(weakref.WeakValueDictionary): + __slots__ = ('x', 'y') + + +class WeakKeyDictionaryWithSlots(weakref.WeakKeyDictionary): + __slots__ = ('x', 'y') + + class Object: def __init__(self, arg): self.arg = arg @@ -1316,6 +1325,55 @@ def test_hashing(self): # If it wasn't hashed when alive, a dead WeakMethod cannot be hashed. self.assertRaises(TypeError, hash, c) + def test_pickle(self): + # The bound method is rebuilt on unpickling, bound to the same + # restored instance regardless of pickling order. A callback does + # not prevent pickling. + o = Object(1) + r = weakref.WeakMethod(o.some_method, lambda arg: None) + for proto in range(pickle.HIGHEST_PROTOCOL + 1): + with self.subTest(proto=proto): + o2, r2 = pickle.loads(pickle.dumps([o, r], proto)) + self.assertIsInstance(r2, weakref.WeakMethod) + self.assertIs(r2().__self__, o2) + self.assertEqual(r2()(), 4) + r2, o2 = pickle.loads(pickle.dumps([r, o], proto)) + self.assertIs(r2().__self__, o2) + self.assertEqual(r2()(), 4) + + def test_pickle_dead(self): + # A WeakMethod to a dead object (or a dead method) cannot be pickled. + o = Object(1) + r = weakref.WeakMethod(o.some_method) + del o + gc.collect() + for proto in range(pickle.HIGHEST_PROTOCOL + 1): + with self.subTest(proto=proto): + self.assertRaises(TypeError, pickle.dumps, r, proto) + + def test_pickle_without_strong_ref(self): + # As for weakref.ref, an instance with no strong reference in the + # pickle is refused. + o = Object(1) + r = weakref.WeakMethod(o.some_method) + for proto in range(pickle.HIGHEST_PROTOCOL + 1): + with self.subTest(proto=proto): + self.assertRaises(pickle.PicklingError, pickle.dumps, r, proto) + + def test_copy(self): + # Copying rebuilds an equal weak method bound to the same instance + # (weak references are atomic under copy). + o = Object(1) + r = weakref.WeakMethod(o.some_method) + for copyfunc in copy.copy, copy.deepcopy: + with self.subTest(copyfunc=copyfunc): + r2 = copyfunc(r) + self.assertIsInstance(r2, weakref.WeakMethod) + self.assertIsNot(r2, r) + self.assertEqual(r2, r) + self.assertIs(r2().__self__, o) + self.assertEqual(r2()(), 4) + class MappingTestCase(TestBase): @@ -1656,8 +1714,8 @@ def test_make_weak_keyed_dict_from_weak_keyed_dict(self): dict2 = weakref.WeakKeyDictionary(dict) self.assertEqual(dict[o], 364) - def make_weak_keyed_dict(self): - dict = weakref.WeakKeyDictionary() + def make_weak_keyed_dict(self, cls=weakref.WeakKeyDictionary): + dict = cls() objects = list(map(Object, range(self.COUNT))) for o in objects: dict[o] = o.arg @@ -1686,13 +1744,59 @@ def test_make_weak_valued_dict_misc(self): self.assertEqual(list(d.keys()), [kw]) self.assertEqual(d[kw], o) - def make_weak_valued_dict(self): - dict = weakref.WeakValueDictionary() + def make_weak_valued_dict(self, cls=weakref.WeakValueDictionary): + dict = cls() objects = list(map(Object, range(self.COUNT))) for o in objects: dict[o.arg] = o return dict, objects + def test_pickling(self): + # The weak values/keys must be strongly referenced elsewhere in the + # pickle. 'x' is a slot attribute in the *WithSlots subclasses and a + # dict attribute in the bases; 'z' is a dict attribute; 'y' stays + # unset. + for proto in range(pickle.HIGHEST_PROTOCOL + 1): + for cls in weakref.WeakValueDictionary, WeakValueDictionaryWithSlots: + with self.subTest(proto=proto, cls=cls): + d, objects = self.make_weak_valued_dict(cls) + d.x = ['x'] + d.z = ['z'] + d2, objects2 = pickle.loads(pickle.dumps((d, objects), + proto)) + self.assertIsInstance(d2, cls) + self.assertEqual(dict(d2), dict(d)) + self.assertEqual(set(map(id, d2.values())), + set(map(id, objects2))) + self.assertEqual(d2.x, ['x']) + self.assertEqual(d2.z, ['z']) + self.assertNotHasAttr(d2, 'y') + + for cls in weakref.WeakKeyDictionary, WeakKeyDictionaryWithSlots: + with self.subTest(proto=proto, cls=cls): + d, objects = self.make_weak_keyed_dict(cls) + d.x = ['x'] + d.z = ['z'] + d2, objects2 = pickle.loads(pickle.dumps((d, objects), + proto)) + self.assertIsInstance(d2, cls) + self.assertEqual(dict(d2), dict(d)) + self.assertEqual(set(map(id, d2.keys())), + set(map(id, objects2))) + self.assertEqual(d2.x, ['x']) + self.assertEqual(d2.z, ['z']) + self.assertNotHasAttr(d2, 'y') + + def test_pickling_without_strong_ref(self): + # A key/value with no strong reference in the pickle is refused. + for proto in range(pickle.HIGHEST_PROTOCOL + 1): + d, objects = self.make_weak_valued_dict() + with self.assertRaises(pickle.PicklingError): + pickle.dumps(d, proto) + d, objects = self.make_weak_keyed_dict() + with self.assertRaises(pickle.PicklingError): + pickle.dumps(d, proto) + def check_popitem(self, klass, key1, value1, key2, value2): weakdict = klass() weakdict[key1] = value1 diff --git a/Lib/test/test_weakset.py b/Lib/test/test_weakset.py index c1e4f9c8366e58..60e1a154b3ec44 100644 --- a/Lib/test/test_weakset.py +++ b/Lib/test/test_weakset.py @@ -1,6 +1,7 @@ import unittest from weakref import WeakSet import copy +import pickle import string from collections import UserString as ustr from collections.abc import Set, MutableSet @@ -478,6 +479,34 @@ def test_copying(self): self.assertIsNot(dup.z, s.z) self.assertNotHasAttr(dup, 'y') + def test_pickling(self): + for cls in WeakSet, WeakSetWithSlots: + s = cls(self.items) + # 'x' is a slot attribute in WeakSetWithSlots and a dict + # attribute in WeakSet; 'z' is a dict attribute; 'y' stays unset. + s.x = ['x'] + s.z = ['z'] + for proto in range(pickle.HIGHEST_PROTOCOL + 1): + with self.subTest(proto=proto, cls=cls): + # The elements must be strongly referenced elsewhere + # in the pickle. + dup, items = pickle.loads( + pickle.dumps((s, self.items), proto)) + self.assertIsInstance(dup, cls) + self.assertEqual(dup, WeakSet(items)) + self.assertEqual(dup.x, ['x']) + self.assertEqual(dup.z, ['z']) + self.assertNotHasAttr(dup, 'y') + # The elements are shared with the strongly-pickled copy. + self.assertEqual(set(map(id, dup)), set(map(id, items))) + + def test_pickling_without_strong_ref(self): + # An element with no strong reference in the pickle is refused. + s = WeakSet(self.items) + for proto in range(pickle.HIGHEST_PROTOCOL + 1): + with self.assertRaises(pickle.PicklingError): + pickle.dumps(s, proto) + if __name__ == "__main__": unittest.main() diff --git a/Lib/weakref.py b/Lib/weakref.py index af7244553c908c..fc895d4033300b 100644 --- a/Lib/weakref.py +++ b/Lib/weakref.py @@ -22,6 +22,7 @@ from _weakrefset import WeakSet import _collections_abc # Import after _weakref to avoid circular import. +import copyreg import sys import itertools @@ -35,6 +36,12 @@ _collections_abc.MutableSet.register(WeakSet) + +# Pickle reconstructs a reference to a dead object as this singleton -- all +# dead references are interchangeable. Born dead: nothing else holds the set. +_dead_ref = ref(set()) + + class WeakMethod(ref): """ A custom `weakref.ref` subclass which simulates a weak reference to @@ -89,12 +96,34 @@ def __ne__(self, other): __hash__ = ref.__hash__ +# Pickling support for weak methods. A WeakMethod is reduced to a weak +# reference to its instance; the bound method is rebuilt on unpickling by +# looking the function up on the instance. The callback is not preserved. + +def _restore_weakmethod(obj_ref, name): + return WeakMethod(getattr(obj_ref(), name)) + + +def _reduce_weakmethod(self): + meth = self() + if meth is None: + raise TypeError( + "cannot pickle a weak reference to a dead object") + return (_restore_weakmethod, (ref(meth.__self__), meth.__func__.__name__)) + +copyreg.pickle(WeakMethod, _reduce_weakmethod) + + class WeakValueDictionary(_collections_abc.MutableMapping): """Mapping class that references values weakly. Entries in the dictionary will be discarded when no strong reference to the value exists anymore """ + # Slots give __getstate__() a uniform (dict, slots) shape for __reduce__. + # __dict__ carries user attributes; the removal callback needs __weakref__. + __slots__ = '__dict__', '__weakref__', 'data', '_remove' + # We inherit the constructor without worrying about the input # dictionary; since it uses our .update() method, we get the right # checks (if the other dictionary is a WeakValueDictionary, @@ -157,6 +186,16 @@ def __deepcopy__(self, memo): new[deepcopy(key, memo)] = o return new + def __reduce__(self): + # Pickle the values as weak references; data and _remove are rebuilt + # by the reconstructor. Copying uses __copy__/__deepcopy__ instead. + dict, slots = self.__getstate__() + for name in ('data', '_remove'): + slots.pop(name, None) + state = (dict, slots) if dict or slots else None + items = [(key, ref(o)) for key, o in self.items()] + return (_restore_weakvaluedict, (self.__class__, items), state) + def get(self, key, default=None): try: wr = self.data[key] @@ -274,6 +313,16 @@ def __ror__(self, other): return NotImplemented +def _restore_weakvaluedict(cls, items): + self = cls.__new__(cls) + WeakValueDictionary.__init__(self) + for key, r in items: + o = r() + if o is not None: + self[key] = o + return self + + class KeyedRef(ref): """Specialized reference that includes a key corresponding to the value. @@ -305,6 +354,8 @@ class WeakKeyDictionary(_collections_abc.MutableMapping): can be especially useful with objects that override attribute accesses. """ + # See the comment on WeakValueDictionary.__slots__. + __slots__ = '__dict__', '__weakref__', 'data', '_remove' def __init__(self, dict=None): self.data = {} @@ -353,6 +404,16 @@ def __deepcopy__(self, memo): new[o] = deepcopy(value, memo) return new + def __reduce__(self): + # Pickle the keys as weak references; data and _remove are rebuilt + # by the reconstructor. Copying uses __copy__/__deepcopy__ instead. + dict, slots = self.__getstate__() + for name in ('data', '_remove'): + slots.pop(name, None) + state = (dict, slots) if dict or slots else None + items = [(ref(key), value) for key, value in self.items()] + return (_restore_weakkeydict, (self.__class__, items), state) + def get(self, key, default=None): return self.data.get(ref(key),default) @@ -439,6 +500,16 @@ def __ror__(self, other): return NotImplemented +def _restore_weakkeydict(cls, items): + self = cls.__new__(cls) + WeakKeyDictionary.__init__(self) + for r, value in items: + key = r() + if key is not None: + self[key] = value + return self + + class finalize: """Class for finalization of weakrefable objects diff --git a/Misc/NEWS.d/next/Library/2026-07-16-21-30-00.gh-issue-74876.wR3fXk.rst b/Misc/NEWS.d/next/Library/2026-07-16-21-30-00.gh-issue-74876.wR3fXk.rst new file mode 100644 index 00000000000000..aef5cea67047fc --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-16-21-30-00.gh-issue-74876.wR3fXk.rst @@ -0,0 +1,8 @@ +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:`weakref.WeakMethod` and :class:`weakref.WeakSet` +can now also be copied. diff --git a/Modules/_pickle.c b/Modules/_pickle.c index 034dd06b6ab460..f4b2fe3049816d 100644 --- a/Modules/_pickle.c +++ b/Modules/_pickle.c @@ -603,6 +603,13 @@ typedef struct { Py_ssize_t me_value; } PyMemoEntry; +/* The top bit of me_value flags an object reached so far only through a weak + reference (memo indices are small non-negative ints, so it is free); + MEMO_INDEX() masks it out. Bit math in size_t; the store back into the + signed field relies on two's complement. */ +#define MEMO_WEAK_ONLY ((size_t)1 << (sizeof(size_t) * 8 - 1)) +#define MEMO_INDEX(v) ((Py_ssize_t)((size_t)(v) & ~MEMO_WEAK_ONLY)) + typedef struct { size_t mt_mask; size_t mt_used; @@ -647,6 +654,10 @@ typedef struct PicklerObject { int running; /* True when a method of Pickler is executing. */ PyObject *fast_memo; PyObject *buffer_callback; /* Callback for out-of-band buffers, or NULL */ + + Py_ssize_t weak_only_count; /* Number of memo entries flagged weak-only; + any left at the end of dump() are pickling + errors. */ } PicklerObject; typedef struct UnpicklerObject { @@ -736,6 +747,8 @@ typedef struct { /* Forward declarations */ static int save(PickleState *state, PicklerObject *, PyObject *, int); +static int save_object(PickleState *state, PicklerObject *, PyObject *, + int pers_save, int weak); static int save_reduce(PickleState *, PicklerObject *, PyObject *, PyObject *); #include "clinic/_pickle.c.h" @@ -1166,6 +1179,7 @@ _Pickler_New(PickleState *st) self->running = 0; self->fast_memo = NULL; self->buffer_callback = NULL; + self->weak_only_count = 0; PyObject_GC_Track(self); return self; @@ -1836,25 +1850,26 @@ memo_get(PickleState *st, PicklerObject *self, PyObject *key) PyErr_SetObject(PyExc_KeyError, key); return -1; } + Py_ssize_t idx = MEMO_INDEX(*value); /* mask off the weak-only flag bit */ if (!self->bin) { pdata[0] = GET; PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, - "%zd\n", *value); + "%zd\n", idx); len = strlen(pdata); } else { - if (*value < 256) { + if (idx < 256) { pdata[0] = BINGET; - pdata[1] = (unsigned char)(*value & 0xff); + pdata[1] = (unsigned char)(idx & 0xff); len = 2; } - else if ((size_t)*value <= 0xffffffffUL) { + else if ((size_t)idx <= 0xffffffffUL) { pdata[0] = LONG_BINGET; - pdata[1] = (unsigned char)(*value & 0xff); - pdata[2] = (unsigned char)((*value >> 8) & 0xff); - pdata[3] = (unsigned char)((*value >> 16) & 0xff); - pdata[4] = (unsigned char)((*value >> 24) & 0xff); + pdata[1] = (unsigned char)(idx & 0xff); + pdata[2] = (unsigned char)((idx >> 8) & 0xff); + pdata[3] = (unsigned char)((idx >> 16) & 0xff); + pdata[4] = (unsigned char)((idx >> 24) & 0xff); len = 5; } else { /* unlikely */ @@ -4548,8 +4563,154 @@ save_reduce(PickleState *st, PicklerObject *self, PyObject *args, return 0; } +/* Weak-only objects are flagged in place in the memo (the MEMO_WEAK_ONLY bit + of the value, see PyMemoEntry); self->weak_only_count tracks how many. The + value pointer comes from PyMemoTable_Get() and must not be held across an + insertion into the table (it may rehash). */ + +static void +memo_set_weak_only(PicklerObject *self, Py_ssize_t *value) +{ + if (*value >= 0) { /* not already flagged */ + *value = (Py_ssize_t)((size_t)*value | MEMO_WEAK_ONLY); + self->weak_only_count++; + } +} + +static void +memo_clear_weak_only(PicklerObject *self, Py_ssize_t *value) +{ + if (*value < 0) { /* flagged */ + *value = MEMO_INDEX(*value); + self->weak_only_count--; + } +} + +/* Pickle a weakref.ref as weakref.ref(referent). The referent is pickled + *weakly* (no anchoring) and tagged weak-only only after its whole subtree, + so a strong cycle inside it cannot anchor it. See Lib/pickle.py. */ static int -save(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save) +save_weakref(PickleState *st, PicklerObject *self, PyObject *obj) +{ + const char tuple1_op = TUPLE1; + const char tuple_op = TUPLE; + const char mark_op = MARK; + const char reduce_op = REDUCE; + const char pop_op = POP; + int status = -1; + + PyObject *referent; + if (PyWeakref_GetRef(obj, &referent) < 0) + return -1; + if (referent == NULL) { + /* All dead references are interchangeable: pickle by name as the + weakref._dead_ref singleton. */ + const char global_op = GLOBAL; + const char dead_ref_name[] = "weakref\n_dead_ref\n"; + if (_Pickler_Write(self, &global_op, 1) < 0 || + _Pickler_Write(self, dead_ref_name, sizeof(dead_ref_name) - 1) < 0) + { + return -1; + } + return 0; + } + + /* Py_TYPE(obj) is exactly weakref.ref, the reconstructor. */ + if (save(st, self, (PyObject *)Py_TYPE(obj), 0) < 0) + goto done; + + /* A referent already in the memo is anchored or already tracked; + otherwise it is tagged below, until a strong save of it. */ + int anchored = PyMemoTable_Get(self->memo, referent) != NULL; + + /* Build the 1-tuple (referent,); TUPLE1 needs protocol 2. */ + if (self->proto < 2 && _Pickler_Write(self, &mark_op, 1) < 0) + goto done; + + if (save_object(st, self, referent, 0, 1) < 0) + goto done; + + if (!anchored) { + Py_ssize_t *value = PyMemoTable_Get(self->memo, referent); + if (value != NULL) + memo_set_weak_only(self, value); + } + + if (_Pickler_Write(self, self->proto >= 2 ? &tuple1_op : &tuple_op, 1) < 0 || + _Pickler_Write(self, &reduce_op, 1) < 0) + goto done; + + /* The referent's subtree may have pickled this very (cached) weakref + already; reuse the memoized one. */ + if (PyMemoTable_Get(self->memo, obj)) { + if (_Pickler_Write(self, &pop_op, 1) < 0) + goto done; + if (memo_get(st, self, obj) < 0) + goto done; + } + else if (memo_put(st, self, obj) < 0) { + goto done; + } + status = 0; + + done: + Py_DECREF(referent); + return status; +} + +/* Build a sorted ", "-joined string of the distinct type names of the memo + objects still flagged weak-only, for the dump() error. */ +static PyObject * +weak_only_type_names(PicklerObject *self) +{ + PyObject *names = PySet_New(NULL); + if (names == NULL) + return NULL; + + PyMemoTable *memo = self->memo; + if (memo && memo->mt_table) { + for (size_t i = 0; i < memo->mt_allocated; i++) { + PyObject *key = memo->mt_table[i].me_key; + if (key == NULL || memo->mt_table[i].me_value >= 0) + continue; /* absent or not flagged weak-only */ + PyObject *tname = PyType_GetFullyQualifiedName(Py_TYPE(key)); + if (tname == NULL) + goto error; + int r = PySet_Add(names, tname); + Py_DECREF(tname); + if (r < 0) + goto error; + } + } + + PyObject *lst = PySequence_List(names); + Py_CLEAR(names); + if (lst == NULL) + return NULL; + if (PyList_Sort(lst) < 0) { + Py_DECREF(lst); + return NULL; + } + PyObject *sep = PyUnicode_FromString(", "); + if (sep == NULL) { + Py_DECREF(lst); + return NULL; + } + PyObject *joined = PyUnicode_Join(sep, lst); + Py_DECREF(sep); + Py_DECREF(lst); + return joined; + + error: + Py_DECREF(names); + return NULL; +} + +/* `weak` is true only for save_weakref's immediate save of its referent; it + does not propagate to nested saves. All other callers go through save(). */ +static int +save_object(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save, + int weak) { PyTypeObject *type; PyObject *reduce_func = NULL; @@ -4596,7 +4757,11 @@ save(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save) /* Check the memo to see if it has the object. If so, generate a GET (or BINGET) opcode, instead of pickling the object once again. */ - if (PyMemoTable_Get(self->memo, obj)) { + Py_ssize_t *memo_value = PyMemoTable_Get(self->memo, obj); + if (memo_value != NULL) { + /* A strong reference anchors the object: no longer weak-only. */ + if (!weak) + memo_clear_weak_only(self, memo_value); return memo_get(st, self, obj); } @@ -4613,6 +4778,11 @@ save(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save) return -1; } + if (PyWeakref_CheckRefExact(obj)) { + status = save_weakref(st, self, obj); + goto done; + } + if (type == &PyDict_Type) { status = save_dict(st, self, obj); goto done; @@ -4755,6 +4925,14 @@ save(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save) if (status < 0) { _PyErr_FormatNote("when serializing %T object", obj); } + else if (!weak) { + /* 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. */ + Py_ssize_t *value = PyMemoTable_Get(self->memo, obj); + if (value != NULL) + memo_clear_weak_only(self, value); + } if (0) { error: @@ -4769,6 +4947,12 @@ save(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save) return status; } +static int +save(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save) +{ + return save_object(st, self, obj, pers_save, 0); +} + static PyObject * persistent_id(PyObject *self, PyObject *obj) { @@ -4814,8 +4998,22 @@ dump(PickleState *state, PicklerObject *self, PyObject *obj) self->framing = 1; } - if (save(state, self, obj, 0) < 0 || - _Pickler_Write(self, &stop_op, 1) < 0 || + if (save(state, self, obj, 0) < 0) + goto error; + + if (self->weak_only_count > 0) { + /* These referents would be collected right after loading. */ + PyObject *names = weak_only_type_names(self); + if (names == NULL) + goto error; + PyErr_Format(state->PicklingError, + "cannot pickle weak reference(s) whose referent(s) have " + "no strong reference in the pickle: %U", names); + Py_DECREF(names); + goto error; + } + + if (_Pickler_Write(self, &stop_op, 1) < 0 || _Pickler_CommitFrame(self) < 0) goto error; @@ -4854,6 +5052,7 @@ _pickle_Pickler_clear_memo_impl(PicklerObject *self) { if (self->memo) PyMemoTable_Clear(self->memo); + self->weak_only_count = 0; Py_RETURN_NONE; } @@ -4956,6 +5155,8 @@ Pickler_clear(PyObject *op) self->memo = NULL; PyMemoTable_Del(memo); } + /* The weak-only flags lived in the memo entries just freed. */ + self->weak_only_count = 0; return 0; } @@ -5143,7 +5344,7 @@ _pickle_PicklerMemoProxy_copy_impl(PicklerMemoProxyObject *self) if (key == NULL) { goto error; } - value = Py_BuildValue("nO", entry.me_value, entry.me_key); + value = Py_BuildValue("nO", MEMO_INDEX(entry.me_value), entry.me_key); if (value == NULL) { Py_DECREF(key); goto error;