python · david-caro · Jul 8, 2020 · Aug 8, 2020 · Jul 11, 2022 · Jan 7, 2023
@@ -647,6 +647,23 @@ The :mod:`functools` module defines the following functions:
    on the wrapper function). :exc:`AttributeError` is still raised if the
    wrapper function itself is missing any attributes named in *updated*.
 
+   When assigning the ``__annotations__`` attribute (when it's part of the
+   ``assigned`` list parameter), it will be populated only if missing some
+   annotations as follows:
+
+   * If the wrapper function defines it's own return type and parameter
-   * If the wrapper function defines it's own return type and parameter
+   * If the wrapper function defines its own return type and parameter
-   * If the wrapper function defines it's own return type and parameter
+   * If the wrapper function defines its own return type and parameter
+     annotations, those will be preserved.
+
+   * If the wrapper function does not annotate the return type, it will be
+     copied from the wrapped if any.
-     copied from the wrapped if any.
+     copied from the wrapped function, if it specifies a return type.
-     copied from the wrapped if any.
+     copied from the wrapped function, if it specifies a return type.
+
+   * If the wrapped function does not annotate any parameters, the annotations
+     for them will be copied over from the wrapped if any.
-     for them will be copied over from the wrapped if any.
+     for them will be copied over from the wrapped function, if it has any.
-     for them will be copied over from the wrapped if any.
+     for them will be copied over from the wrapped function, if it has any.
+
+   That ensures that any custom annotations from the wrapper will not be
+   overridden by the wrapped, allowing to change the annotated types of
+   the wrapped function.
+
    .. versionadded:: 3.2
       Automatic addition of the ``__wrapped__`` attribute.
 

@@ -46,14 +46,45 @@ 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)
+
+       There's a special treatment for the `__annotations__` attribute:
+       * If the wrapper defines a return type, then that will be used, if not,
+         the one from the wrapped function will be used if any.
+       * If the wrapper defines any parameter types, those will be used, if
+         not, the ones from the wrapped function will be used if any.
     """
     for attr in assigned:
         try:
             value = getattr(wrapped, attr)
         except AttributeError:
             pass
         else:
+            if attr == '__annotations__':
+                continue
             setattr(wrapper, attr, value)
+
+    if '__annotations__' in assigned:
+        try:
+            # Issue #41231: copy the annotations from wrapped, only if they are
+            #  not already defined in the wrapper
+            if not wrapper.__annotations__:
+                wrapper.__annotations__ = wrapped.__annotations__
+            # if the wrapper does not annotate the parameters, copy their
+            # annotations over from the wrapped
+            elif ('return' in wrapper.__annotations__
+                  and len(wrapper.__annotations__) == 1):
+                wrapper.__annotations__ = {
+                    **wrapped.__annotations__,
+                    **wrapper.__annotations__,
+                }
+            # if the wrapper does not annotate the return type, copy the
+            # annotation from the wrapped
+            elif ('return' not in wrapper.__annotations__
+                  and 'return' in wrapped.__annotations__):
+                wrapper.__annotations__['return'] = wrapped.__annotations__['return']
+        except AttributeError:
+            pass
+
     for attr in updated:
         getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
     # Issue #17482: set __wrapped__ last so we don't inadvertently copy it

diff --git a/Lib/test/ann_module.py b/Lib/test/ann_module.py
@@ -59,4 +59,22 @@ def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
     return wrapper
 
+def dec_with_return_type(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs) -> bool:
+        return bool(func(*args, **kwargs))
+    return wrapper
+
+def dec_with_arg_types(func):
+    @wraps(func)
+    def wrapper(self, custom_arg: int, *args, **kwargs):
+        return func(self, custom_arg, *args, **kwargs)
+    return wrapper
+
+def dec_with_args_and_return_types(func):
+    @wraps(func)
+    def wrapper(self, custom_arg: int, *args, **kwargs) -> bool:
+        return bool(func(self, custom_arg, *args, **kwargs))
+    return wrapper
+
 u: int | float
@@ -599,10 +599,17 @@ class TestUpdateWrapper(unittest.TestCase):
 
     def check_wrapper(self, wrapper, wrapped,
                       assigned=functools.WRAPPER_ASSIGNMENTS,
-                      updated=functools.WRAPPER_UPDATES):
+                      updated=functools.WRAPPER_UPDATES,
+                      expected_annotations=None,
+                      ):
         # Check attributes were assigned
         for name in assigned:
-            self.assertIs(getattr(wrapper, name), getattr(wrapped, name))
+            if expected_annotations is not None and name == '__annotations__':
+                self.assertEqual(expected_annotations, wrapper.__annotations__)
+            elif name == '__annotations__':
+                self.assertEqual(getattr(wrapper, name), getattr(wrapped, name))
+            else:
+                self.assertIs(getattr(wrapper, name), getattr(wrapped, name))
         # Check attributes were updated
         for name in updated:
             wrapper_attr = getattr(wrapper, name)
@@ -629,13 +636,14 @@ def wrapper(b:'This is the prior annotation'):
 
     def test_default_update(self):
         wrapper, f = self._default_update()
-        self.check_wrapper(wrapper, f)
+        self.check_wrapper(wrapper, f,
+                           expected_annotations={'b': 'This is the prior annotation'})
         self.assertIs(wrapper.__wrapped__, f)
         self.assertEqual(wrapper.__name__, 'f')
         self.assertEqual(wrapper.__qualname__, f.__qualname__)
         self.assertEqual(wrapper.attr, 'This is also a test')
-        self.assertEqual(wrapper.__annotations__['a'], 'This is a new annotation')
-        self.assertNotIn('b', wrapper.__annotations__)
+        self.assertEqual(wrapper.__annotations__['b'], 'This is the prior annotation')
+        self.assertNotIn('a', wrapper.__annotations__)
 
     @unittest.skipIf(sys.flags.optimize >= 2,
                      "Docstrings are omitted with -O2 and above")
@@ -1638,7 +1646,10 @@ def f(zomg: 'zomg_annotation'):
             return 42
         g = self.module.lru_cache()(f)
         for attr in self.module.WRAPPER_ASSIGNMENTS:
-            self.assertEqual(getattr(g, attr), getattr(f, attr))
+            try:
+                self.assertEqual(getattr(g, attr), getattr(f, attr))
+            except AttributeError:
+                pass
 
     @threading_helper.requires_working_threading()
     def test_lru_cache_threaded(self):

diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
@@ -3038,8 +3038,9 @@ def bar(self, a, b):
         self.assertEqual(self.signature(Foo.bar, follow_wrapped=False),
                          ((('args', ..., ..., "var_positional"),
                            ('kwargs', ..., ..., "var_keyword")),
-                          ...)) # functools.wraps will copy __annotations__
-                                # from "func" to "wrapper", hence no
+                          int)) # functools.wraps will copy __annotations__
+                                # from "func" to "wrapper" only if wrapper does
+                                # not have some already, hence the
                                 # return_annotation
 
         self.assertEqual(self.signature(bar),

@@ -5009,6 +5009,33 @@ def test_get_type_hints_wrapped_decoratored_func(self):
         self.assertEqual(gth(ForRefExample.func), expects)
         self.assertEqual(gth(ForRefExample.nested), expects)
 
+    def test_get_type_hints_wrapped_decoratored_func_only_with_return_type(self):
+        expects = {'arg1': str, 'return': bool}
+
+        @ann_module.dec_with_return_type
+        def func(arg1: str) -> int:
+            return 0
+
+        self.assertEqual(gth(func), expects)
+
+    def test_get_type_hints_wrapped_decoratored_func_only_with_arg_types(self):
+        expects = {'custom_arg': int, 'return': int}
+
+        @ann_module.dec_with_arg_types
+        def func(arg1: str) -> int:
+            return 0
+
+        self.assertEqual(gth(func), expects)
+
+    def test_get_type_hints_wrapped_decoratored_func_with_args_and_return_type(self):
+        expects = {'custom_arg': int, 'return': bool}
+
+        @ann_module.dec_with_args_and_return_types
+        def func(arg1: str) -> int:
+            return 0
+
+        self.assertEqual(gth(func), expects)
+
     def test_get_type_hints_annotated(self):
         def foobar(x: List['X']): ...
         X = Annotated[int, (1, 10)]

diff --git a/Misc/ACKS b/Misc/ACKS
@@ -281,6 +281,7 @@ Brett Cannon
 Joshua Cannon
 Tristan Carel
 Mike Carlton
+David Caro
 Pierre Carrier
 Terry Carroll
 Edward Catmur

diff --git a/Misc/NEWS.d/next/Library/2020-07-08-11-00-41.bpo-41231.b_-THN.rst b/Misc/NEWS.d/next/Library/2020-07-08-11-00-41.bpo-41231.b_-THN.rst
@@ -0,0 +1,3 @@
+``update_wrapper``: Don't override the ``__annotations__`` of the wrapper
+function if it has any as it might have a different signature than the
+wrapped one, and using the wrapped might mask the actual types.