Repository: smarie/python-makefun

Issue:
Forward references in annotations with a signature changing decorator
`typing.get_type_hints` uses the `__wrapped__` attribute for finding the namespace of the original (undecorated) method was defined in. See [here](https://github.com/python/cpython/blob/605e9c66ad367b54a847f9fc65447a071742f554/Lib/typing.py#L2314-L2315) for the implementation in the standard library. This allows dereferencing of forward references, e.g.:

**foo.py:**
```python
def foo(bar: "Bar", baz: str):
    pass
class Bar:
    pass
```

**othermodule.py:**
```python
from foo import foo
import makefun, typing
@makefun.wraps(foo, remove_args=["baz"])
def decorated(*args, **kwargs):
    return foo(*args, **kwargs, baz="baz")

typing.get_type_hints(foo)  # fine
typing.get_type_hints(decorated)  # NameError
```

This then fails with:
```
$ PYTONPATH=.:$PYTHONPATH python ./mymodule.py
Traceback (most recent call last):
  File "/private/tmp/asdf/./mymodule.py", line 7, in <module>
    typing.get_type_hints(decorated)
  File "/Users/lucaswiman/.pyenv/versions/3.9.1/lib/python3.9/typing.py", line 1442, in get_type_hints
    value = _eval_type(value, globalns, localns)
  File "/Users/lucaswiman/.pyenv/versions/3.9.1/lib/python3.9/typing.py", line 277, in _eval_type
    return t._evaluate(globalns, localns, recursive_guard)
  File "/Users/lucaswiman/.pyenv/versions/3.9.1/lib/python3.9/typing.py", line 533, in _evaluate
    eval(self.__forward_code__, globalns, localns),
  File "<string>", line 1, in <module>
NameError: name 'Bar' is not defined
```

## Things that did not work
1. Setting `__wrapped__` to the original function.  This gets rid of the `NameError`, but also changes the signature back to its unaltered form.
2. Setting `__globals__`. This is a readonly attribute. Also this might mess up execution of the decorated method. 

## Possible approach
If there were some way to define a method with the right signature and the right `__globals__`, then you could set the `__wrapped__` attribute to that method and things would work correctly.

Discussion:

Patch:
diff --git a/docs/api_reference.md b/docs/api_reference.md
index 2300796..fbc6040 100644
--- a/docs/api_reference.md
+++ b/docs/api_reference.md
@@ -166,7 +166,7 @@ Comparison with `@with_signature`: `@wraps(f)` is equivalent to
 
 In other words, as opposed to `@with_signature`, the metadata (doc, module name, etc.) is provided by the wrapped `wrapped_fun`, so that the created function seems to be identical (except possiblyfor the signature). Note that all options in `with_signature` can still be overrided using parameters of `@wraps`.
 
-If the signature is *not* modified through `new_sig`, `remove_args`, `append_args` or `prepend_args`, the additional `__wrapped__` attribute  on the created function, to stay consistent with the `functools.wraps` behaviour.
+The additional `__wrapped__` attribute is added on the created function, to stay consistent with the `functools.wraps` behaviour. If the signature is modified through `new_sig`, `remove_args`, `append_args` or `prepend_args`, the `__signature__` attribute will be added per [PEP 362](https://peps.python.org/pep-0362/).
 
 See also [python documentation on @wraps](https://docs.python.org/3/library/functools.html#functools.wraps)
 
diff --git a/src/makefun/main.py b/src/makefun/main.py
index 14c7b0e..6228ec7 100644
--- a/src/makefun/main.py
+++ b/src/makefun/main.py
@@ -278,6 +278,11 @@ def create_function(func_signature,             # type: Union[str, Signature]
     else:
         raise TypeError("Invalid type for `func_signature`: %s" % type(func_signature))
 
+    if isinstance(attrs.get('__signature__'), str):
+        # __signature__ must be a Signature object, so if it is a string,
+        # we need to evaluate it.
+        attrs['__signature__'] = get_signature_from_string(attrs['__signature__'], evaldict)[1]
+
     # extract all information needed from the `Signature`
     params_to_kw_assignment_mode = get_signature_params(func_signature)
     params_names = list(params_to_kw_assignment_mode.keys())
@@ -819,9 +824,11 @@ def wraps(wrapped_fun,
     `wrapped_fun`, so that the created function seems to be identical (except possiblyfor the signature).
     Note that all options in `with_signature` can still be overrided using parameters of `@wraps`.
 
-    If the signature is *not* modified through `new_sig`, `remove_args`, `append_args` or `prepend_args`, the
-    additional `__wrapped__` attribute  on the created function, to stay consistent with the `functools.wraps`
-    behaviour.
+    The additional `__wrapped__` attribute is set on the created function, to stay consistent
+    with the `functools.wraps` behaviour. If the signature is modified through `new_sig`,
+    `remove_args`, `append_args` or `prepend_args`, the additional
+    `__signature__` attribute will be set so that `inspect.signature` and related functionality
+    works as expected. See PEP 362 for more detail on `__wrapped__` and `__signature__`.
 
     See also [python documentation on @wraps](https://docs.python.org/3/library/functools.html#functools.wraps)
 
@@ -960,15 +967,10 @@ def _get_args_for_wrapping(wrapped, new_sig, remove_args, prepend_args, append_a
 
     # attributes: start from the wrapped dict, add '__wrapped__' if needed, and override with all attrs.
     all_attrs = copy(getattr_partial_aware(wrapped, '__dict__'))
+    # PEP362: always set `__wrapped__`, and if signature was changed, set `__signature__` too
+    all_attrs["__wrapped__"] = wrapped
     if has_new_sig:
-        # change of signature: delete the __wrapped__ attribute if any
-        try:
-            del all_attrs['__wrapped__']
-        except KeyError:
-            pass
-    else:
-        # no change of signature: we can safely set the __wrapped__ attribute
-        all_attrs['__wrapped__'] = wrapped
+        all_attrs["__signature__"] = func_sig
     all_attrs.update(attrs)
 
     return func_name, func_sig, doc, qualname, co_name, module_name, all_attrs


Test patch:
diff --git a/tests/_issue_85_module.py b/tests/_issue_85_module.py
new file mode 100644
index 0000000..7ab3790
--- /dev/null
+++ b/tests/_issue_85_module.py
@@ -0,0 +1,7 @@
+def forwardref_method(foo: "ForwardRef", bar: str) -> "ForwardRef":
+    return ForwardRef(foo.x + bar)
+
+
+class ForwardRef:
+    def __init__(self, x="default"):
+        self.x = x
diff --git a/tests/test_issues.py b/tests/test_issues.py
index 912f9df..87a6c1a 100644
--- a/tests/test_issues.py
+++ b/tests/test_issues.py
@@ -181,11 +181,10 @@ def test_issue_66():
     def second_wrapper(foo, bar):
         return wrapper(foo) + bar
 
+    assert second_wrapper.__wrapped__ is wrapper
+    assert "bar" in signature(second_wrapper).parameters
     assert second_wrapper(1, -1) == 0
 
-    with pytest.raises(AttributeError):
-        second_wrapper.__wrapped__
-
 
 def test_issue_pr_67():
     """Test handcrafted for https://github.com/smarie/python-makefun/pull/67"""
@@ -253,3 +252,25 @@ def test_issue_77_async_generator_partial():
     assert inspect.isasyncgenfunction(f_partial)
 
     assert asyncio.get_event_loop().run_until_complete(asyncio.ensure_future(f_partial().__anext__())) == 1
+
+
+
+@pytest.mark.skipif(sys.version_info < (3, 7, 6), reason="The __wrapped__ behavior in get_type_hints being tested was not added until python 3.7.6.")
+def test_issue_85_wrapped_forwardref_annotation():
+    import typing
+    from . import _issue_85_module
+
+    @wraps(_issue_85_module.forwardref_method, remove_args=["bar"])
+    def wrapper(**kwargs):
+        kwargs["bar"] = "x"  # python 2 syntax to prevent syntax error.
+        return _issue_85_module.forwardref_method(**kwargs)
+
+    # Make sure the wrapper function works as expected
+    assert wrapper(_issue_85_module.ForwardRef()).x == "defaultx"
+
+    # Check that the type hints of the wrapper are ok with the forward reference correctly resolved
+    expected_annotations = {
+        "foo": _issue_85_module.ForwardRef,
+        "return": _issue_85_module.ForwardRef,
+    }
+    assert typing.get_type_hints(wrapper) == expected_annotations