Merge pull request #664 from reinhrst/typehints-in-comments

simple typehints in comments

jedi/evaluate/finder.py

@@ -304,6 +304,14 @@ class NameFinder(object):
 @memoize_default(set(), evaluator_is_first_arg=True)
 def _name_to_types(evaluator, name, scope):
     typ = name.get_definition()
+    if typ.isinstance(tree.ForStmt):
+        types = pep0484.find_type_from_comment_hint_for(evaluator, typ, name)
+        if types:
+            return types
+    if typ.isinstance(tree.WithStmt):
+        types = pep0484.find_type_from_comment_hint_with(evaluator, typ, name)
+        if types:
+            return types
     if typ.isinstance(tree.ForStmt, tree.CompFor):
         container_types = evaluator.eval_element(typ.children[3])
         for_types = iterable.py__iter__types(evaluator, container_types, typ.children[3])
@@ -355,6 +363,10 @@ def _remove_statements(evaluator, stmt, name):
         check_instance = stmt.instance
         stmt = stmt.var
 
+    pep0484types = \
+        pep0484.find_type_from_comment_hint_assign(evaluator, stmt, name)
+    if pep0484types:
+        return pep0484types
     types |= evaluator.eval_statement(stmt, seek_name=name)
 
     if check_instance is not None:

jedi/evaluate/pep0484.py

@@ -9,13 +9,14 @@ v Function parameter annotations with builtin/custom type classes
 v Function returntype annotations with builtin/custom type classes
 v Function parameter annotations with strings (forward reference)
 v Function return type annotations with strings (forward reference)
-x Local variable type hints
+v Local variable type hints
 v Assigned types: `Url = str\ndef get(url:Url) -> str:`
-x Type hints in `with` statements
+v Type hints in `with` statements
 x Stub files support
 x support `@no_type_check` and `@no_type_check_decorator`
-x support for type hint comments `# type: (int, str) -> int`. See comment from
+x support for typing.cast() operator
-    Guido https://github.com/davidhalter/jedi/issues/662
+x support for type hint comments for functions, `# type: (int, str) -> int`.
+    See comment from Guido https://github.com/davidhalter/jedi/issues/662
 """
 
 import itertools
@@ -28,12 +29,23 @@ from jedi.common import unite
 from jedi.evaluate import compiled
 from jedi import debug
 from jedi import _compatibility
+import re
 
 
-def _evaluate_for_annotation(evaluator, annotation):
+def _evaluate_for_annotation(evaluator, annotation, index=None):
+    """
+    Evaluates a string-node, looking for an annotation
+    If index is not None, the annotation is expected to be a tuple
+    and we're interested in that index
+    """
     if annotation is not None:
         definitions = evaluator.eval_element(
             _fix_forward_reference(evaluator, annotation))
+        if index is not None:
+            definitions = list(itertools.chain.from_iterable(
+                definition.py__getitem__(index) for definition in definitions
+                if definition.type == 'tuple' and
+                len(list(definition.py__iter__())) >= index))
         return list(itertools.chain.from_iterable(
             evaluator.execute(d) for d in definitions))
     else:
@@ -136,3 +148,48 @@ def get_types_for_typing_module(evaluator, typ, node):
 
     result = evaluator.execute_evaluated(factory, compiled_classname, args)
     return result
+
+
+def find_type_from_comment_hint_for(evaluator, node, name):
+    return \
+        _find_type_from_comment_hint(evaluator, node, node.children[1], name)
+
+
+def find_type_from_comment_hint_with(evaluator, node, name):
+    assert len(node.children[1].children) == 3, \
+        "Can only be here when children[1] is 'foo() as f'"
+    return _find_type_from_comment_hint(
+        evaluator, node, node.children[1].children[2], name)
+
+
+def find_type_from_comment_hint_assign(evaluator, node, name):
+    return \
+        _find_type_from_comment_hint(evaluator, node, node.children[0], name)
+
+
+def _find_type_from_comment_hint(evaluator, node, varlist, name):
+    index = None
+    if varlist.type in ("testlist_star_expr", "exprlist"):
+        # something like "a, b = 1, 2"
+        index = 0
+        for child in varlist.children:
+            if child == name:
+                break
+            if child.type == "operator":
+                continue
+            index += 1
+        else:
+            return []
+
+    comment = node.get_following_comment_same_line()
+    if comment is None:
+        return []
+    match = re.match(r"^#\s*type:\s*([^#]*)", comment)
+    if not match:
+        return []
+    annotation = tree.String(
+        tree.zero_position_modifier,
+        repr(str(match.group(1).strip())),
+        node.start_pos)
+    annotation.parent = node.parent
+    return _evaluate_for_annotation(evaluator, annotation, index)

jedi/parser/tree.py

@@ -239,13 +239,35 @@ class Leaf(Base):
             else:
                 node = c[i - 1]
                 break
-
         while True:
             try:
                 node = node.children[-1]
             except AttributeError:  # A Leaf doesn't have children.
                 return node
 
+    def get_next(self):
+        """
+        Returns the next leaf in the parser tree.
+        """
+        node = self
+        while True:
+            c = node.parent.children
+            i = c.index(node)
+            try:
+                node = c[i + 1]
+            except IndexError:
+                node = node.parent
+                if node.parent is None:
+                    raise IndexError('Cannot access the next element of the last one.')
+            else:
+                break
+        while True:
+            try:
+                node = node.children[0]
+            except AttributeError:  # A Leaf doesn't have children.
+                return node
+
+
     def get_code(self, normalized=False):
         if normalized:
             return self.value
@@ -264,6 +286,7 @@ class Leaf(Base):
                 except IndexError:
                     return None
 
+
     def prev_sibling(self):
         """
         The node/leaf immediately preceding the invocant in their parent's
@@ -277,6 +300,7 @@ class Leaf(Base):
                     return None
                 return self.parent.children[i - 1]
 
+
     def nodes_to_execute(self, last_added=False):
         return []
 
@@ -488,6 +512,39 @@ class BaseNode(Base):
         except AttributeError:
             return self.children[0]
 
+    def last_leaf(self):
+        try:
+            return self.children[-1].last_leaf()
+        except AttributeError:
+            return self.children[-1]
+
+    def get_following_comment_same_line(self):
+        """
+        returns (as string) any comment that appears on the same line,
+        after the node, including the #
+        """
+        try:
+            if self.isinstance(ForStmt):
+                whitespace = self.children[5].first_leaf().prefix
+            elif self.isinstance(WithStmt):
+                whitespace = self.children[3].first_leaf().prefix
+            else:
+                whitespace = self.last_leaf().get_next().prefix
+        except AttributeError:
+            return None
+        except ValueError:
+            # in some particular cases, the tree doesn't seem to be linked
+            # correctly
+            return None
+        if "#" not in whitespace:
+            return None
+        comment = whitespace[whitespace.index("#"):]
+        if "\r" in comment:
+            comment = comment[:comment.index("\r")]
+        if "\n" in comment:
+            comment = comment[:comment.index("\n")]
+        return comment
+
     @utf8_repr
     def __repr__(self):
         code = self.get_code().replace('\n', ' ').strip()

test/completion/pep0484_comments.py

@@ -0,0 +1,105 @@
+a = 3  # type: str
+#? str()
+a
+
+b = 3  # type: str but I write more
+#? int()
+b
+
+c = 3  # type: str # I comment more
+#? str()
+c
+
+d = "It should not read comments from the next line"
+# type: int
+#? str()
+d
+
+# type: int
+e = "It should not read comments from the previous line"
+#? str()
+e
+
+class BB: pass
+
+def test(a, b):
+    a = a  # type: BB
+    c = a  # type: str
+    d = a
+    # type: str
+    e = a                 # type: str           # Should ignore long whitespace
+
+    #? BB()
+    a
+    #? str()
+    c
+    #? BB()
+    d
+    #? str()
+    e
+
+a,b = 1, 2 # type: str, float
+#? str()
+a
+#? float()
+b
+
+class Employee:
+    pass
+
+from typing import List
+x = []   # type: List[Employee]
+#? Employee()
+x[1]
+x, y, z = [], [], []  # type: List[int], List[int], List[str]
+#? int()
+y[2]
+x, y, z = [], [], []  # type: (List[float], List[float], List[BB])
+for zi in z:
+    #? BB()
+    zi
+
+x = [
+   1,
+   2,
+]  # type: List[str]
+
+#? str()
+x[1]
+
+
+for bar in foo():  # type: str
+    #? str()
+    bar
+
+for bar, baz in foo():  # type: int, float
+    #? int()
+    bar
+    #? float()
+    baz
+
+for bar, baz in foo():
+    # type: str, str
+    """ type hinting on next line should not work """
+    #?
+    bar
+    #?
+    baz
+
+with foo():  # type: int
+    ...
+
+with foo() as f:  # type: str
+    #? str()
+    f
+
+with foo() as f:
+    # type: str
+    """ type hinting on next line should not work """
+    #?
+    f
+
+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