gh-74876: Add copyreg.Proxy

Doc/library/copyreg.rst

@@ -39,6 +39,30 @@ Such constructors may be factory functions or class instances.
    object or subclass of :class:`pickle.Pickler` can also be used for
    declaring reduction functions.
 
+
+.. class:: Proxy(obj)
+
+   A wrapper for an object which is pickled weakly.
+   Reduction functions can use it to pickle objects
+   to which the reconstructed object only holds weak references.
+   For example, a weak reference is reduced to
+   ``(weakref.ref, (Proxy(referent),))``.
+
+   A proxy is pickled as the wrapped object *obj*,
+   which is unpickled in its place.
+   *obj* must also be strongly referenced elsewhere in the same pickle;
+   otherwise pickling raises :exc:`pickle.PicklingError`.
+
+   :func:`copy.copy` of a proxy returns *obj*,
+   and :func:`copy.deepcopy` returns the deep copy of *obj*.
+
+   .. attribute:: obj
+
+      The wrapped object.
+
+   .. versionadded:: next
+
+
 Example
 -------
 

Doc/library/weakref.rst

@@ -276,8 +276,10 @@ These have the same issues as the :meth:`WeakKeyDictionary.keyrefs` method.
 
    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.
+   A shallow copy contains the same elements as the original.
+   A deep copy refers to the deep copies of the elements,
+   so they must be deep-copied together with it;
+   its user-defined attributes are deep-copied.
 
    .. versionchanged:: next
       :class:`!WeakSet` can now be pickled and copied.
@@ -314,7 +316,9 @@ These have the same issues as the :meth:`WeakKeyDictionary.keyrefs` method.
    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.
+   A copy is an equal weak method bound to the same object.
+   A deep copy is bound to the deep copy of the object,
+   so it must be deep-copied together with it.
 
    .. versionadded:: 3.4
 

Doc/whatsnew/3.16.rst

@@ -193,6 +193,17 @@ concurrent.futures
   (Contributed by xzmeng and Serhiy Storchaka in :gh:`108518`.)
 
 
+copyreg
+-------
+
+* Add :class:`copyreg.Proxy`, a wrapper that is pickled as the wrapped object.
+  Reduction functions can use it to pickle objects
+  to which the reconstructed object only holds weak references.
+  The wrapped object must be strongly referenced elsewhere in the same pickle;
+  otherwise pickling raises :exc:`pickle.PicklingError`.
+  (Contributed by Serhiy Storchaka in :gh:`74876`.)
+
+
 csv
 ---
 

Lib/_weakrefset.py

@@ -3,6 +3,7 @@
 # by abc.py to load everything else at startup.
 
 from _weakref import ref
+from copyreg import Proxy
 from types import GenericAlias
 
 __all__ = ['WeakSet']
@@ -44,14 +45,13 @@ def __contains__(self, item):
         return wr in self.data
 
     def __reduce__(self):
-        # Pickle the elements as weak references; data and _remove are
-        # rebuilt by the reconstructor.  Copying shares the elements.
+        # The elements must be strongly referenced elsewhere in the pickle.
+        # data and _remove are rebuilt by the constructor.
         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)
+        return (self.__class__, ([Proxy(item) for item in self],), state)
 
     def add(self, item):
         self.data.add(ref(item, self._remove))
@@ -62,6 +62,8 @@ def clear(self):
     def copy(self):
         return self.__class__(self)
 
+    __copy__ = copy
+
     def pop(self):
         while True:
             try:
@@ -156,10 +158,3 @@ 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/copy.py

@@ -51,7 +51,7 @@ class instances).
 
 import types
 import weakref
-from copyreg import dispatch_table
+from copyreg import dispatch_table, Proxy
 
 class Error(Exception):
     pass
@@ -71,6 +71,8 @@ def copy(x):
         return x
     if cls in _copy_builtin_containers:
         return cls.copy(x)
+    if cls is Proxy:
+        return x.obj
 
 
     if issubclass(cls, type):
@@ -221,6 +223,10 @@ def _deepcopy_method(x, memo): # Copy instance methods
     return type(x)(x.__func__, deepcopy(x.__self__, memo))
 d[types.MethodType] = _deepcopy_method
 
+def _deepcopy_proxy(x, memo, deepcopy=deepcopy):
+    return deepcopy(x.obj, memo)
+d[Proxy] = _deepcopy_proxy
+
 del d
 
 def _keep_alive(x, memo):

Lib/copyreg.py

@@ -4,7 +4,7 @@
 C, not for instances of user-defined classes.
 """
 
-__all__ = ["pickle", "constructor",
+__all__ = ["pickle", "constructor", "Proxy",
            "add_extension", "remove_extension", "clear_extension_cache"]
 
 dispatch_table = {}
@@ -23,6 +23,36 @@ def constructor(object):
     if not callable(object):
         raise TypeError("constructors must be callable")
 
+class Proxy:
+    """Wrapper for an object that is pickled as a weak reference.
+
+    Pickling a Proxy pickles the wrapped object, which is unpickled in
+    its place.  Like the referent of a weak reference, the wrapped object
+    must also be strongly referenced elsewhere in the same pickle,
+    otherwise pickling fails.  This lets a reduction function pickle
+    objects that the reconstructed object will only refer to weakly.
+    """
+
+    __slots__ = ('obj',)
+
+    def __init__(self, obj):
+        self.obj = obj
+
+    def __repr__(self):
+        return f'{self.__class__.__qualname__}({self.obj!r})'
+
+class _Reducer:
+    # An object that is pickled and copied as the result of a call:
+    # reduce is a (callable, args) pair returned from __reduce__().
+
+    __slots__ = ('_reduce',)
+
+    def __init__(self, reduce):
+        self._reduce = reduce
+
+    def __reduce__(self):
+        return self._reduce
+
 # Example: provide pickling support for complex numbers.
 
 def pickle_complex(c):
@@ -41,6 +71,21 @@ def pickle_super(obj):
 
 pickle(super, pickle_super)
 
+from _weakref import ref as _ref
+
+# All dead weak references are interchangeable, so they are pickled as
+# references to this set.  It is unpickled as a new set which dies with
+# the unpickler's memo, and all dead references as the same reference.
+_dead_referent = set()
+
+def pickle_weakref(r):
+    obj = r()
+    if obj is None:
+        return _ref, (_dead_referent,)
+    return _ref, (Proxy(obj),)
+
+pickle(_ref, pickle_weakref)
+
 # Support for pickling new-style objects
 
 def _reconstructor(cls, base, state):

Lib/pickle.py

@@ -23,7 +23,7 @@
 """
 
 from types import FunctionType
-from copyreg import dispatch_table
+from copyreg import dispatch_table, Proxy
 from copyreg import _extension_registry, _inverted_registry, _extension_cache
 from itertools import batched
 from functools import partial
@@ -32,7 +32,6 @@
 from struct import pack, unpack
 import io
 import codecs
-import weakref
 import _compat_pickle
 
 __all__ = ["PickleError", "PicklingError", "UnpicklingError", "Pickler",
@@ -487,7 +486,7 @@ def __init__(self, file, protocol=None, *, fix_imports=True,
         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.
+        # left at the end of dump() is an error.  See save_proxy.
         self._weak_only = set()
         self.proto = int(protocol)
         self.bin = protocol >= 1
@@ -1286,39 +1285,20 @@ 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
+    def save_proxy(self, obj):
+        # A Proxy is pickled as its wrapped object, saved *weakly* (no
+        # anchoring, see save()).  The object is 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
+        # An object already in the memo is anchored or already tracked;
+        # otherwise it stays tagged until a strong save of it.
+        obj = obj.obj
+        i = id(obj)
+        anchored = i in self.memo
+        self.save(obj, weak=True)
+        if not anchored and i in self.memo:
+            self._weak_only.add(i)
+
+    dispatch[Proxy] = save_proxy
 
 
 # Unpickling machinery

Lib/test/pickletester.py

@@ -1866,6 +1866,17 @@ def test_weakref_reused_pickler(self):
         obj2, wr2 = pickle.loads(f.getvalue())
         self.assertIs(wr2(), obj2)
 
+    def test_proxy_without_strong_ref(self):
+        obj = Object()
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                self.assertRaises(pickle.PicklingError, self.dumps,
+                                  copyreg.Proxy(obj), proto)
+                # Another Proxy is not a strong reference.
+                self.assertRaises(pickle.PicklingError, self.dumps,
+                                  [copyreg.Proxy(obj), copyreg.Proxy(obj)],
+                                  proto)
+
     def test_bad_reduce_result(self):
         obj = REX([print, ()])
         for proto in protocols:
@@ -3684,6 +3695,23 @@ def test_dead_weakref(self):
                 support.gc_collect()
                 self.assertIsNone(a())
 
+    def test_proxy(self):
+        # A Proxy is unpickled as the wrapped object.
+        obj = Object()
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                x, y = self.loads(self.dumps([obj, copyreg.Proxy(obj)], proto))
+                self.assertIs(type(x), Object)
+                self.assertIs(y, x)
+                # Identity holds regardless of pickling order.
+                x, y = self.loads(self.dumps([copyreg.Proxy(obj), obj], proto))
+                self.assertIs(type(x), Object)
+                self.assertIs(y, x)
+                # Objects which are not memoized need no strong reference.
+                x = self.loads(self.dumps([copyreg.Proxy(None),
+                                           copyreg.Proxy(42)], proto))
+                self.assertEqual(x, [None, 42])
+
     def test_newobj_overridden_new(self):
         # Test that Python class with C implemented __new__ is pickleable
         for proto in protocols:

Lib/test/test_copy.py

@@ -869,6 +869,24 @@ def test_copy_weakref(self):
     def test_deepcopy_weakref(self):
         self._check_weakref(copy.deepcopy)
 
+    def test_copy_proxy(self):
+        class C:
+            pass
+        x = C()
+        self.assertIs(copy.copy(copyreg.Proxy(x)), x)
+
+    def test_deepcopy_proxy(self):
+        # A Proxy is replaced by the deep copy of the wrapped object.
+        class C:
+            pass
+        x = C()
+        y, z = copy.deepcopy([x, copyreg.Proxy(x)])
+        self.assertIsNot(y, x)
+        self.assertIs(z, y)
+        z, y = copy.deepcopy([copyreg.Proxy(x), x])
+        self.assertIsNot(y, x)
+        self.assertIs(z, y)
+
     def _check_copy_weakdict(self, _dicttype):
         class C(object):
             pass