pep0484 return type support

jedi/evaluate/pep0484.py

@@ -6,7 +6,7 @@ as annotations in future python versions.
 
 The (initial / probably incomplete) implementation todo list for pep-0484:
 v Function parameter annotations with builtin/custom type classes
-x Function returntype annotations with builtin/custom type classes
+v Function returntype annotations with builtin/custom type classes
 x Function parameter annotations with strings (forward reference)
 x Function return type annotations with strings (forward reference)
 x Local variable type hints

jedi/evaluate/representation.py

@@ -580,6 +580,21 @@ class Function(use_metaclass(CachedMetaClass, Wrapper)):
         else:
             return FunctionExecution(self._evaluator, self, params).get_return_types()
 
+    @memoize_default()
+    def py__annotations__(self):
+        parser_func = self.base
+        return_annotation = parser_func.annotation()
+        if return_annotation:
+            dct = {'return': self._evaluator.eval_element(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)
+        return dct
+
     def py__class__(self):
         return compiled.get_special_object(self._evaluator, 'FUNCTION_CLASS')
 
@@ -639,6 +654,9 @@ 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))
 
         for r in returns:
             check = flow_analysis.break_check(self._evaluator, self, r)

jedi/parser/tree.py

@@ -866,7 +866,9 @@ class Function(ClassOrFunc):
 
     def annotation(self):
         try:
-            return self.children[6]  # 6th element: def foo(...) -> bar
+            if self.children[3] == "->":
+                return self.children[4]
+            assert self.children[3] == ":"
         except IndexError:
             return None
 

test/completion/pep0484.py

@@ -2,17 +2,69 @@
 
 # python >= 3.2
 
-# -----------------
+
-# simple classes
+class A():
-# -----------------
+    pass
 
 
-def typehints(a, b: str, c: int, d: int=4):
+def function_parameters(a: A, b, c: str, d: int=4):
-    #?
+    #? A()
     a
-    #? str()
+    #?
     b
-    #? int()
+    #? str()
     c
     #? int()
     d
+
+
+def return_unspecified():
+    pass
+
+#?
+return_unspecified()
+
+
+def return_none() -> None:
+    """
+    Return type None means the same as no return type as far as jedi
+    is concerned
+    """
+    pass
+
+#?
+return_none()
+
+
+def return_str() -> str:
+    pass
+
+#? str()
+return_str()
+
+
+def return_custom_class() -> A:
+    pass
+
+#? A()
+return_custom_class()
+
+
+def return_annotation_and_docstring() -> str:
+    """
+    :rtype: int
+    """
+    pass
+
+#? str() int()
+return_annotation_and_docstring()
+
+
+def return_annotation_and_docstring_different() -> str:
+    """
+    :rtype: str
+    """
+    pass
+
+#? str()
+return_annotation_and_docstring_different()