Support mypy annotations using comment syntax

This allows us to use mypy annotations for completion in Python 2.

Closes #946

jedi/evaluate/pep0484.py

@@ -22,7 +22,7 @@ x support for type hint comments for functions, `# type: (int, str) -> int`.
 import os
 import re
 
-from parso import ParserSyntaxError
+from parso import ParserSyntaxError, parse
 from parso.python import tree
 
 from jedi._compatibility import unicode, force_unicode
@@ -80,9 +80,66 @@ def _fix_forward_reference(context, node):
         return node
 
 
+def _split_mypy_param_declaration(decl_text):
+    """
+    Split decl_text on commas, but group generic expressions
+    together.
+
+    For example, given "foo, Bar[baz, biz]" we return
+    ['foo', 'Bar[baz, biz]'].
+
+    """
+    node = parse(decl_text).children[0]
+
+    if node.type == 'name':
+        return [node.get_code().strip()]
+
+    params = []
+    for child in node.children:
+        if child.type in ['name', 'atom_expr', 'power']:
+            params.append(child.get_code().strip())
+
+    return params
+
+
 @evaluator_method_cache()
 def infer_param(execution_context, param):
+    """
+    Infers the type of a function parameter, using type annotations.
+    """
     annotation = param.annotation
+    if annotation is None:
+        # If no Python 3-style annotation, look for a Python 2-style comment
+        # annotation.
+        # Identify parameters to function in the same sequence as they would
+        # appear in a type comment.
+        all_params = [child for child in param.parent.children
+                      if child.type == 'param']
+
+        node = param.parent.parent
+        comment = parser_utils.get_following_comment_same_line(node)
+        if comment is None:
+            return NO_CONTEXTS
+
+        match = re.match(r"^#\s*type:\s*\(([^#]*)\)\s*->", comment)
+        if not match:
+            return NO_CONTEXTS
+        params_comments = _split_mypy_param_declaration(match.group(1))
+
+        # Find the specific param being investigated
+        index = all_params.index(param)
+        # If the number of parameters doesn't match length of type comment,
+        # ignore first parameter (assume it's self).
+        if len(params_comments) != len(all_params):
+            if index == 0:
+                # Assume it's self, which is already handled
+                return NO_CONTEXTS
+            else:
+                index -= 1
+        param_comment = params_comments[index]
+        # Construct annotation from type comment
+        annotation = tree.String(repr(param_comment), node.start_pos)
+        annotation.parent = node.parent
     module_context = execution_context.get_root_context()
     return _evaluate_for_annotation(module_context, annotation)
 
@@ -102,7 +159,26 @@ def py__annotations__(funcdef):
 
 @evaluator_method_cache()
 def infer_return_types(function_context):
+    """
+    Infers the type of a function's return value,
+    according to type annotations.
+    """
     annotation = py__annotations__(function_context.tree_node).get("return", None)
+    if annotation is None:
+        # If there is no Python 3-type annotation, look for a Python 2-type annotation
+        node = function_context.tree_node
+        comment = parser_utils.get_following_comment_same_line(node)
+        if comment is None:
+            return NO_CONTEXTS
+
+        match = re.match(r"^#\s*type:\s*\([^#]*\)\s*->\s*([^#]*)", comment)
+        if not match:
+            return NO_CONTEXTS
+
+        annotation = tree.String(
+            repr(str(match.group(1).strip())),
+            node.start_pos)
+        annotation.parent = node.parent
     module_context = function_context.get_root_context()
     return _evaluate_for_annotation(module_context, annotation)
 

jedi/parser_utils.py

@@ -201,6 +201,9 @@ def get_following_comment_same_line(node):
             whitespace = node.children[5].get_first_leaf().prefix
         elif node.type == 'with_stmt':
             whitespace = node.children[3].get_first_leaf().prefix
+        elif node.type == 'funcdef':
+            # actually on the next line
+            whitespace = node.children[4].get_first_leaf().get_next_leaf().prefix
         else:
             whitespace = node.get_last_leaf().get_next_leaf().prefix
     except AttributeError:

test/completion/pep0484_comments.py

@@ -47,7 +47,7 @@ b
 class Employee:
     pass
 
-from typing import List
+from typing import List, Tuple
 x = []   # type: List[Employee]
 #? Employee()
 x[1]
@@ -103,3 +103,67 @@ aaa = some_extremely_long_function_name_that_doesnt_leave_room_for_hints() \
     # type: float # We should be able to put hints on the next line with a \
 #? float()
 aaa
+
+# Test instance methods
+class Dog:
+    def __init__(self, age, friends, name):
+        # type: (int, List[Tuple[str, Dog]], str) -> None
+        #? int()
+        self.age = age
+        self.friends = friends
+
+        #? Dog()
+        friends[0][1]
+
+        #? str()
+        self.name = name
+
+    def friend_for_name(self, name):
+        # type: (str) -> Dog
+        for (friend_name, friend) in self.friends:
+            if friend_name == name:
+                return friend
+        raise ValueError()
+
+    def bark(self):
+        pass
+
+buddy = Dog(UNKNOWN_NAME1, UNKNOWN_NAME2, UNKNOWN_NAME3)
+friend = buddy.friend_for_name('buster')
+# type of friend is determined by function return type
+#! 9 ['def bark']
+friend.bark()
+
+friend = buddy.friends[0][1]
+# type of friend is determined by function parameter type
+#! 9 ['def bark']
+friend.bark()
+
+# type is determined by function parameter type following nested generics
+#? str()
+friend.name
+
+# Mypy comment describing function return type.
+def annot():
+    # type: () -> str
+    pass
+
+#? str()
+annot()
+
+# Mypy variable type annotation.
+x = UNKNOWN_NAME2  # type: str
+
+#? str()
+x
+
+class Cat(object):
+    def __init__(self, age, friends, name):
+        # type: (int, List[Dog], str) -> None
+        self.age = age
+        self.friends = friends
+        self.name = name
+
+cat = Cat(UNKNOWN_NAME4, UNKNOWN_NAME5, UNKNOWN_NAME6)
+#? str()
+cat.name