From aa2dc6be09da315ff8ca28553fde0bda01ee0132 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sat, 22 Jun 2019 00:15:20 +0200
Subject: [PATCH] Return annotations for compiled objects now help to infer

However only if it's a type, if it's a string, it doesn't work, yet

Fixes #1347
---
 jedi/evaluate/compiled/access.py  | 11 +++++++++++
 jedi/evaluate/compiled/context.py |  5 +++++
 test/test_api/test_interpreter.py |  5 ++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/jedi/evaluate/compiled/access.py b/jedi/evaluate/compiled/access.py
index 04d00d94..b7aa3874 100644
--- a/jedi/evaluate/compiled/access.py
+++ b/jedi/evaluate/compiled/access.py
@@ -448,6 +448,17 @@ class DirectObjectAccess(object):
             # the signature. In that case we just want a simple escape for now.
             raise ValueError
 
+    def get_return_annotation(self):
+        try:
+            o = self._obj.__annotations__.get('return')
+        except AttributeError:
+            return None
+
+        if o is None:
+            return None
+
+        return self._create_access_path(o)
+
     def negate(self):
         return self._create_access_path(-self._obj)
 
diff --git a/jedi/evaluate/compiled/context.py b/jedi/evaluate/compiled/context.py
index bd9f0d19..aa0090b8 100644
--- a/jedi/evaluate/compiled/context.py
+++ b/jedi/evaluate/compiled/context.py
@@ -46,6 +46,11 @@ class CompiledObject(Context):
         self.access_handle = access_handle
 
     def py__call__(self, arguments):
+        return_annotation = self.access_handle.get_return_annotation()
+        if return_annotation is not None:
+            # TODO the return annotation may also be a string.
+            return create_from_access_path(self.evaluator, return_annotation).execute_annotation()
+
         try:
             self.access_handle.getattr_paths(u'__call__')
         except AttributeError:
diff --git a/test/test_api/test_interpreter.py b/test/test_api/test_interpreter.py
index 1bf73713..6b580f96 100644
--- a/test/test_api/test_interpreter.py
+++ b/test/test_api/test_interpreter.py
@@ -262,7 +262,7 @@ def test_completion_params():
 def test_completion_param_annotations():
     # Need to define this function not directly in Python. Otherwise Jedi is to
     # clever and uses the Python code instead of the signature object.
-    code = 'def foo(a: 1, b: str, c: int = 1.0): pass'
+    code = 'def foo(a: 1, b: str, c: int = 1.0) -> bytes: pass'
     exec_(code, locals())
     script = jedi.Interpreter('foo', [locals()])
     c, = script.completions()
@@ -271,6 +271,9 @@ def test_completion_param_annotations():
     assert [d.name for d in b.infer()] == ['str']
     assert {d.name for d in c.infer()} == {'int', 'float'}
 
+    d, = jedi.Interpreter('foo()', [locals()]).goto_definitions()
+    assert d.name == 'bytes'
+
 
 def test_keyword_argument():
     def f(some_keyword_argument):