diff --git a/jedi/evaluate/context/decorator.py b/jedi/evaluate/context/decorator.py new file mode 100644 index 00000000..317c5f43 --- /dev/null +++ b/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__() diff --git a/jedi/evaluate/syntax_tree.py b/jedi/evaluate/syntax_tree.py index 99a5786b..083c451e 100644 --- a/jedi/evaluate/syntax_tree.py +++ b/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 diff --git a/test/test_evaluate/test_docstring.py b/test/test_evaluate/test_docstring.py index 676f4e8f..007289f8 100644 --- a/test/test_evaluate/test_docstring.py +++ b/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'