gh-124342: Wrap __annotate__ functions in functools.update_wrapper

Doc/library/functools.rst

@@ -638,18 +638,21 @@ The :mod:`functools` module defines the following functions:
    .. versionadded:: 3.8
 
 
-.. function:: update_wrapper(wrapper, wrapped, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES)
+.. function:: update_wrapper(wrapper, wrapped, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES, *, delegated=WRAPPER_DELEGATIONS)
 
    Update a *wrapper* function to look like the *wrapped* function. The optional
-   arguments are tuples to specify which attributes of the original function are
-   assigned directly to the matching attributes on the wrapper function and which
+   arguments are tuples to specify how particular attributes are handled.
+
+   The attributes in *assigned* are assigned directly to the matching attributes
+   on the wrapper function. The default value is the module-level constant
+   ``WRAPPER_ASSIGNMENTS``, which contains ``__module__``, ``__name__``, ``__qualname__``,
+   ``__type_params__``, and ``__doc__``. For names in *updated*,
    attributes of the wrapper function are updated with the corresponding attributes
-   from the original function. The default values for these arguments are the
-   module level constants ``WRAPPER_ASSIGNMENTS`` (which assigns to the wrapper
-   function's ``__module__``, ``__name__``, ``__qualname__``, ``__annotations__``,
-   ``__type_params__``, and ``__doc__``, the documentation string)
-   and ``WRAPPER_UPDATES`` (which
-   updates the wrapper function's ``__dict__``, i.e. the instance dictionary).
+   from the original function. It defaults to ``WRAPPER_UPDATES``, which contains
+   ``__dict__``, i.e. the instance dictionary. An attribute-specific delegation
+   mechanism is used for attributes in *delegated*. The default value,
+   ``WRAPPER_DELEGATIONS``, contains only ``__annotate__``, for which a wrapper
+   :term:`annotate function` is generated. Other attributes cannot be added to *delegated*.
 
    To allow access to the original function for introspection and other purposes
    (e.g. bypassing a caching decorator such as :func:`lru_cache`), this function
@@ -681,12 +684,16 @@ The :mod:`functools` module defines the following functions:
    .. versionchanged:: 3.12
       The ``__type_params__`` attribute is now copied by default.
 
+   .. versionchanged:: 3.14
+      The ``__annotations__`` attribute is no longer copied by default. Instead,
+      the ``__annotate__`` attribute is delegated. See :pep:`749`.
+
 
-.. decorator:: wraps(wrapped, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES)
+.. decorator:: wraps(wrapped, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES, *, delegated=WRAPPER_DELEGATIONS)
 
    This is a convenience function for invoking :func:`update_wrapper` as a
    function decorator when defining a wrapper function.  It is equivalent to
-   ``partial(update_wrapper, wrapped=wrapped, assigned=assigned, updated=updated)``.
+   ``partial(update_wrapper, wrapped=wrapped, assigned=assigned, updated=updated, delegated=delegated)``.
    For example::
 
       >>> from functools import wraps

Doc/whatsnew/3.14.rst

@@ -168,6 +168,14 @@ Added support for converting any objects that have the
 (Contributed by Serhiy Storchaka in :gh:`82017`.)
 
 
+functools
+---------
+
+* :func:`functools.update_wrapper` and :func:`functools.wraps` now have
+  special support for wrapping :term:`annotate functions <annotate function>`.
+  A new parameter *delegated* was added to both. (Contributed by Jelle Zijlstra
+  in :gh:`124342`.)
+
 http
 ----
 

Lib/functools.py

@@ -32,12 +32,15 @@
 # wrapper functions that can handle naive introspection
 
 WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__qualname__', '__doc__',
-                       '__annotate__', '__type_params__')
+                       '__type_params__')
 WRAPPER_UPDATES = ('__dict__',)
+WRAPPER_DELEGATIONS = ('__annotate__',)
 def update_wrapper(wrapper,
                    wrapped,
                    assigned = WRAPPER_ASSIGNMENTS,
-                   updated = WRAPPER_UPDATES):
+                   updated = WRAPPER_UPDATES,
+                   *,
+                   delegated = WRAPPER_DELEGATIONS):
     """Update a wrapper function to look like the wrapped function
 
        wrapper is the function to be updated
@@ -48,6 +51,11 @@ def update_wrapper(wrapper,
        updated is a tuple naming the attributes of the wrapper that
        are updated with the corresponding attribute from the wrapped
        function (defaults to functools.WRAPPER_UPDATES)
+       delegated is a tuple naming attributes of the wrapper for which
+       resolution should be delegated to the wrapped object using custom
+       logic (defaults to functools.WRAPPER_DELEGATIONS). Only the
+       __annotate__ attribute is supported; any other attribute will raise
+       an error.
     """
     for attr in assigned:
         try:
@@ -58,6 +66,20 @@ def update_wrapper(wrapper,
             setattr(wrapper, attr, value)
     for attr in updated:
         getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
+    for attr in delegated:
+        if attr == "__annotate__":
+            def __annotate__(format):
+                if format == 1:  # VALUE
+                    return getattr(wrapped, "__annotations__", {})
+                get_annotate_function = _get_get_annotate_function()
+                call_annotate_function = _get_call_annotate_function()
+                annotate_function = get_annotate_function(wrapped)
+                if annotate_function is None:
+                    return {}
+                return call_annotate_function(annotate_function, format=format)
+            wrapper.__annotate__ = __annotate__
+        else:
+            raise ValueError(f"Unsupported delegated attribute {attr!r}")
     # Issue #17482: set __wrapped__ last so we don't inadvertently copy it
     # from the wrapped function when updating __dict__
     wrapper.__wrapped__ = wrapped
@@ -66,7 +88,9 @@ def update_wrapper(wrapper,
 
 def wraps(wrapped,
           assigned = WRAPPER_ASSIGNMENTS,
-          updated = WRAPPER_UPDATES):
+          updated = WRAPPER_UPDATES,
+          *,
+          delegated = WRAPPER_DELEGATIONS):
     """Decorator factory to apply update_wrapper() to a wrapper function
 
        Returns a decorator that invokes update_wrapper() with the decorated
@@ -76,7 +100,7 @@ def wraps(wrapped,
        update_wrapper().
     """
     return partial(update_wrapper, wrapped=wrapped,
-                   assigned=assigned, updated=updated)
+                   assigned=assigned, updated=updated, delegated=delegated)
 
 
 ################################################################################
@@ -1048,3 +1072,25 @@ def __get__(self, instance, owner=None):
         return val
 
     __class_getitem__ = classmethod(GenericAlias)
+
+
+################################################################################
+### internal helpers
+################################################################################
+
+@cache
+def _get_get_annotations():
+    from annotationlib import get_annotations
+    return get_annotations
+
+
+@cache
+def _get_get_annotate_function():
+    from annotationlib import get_annotate_function
+    return get_annotate_function
+
+
+@cache
+def _get_call_annotate_function():
+    from annotationlib import call_annotate_function
+    return call_annotate_function

Lib/test/test_functools.py

@@ -1,4 +1,5 @@
 import abc
+import annotationlib
 import builtins
 import collections
 import collections.abc
@@ -747,20 +748,39 @@ def wrapper(*args): pass
 
         functools.update_wrapper(wrapper, inner)
         self.assertEqual(wrapper.__annotations__, {'x': int})
-        self.assertIs(wrapper.__annotate__, inner.__annotate__)
+        self.check_annotate_matches(wrapper, inner)
 
         def with_forward_ref(x: undefined): pass
         def wrapper(*args): pass
 
         functools.update_wrapper(wrapper, with_forward_ref)
 
-        self.assertIs(wrapper.__annotate__, with_forward_ref.__annotate__)
+        # VALUE raises NameError
+        self.check_annotate_matches(wrapper, with_forward_ref, skip=(annotationlib.Format.VALUE,))
         with self.assertRaises(NameError):
             wrapper.__annotations__
 
         undefined = str
         self.assertEqual(wrapper.__annotations__, {'x': undefined})
 
+    def test_update_wrapper_with_modified_annotations(self):
+        def inner(x: int): pass
+        def wrapper(*args): pass
+
+        inner.__annotations__["x"] = str
+        functools.update_wrapper(wrapper, inner)
+        self.assertEqual(wrapper.__annotations__, {'x': str})
+        self.check_annotate_matches(wrapper, inner)
+
+    def check_annotate_matches(self, wrapper, wrapped, skip=()):
+        for format in annotationlib.Format:
+            if format in skip:
+                continue
+            with self.subTest(format=format):
+                wrapper_annos = annotationlib.get_annotations(wrapper, format=format)
+                wrapped_annos = annotationlib.get_annotations(wrapped, format=format)
+                self.assertEqual(wrapper_annos, wrapped_annos)
+
 
 class TestWraps(TestUpdateWrapper):
 

Misc/NEWS.d/next/Library/2024-09-23-08-25-48.gh-issue-124342.7QzKqp.rst

@@ -0,0 +1,2 @@
+Add custom support for wrapping ``__annotate__`` functions to
+:func:`functools.update_wrapper` and :func:`functools.wraps`.