Fix docstrings for wrapped functions, fixes #906

jedi/evaluate/context/decorator.py

@@ -0,0 +1,15 @@
+'''
+Decorators are not really contexts, however we need some wrappers to improve
+docstrings and other things around decorators.
+'''
+
+from jedi.evaluate.base_context import ContextWrapper
+
+
+class Decoratee(ContextWrapper):
+    def __init__(self, wrapped_context, original_context):
+        self._wrapped_context = wrapped_context
+        self._original_context = original_context
+
+    def py__doc__(self):
+        return self._original_context.py__doc__()

jedi/evaluate/syntax_tree.py

@@ -26,6 +26,7 @@ from jedi.evaluate.compiled.access import COMPARISON_OPERATORS
 from jedi.evaluate.cache import evaluator_method_cache
 from jedi.evaluate.gradual.stub_context import VersionInfo
 from jedi.evaluate.gradual import annotation
+from jedi.evaluate.context.decorator import Decoratee
 
 
 def _limit_context_infers(func):
@@ -666,6 +667,8 @@ def _apply_decorators(context, node):
                 return initial
 
         debug.dbg('decorator end %s', values, color="MAGENTA")
+    if values != initial:
+        return ContextSet([Decoratee(c, decoratee_context) for c in values])
     return values
 
 

test/test_evaluate/test_docstring.py

@@ -385,3 +385,26 @@ def test_numpy_comp_returns():
     )
     names = [c.name for c in jedi.Script(s).completions()]
     assert 'diagonal' in names
+
+
+def test_decorator(Script):
+    code = dedent('''
+        def decorator(name=None):
+            def _decorate(func):
+                @wraps(func)
+                def wrapper(*args, **kwargs):
+                    """wrapper docstring"""
+                    return func(*args, **kwargs)
+                return wrapper
+            return _decorate
+
+
+        @decorator('testing')
+        def check_user(f):
+            """Nice docstring"""
+            pass
+
+        check_user''')
+
+    d, = Script(code).goto_definitions()
+    assert d.docstring(raw=True) == 'Nice docstring'