gh-153612: Add copyreg.copy() and copyreg.deepcopy()

Doc/library/copy.rst

@@ -87,6 +87,14 @@ pickling.  See the description of module :mod:`pickle` for information on these
 methods.  In fact, the :mod:`!copy` module uses the registered
 pickle functions from the :mod:`copyreg` module.
 
+Copying of instances of a specific type can also be customized
+without modifying the class:
+see :func:`copyreg.copy` and :func:`copyreg.deepcopy`.
+
+.. versionchanged:: next
+   Added support for functions registered
+   with :func:`copyreg.copy` and :func:`copyreg.deepcopy`.
+
 .. index::
    single: __copy__() (copy protocol)
    single: __deepcopy__() (copy protocol)

Doc/library/copyreg.rst

@@ -1,8 +1,10 @@
-:mod:`!copyreg` --- Register :mod:`!pickle` support functions
-=============================================================
+.. _copyreg-register-pickle-support-functions:
+
+:mod:`!copyreg` --- Register :mod:`!pickle` and :mod:`!copy` support functions
+==============================================================================
 
 .. module:: copyreg
-   :synopsis: Register pickle support functions.
+   :synopsis: Register pickle and copy support functions.
 
 **Source code:** :source:`Lib/copyreg.py`
 
@@ -12,10 +14,12 @@
 
 --------------
 
-The :mod:`!copyreg` module offers a way to define functions used while pickling
-specific objects.  The :mod:`pickle` and :mod:`copy` modules use those functions
-when pickling/copying those objects.  The module provides configuration
-information about object constructors which are not classes.
+The :mod:`!copyreg` module offers a way to define functions
+used while pickling and copying specific objects.
+The :mod:`pickle` and :mod:`copy` modules use those functions
+when pickling/copying those objects.
+The module provides configuration information
+about object constructors which are not classes.
 Such constructors may be factory functions or class instances.
 
 
@@ -39,6 +43,42 @@ Such constructors may be factory functions or class instances.
    object or subclass of :class:`pickle.Pickler` can also be used for
    declaring reduction functions.
 
+
+.. function:: copy(type, function)
+
+   Declares that *function* should be used as the shallow copy function
+   for objects of type *type*.
+   *function* is called with the object as its only argument
+   and must return the copy,
+   like the :meth:`~object.__copy__` method.
+   Registration is by exact type:
+   it does not apply to subclasses of *type*.
+
+   :func:`copy.copy` uses the registered function
+   in preference to the :meth:`~object.__copy__` method
+   and the pickle interfaces.
+   Unlike a reduction function registered with :func:`pickle`,
+   it affects only shallow copying.
+   The registered functions are stored in ``copy_dispatch_table``,
+   the same table that holds the handlers for the built-in
+   container types, so registering a function for a built-in
+   container type overrides its default copying.
+
+   .. versionadded:: next
+
+
+.. function:: deepcopy(type, function)
+
+   Like :func:`copy`, but registers the deep copy function
+   used by :func:`copy.deepcopy`.
+   *function* is called with the object and the memo dictionary
+   as its two arguments,
+   like the :meth:`~object.__deepcopy__` method.
+   The registered functions are stored in ``deepcopy_dispatch_table``.
+
+   .. versionadded:: next
+
+
 Example
 -------
 

Lib/copy.py

@@ -51,7 +51,7 @@ class instances).
 
 import types
 import weakref
-from copyreg import dispatch_table
+from copyreg import dispatch_table, copy_dispatch_table, deepcopy_dispatch_table
 
 class Error(Exception):
     pass
@@ -69,9 +69,10 @@ def copy(x):
 
     if cls in _copy_atomic_types:
         return x
-    if cls in _copy_builtin_containers:
-        return cls.copy(x)
 
+    copier = copy_dispatch_table.get(cls)
+    if copier is not None:
+        return copier(x)
 
     if issubclass(cls, type):
         # treat it as a regular class:
@@ -105,7 +106,13 @@ def copy(x):
           types.BuiltinFunctionType, types.EllipsisType,
           types.NotImplementedType, types.FunctionType, types.CodeType,
           weakref.ref, super})
-_copy_builtin_containers = frozenset({list, dict, set, bytearray})
+
+# The handlers for the built-in container types are registered in the
+# public copyreg.copy_dispatch_table.
+copy_dispatch_table.setdefault(list, list.copy)
+copy_dispatch_table.setdefault(dict, dict.copy)
+copy_dispatch_table.setdefault(set, set.copy)
+copy_dispatch_table.setdefault(bytearray, bytearray.copy)
 
 def deepcopy(x, memo=None):
     """Deep copy operation on arbitrary Python objects.
@@ -126,7 +133,7 @@ def deepcopy(x, memo=None):
         if y is not None:
             return y
 
-    copier = _deepcopy_dispatch.get(cls)
+    copier = deepcopy_dispatch_table.get(cls)
     if copier is not None:
         y = copier(x, memo)
     else:
@@ -166,17 +173,14 @@ def deepcopy(x, memo=None):
           int, float, bool, complex, bytes, str, types.CodeType, type, range,
           types.BuiltinFunctionType, types.FunctionType, weakref.ref, property})
 
-_deepcopy_dispatch = d = {}
-
-
 def _deepcopy_list(x, memo, deepcopy=deepcopy):
     y = []
     memo[id(x)] = y
     append = y.append
     for a in x:
         append(deepcopy(a, memo))
     return y
-d[list] = _deepcopy_list
+deepcopy_dispatch_table.setdefault(list, _deepcopy_list)
 
 def _deepcopy_tuple(x, memo, deepcopy=deepcopy):
     y = [deepcopy(a, memo) for a in x]
@@ -193,15 +197,15 @@ def _deepcopy_tuple(x, memo, deepcopy=deepcopy):
     else:
         y = x
     return y
-d[tuple] = _deepcopy_tuple
+deepcopy_dispatch_table.setdefault(tuple, _deepcopy_tuple)
 
 def _deepcopy_dict(x, memo, deepcopy=deepcopy):
     y = {}
     memo[id(x)] = y
     for key, value in x.items():
         y[deepcopy(key, memo)] = deepcopy(value, memo)
     return y
-d[dict] = _deepcopy_dict
+deepcopy_dispatch_table.setdefault(dict, _deepcopy_dict)
 
 def _deepcopy_frozendict(x, memo, deepcopy=deepcopy):
     y = {}
@@ -216,13 +220,12 @@ def _deepcopy_frozendict(x, memo, deepcopy=deepcopy):
     except KeyError:
         pass
     return frozendict(y)
-d[frozendict] = _deepcopy_frozendict
+deepcopy_dispatch_table.setdefault(frozendict, _deepcopy_frozendict)
 
 def _deepcopy_method(x, memo): # Copy instance methods
     return type(x)(x.__func__, deepcopy(x.__self__, memo))
-d[types.MethodType] = _deepcopy_method
+deepcopy_dispatch_table.setdefault(types.MethodType, _deepcopy_method)
 
-del d
 
 def _keep_alive(x, memo):
     """Keeps a reference to the object x in the memo.

Lib/copyreg.py

@@ -1,10 +1,11 @@
-"""Helper to provide extensibility for pickle.
+"""Helper to provide extensibility for pickle and copy.
 
-This is only useful to add pickle support for extension types defined in
-C, not for instances of user-defined classes.
+This is only useful to add support for types that cannot be modified,
+such as extension types defined in C, not for instances of user-defined
+classes.
 """
 
-__all__ = ["pickle", "constructor",
+__all__ = ["pickle", "constructor", "copy", "deepcopy",
            "add_extension", "remove_extension", "clear_extension_cache"]
 
 dispatch_table = {}
@@ -23,6 +24,22 @@ def constructor(object):
     if not callable(object):
         raise TypeError("constructors must be callable")
 
+# Support for the copy module
+
+copy_dispatch_table = {}
+
+def copy(ob_type, copy_function):
+    if not callable(copy_function):
+        raise TypeError("copy functions must be callable")
+    copy_dispatch_table[ob_type] = copy_function
+
+deepcopy_dispatch_table = {}
+
+def deepcopy(ob_type, deepcopy_function):
+    if not callable(deepcopy_function):
+        raise TypeError("deepcopy functions must be callable")
+    deepcopy_dispatch_table[ob_type] = deepcopy_function
+
 # Example: provide pickling support for complex numbers.
 
 def pickle_complex(c):

Lib/test/test_copy.py

@@ -8,6 +8,7 @@
 
 import unittest
 from test import support
+from test.support import script_helper
 
 order_comparisons = le, lt, ge, gt
 equality_comparisons = eq, ne
@@ -55,6 +56,67 @@ def pickle_C(obj):
         self.assertEqual(type(y), C)
         self.assertEqual(y.foo, x.foo)
 
+    def test_copy_dispatch_table(self):
+        class C:
+            def __copy__(self):
+                return "__copy__"
+        x = C()
+        copyreg.copy(C, lambda obj: (obj, "registered"))
+        try:
+            self.assertEqual(copy.copy(x), (x, "registered"))
+        finally:
+            del copyreg.copy_dispatch_table[C]
+        self.assertEqual(copy.copy(x), "__copy__")
+
+    def test_copy_dispatch_table_exact_type(self):
+        class C:
+            def __copy__(self):
+                return "__copy__"
+        class D(C):
+            pass
+        copyreg.copy(C, lambda obj: "registered")
+        try:
+            self.assertEqual(copy.copy(C()), "registered")
+            self.assertEqual(copy.copy(D()), "__copy__")
+        finally:
+            del copyreg.copy_dispatch_table[C]
+
+    def test_copy_dispatch_table_independent(self):
+        # A registered copy function affects neither deepcopy nor pickling.
+        class C:
+            pass
+        copyreg.copy(C, lambda obj: "copy")
+        try:
+            x = C()
+            self.assertEqual(copy.copy(x), "copy")
+            y = copy.deepcopy(x)
+            self.assertIsInstance(y, C)
+        finally:
+            del copyreg.copy_dispatch_table[C]
+
+    def test_copy_dispatch_table_builtin(self):
+        # The registry is the same table that holds the built-in
+        # handlers, so they can be overridden.
+        orig = copyreg.copy_dispatch_table[list]
+        copyreg.copy(list, lambda obj: "registered")
+        try:
+            self.assertEqual(copy.copy([1]), "registered")
+        finally:
+            copyreg.copy_dispatch_table[list] = orig
+        self.assertEqual(copy.copy([1]), [1])
+
+    def test_copy_register_before_import(self):
+        # Registrations made before importing the copy module are
+        # preserved when it registers the built-in handlers.
+        code = """if True:
+            import copyreg
+            copyreg.copy(list, lambda obj: "registered")
+            import copy
+            print(copy.copy([1]))
+            """
+        rc, out, err = script_helper.assert_python_ok('-c', code)
+        self.assertEqual(out.strip(), b"registered")
+
     def test_copy_reduce_ex(self):
         class C(object):
             def __reduce_ex__(self, proto):
@@ -308,6 +370,73 @@ def __deepcopy__(self, memo=None):
         self.assertEqual(y.__class__, x.__class__)
         self.assertEqual(y.foo, x.foo)
 
+    def test_deepcopy_dispatch_table(self):
+        class C:
+            def __deepcopy__(self, memo):
+                return "__deepcopy__"
+        x = C()
+        def copier(obj, memo):
+            self.assertIsInstance(memo, dict)
+            return (obj, "registered")
+        copyreg.deepcopy(C, copier)
+        try:
+            self.assertEqual(copy.deepcopy(x), (x, "registered"))
+        finally:
+            del copyreg.deepcopy_dispatch_table[C]
+        self.assertEqual(copy.deepcopy(x), "__deepcopy__")
+
+    def test_deepcopy_dispatch_table_memo(self):
+        class C:
+            pass
+        calls = []
+        def copier(obj, memo):
+            calls.append(obj)
+            return C()
+        copyreg.deepcopy(C, copier)
+        try:
+            x = C()
+            y = copy.deepcopy([x, x])
+            self.assertIs(y[0], y[1])
+            self.assertEqual(len(calls), 1)
+        finally:
+            del copyreg.deepcopy_dispatch_table[C]
+
+    def test_deepcopy_dispatch_table_exact_type(self):
+        class C:
+            def __deepcopy__(self, memo):
+                return "__deepcopy__"
+        class D(C):
+            pass
+        copyreg.deepcopy(C, lambda obj, memo: "registered")
+        try:
+            self.assertEqual(copy.deepcopy(C()), "registered")
+            self.assertEqual(copy.deepcopy(D()), "__deepcopy__")
+        finally:
+            del copyreg.deepcopy_dispatch_table[C]
+
+    def test_deepcopy_dispatch_table_builtin(self):
+        # The registry is the same table that holds the built-in
+        # handlers, so they can be overridden.
+        orig = copyreg.deepcopy_dispatch_table[list]
+        copyreg.deepcopy(list, lambda obj, memo: "registered")
+        try:
+            self.assertEqual(copy.deepcopy([1]), "registered")
+        finally:
+            copyreg.deepcopy_dispatch_table[list] = orig
+        self.assertEqual(copy.deepcopy([1]), [1])
+
+    def test_deepcopy_register_before_import(self):
+        # Registrations made before importing the copy module are
+        # preserved when it registers the built-in handlers.
+        code = """if True:
+            import copyreg
+            copyreg.deepcopy(list, lambda obj, memo: "registered")
+            import copy
+            print(copy.deepcopy([1]))
+            """
+        rc, out, err = script_helper.assert_python_ok('-c', code)
+        self.assertEqual(out.strip(), b"registered")
+
     def test_deepcopy_registry(self):
         class C(object):
             def __new__(cls, foo):
@@ -1010,7 +1139,11 @@ class C:
 
 class MiscTestCase(unittest.TestCase):
     def test__all__(self):
-        support.check__all__(self, copy, not_exported={"dispatch_table", "error"})
+        support.check__all__(self, copy,
+                             not_exported={"dispatch_table",
+                                           "copy_dispatch_table",
+                                           "deepcopy_dispatch_table",
+                                           "error"})
 
 def global_foo(x, y): return x+y