diff --git a/jedi/evaluate/pep0484.py b/jedi/evaluate/pep0484.py index a9dc748c..b9008a94 100644 --- a/jedi/evaluate/pep0484.py +++ b/jedi/evaluate/pep0484.py @@ -21,12 +21,22 @@ from itertools import chain from jedi.evaluate.cache import memoize_default -@memoize_default(None, evaluator_is_first_arg=True) -def follow_param(evaluator, param): - annotation = param.annotation() - if annotation: +def _evaluate_for_annotation(evaluator, annotation): + if annotation is not None: definitions = evaluator.eval_element(annotation) return list(chain.from_iterable( evaluator.execute(d) for d in definitions)) else: return [] + + +@memoize_default(None, evaluator_is_first_arg=True) +def follow_param(evaluator, param): + annotation = param.annotation() + return _evaluate_for_annotation(evaluator, annotation) + + +@memoize_default(None, evaluator_is_first_arg=True) +def find_return_types(evaluator, func): + annotation = func.py__annotations__().get("return", None) + return _evaluate_for_annotation(evaluator, annotation) diff --git a/jedi/evaluate/representation.py b/jedi/evaluate/representation.py index 173cf9be..9d0026c0 100644 --- a/jedi/evaluate/representation.py +++ b/jedi/evaluate/representation.py @@ -49,6 +49,7 @@ from jedi.evaluate import compiled from jedi.evaluate import recursion from jedi.evaluate import iterable from jedi.evaluate import docstrings +from jedi.evaluate import pep0484 from jedi.evaluate import helpers from jedi.evaluate import param from jedi.evaluate import flow_analysis @@ -585,14 +586,13 @@ class Function(use_metaclass(CachedMetaClass, Wrapper)): parser_func = self.base return_annotation = parser_func.annotation() if return_annotation: - dct = {'return': self._evaluator.eval_element(return_annotation)} + dct = {'return': return_annotation} else: dct = {} for function_param in parser_func.params: param_annotation = function_param.annotation() - if param_annotation: - dct[function_param.name.value] = \ - self._evaluator.eval_element(param_annotation) + if param_annotation is not None: + dct[function_param.name.value] = param_annotation return dct def py__class__(self): @@ -654,9 +654,7 @@ class FunctionExecution(Executed): else: returns = self.returns types = set(docstrings.find_return_types(self._evaluator, func)) - annotations = func.py__annotations__().get("return", []) - types |= set(chain.from_iterable( - self._evaluator.execute(d) for d in annotations)) + types |= set(pep0484.find_return_types(self._evaluator, func)) for r in returns: check = flow_analysis.break_check(self._evaluator, self, r)