typehints for variables in comments

jedi/evaluate/pep0484.py

@@ -140,27 +140,15 @@ def get_types_for_typing_module(evaluator, typ, node):
 
 
 def find_type_from_comment_hint(evaluator, stmt):
-    try:
-        stmtpos = stmt.parent.children.index(stmt)
-    except ValueError:
-        return []
-    try:
-        next_sibling = stmt.parent.children[stmtpos + 1]
-    except IndexError:
-        return []
-    if not isinstance(next_sibling, tree.Whitespace):
-        return []
-    comment = next_sibling.get_pre_comment()
+    comment = stmt.get_following_comment_same_line()
     if comment is None:
         return []
-    match = re.match(r"\s*type:\s*([^#]*)", comment)
+    match = re.match(r"^#\s*type:\s*([^#]*)", comment)
     if not match:
         return []
-    start_pos = (next_sibling.start_pos[0],
-                 next_sibling.start_pos[1] - len(comment))
     annotation = tree.String(
         tree.zero_position_modifier,
         repr(str(match.group(1).strip())),
-        start_pos)
+        stmt.start_pos)
     annotation.parent = stmt.parent
     return _evaluate_for_annotation(evaluator, annotation)

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,18 +300,10 @@ class Leaf(Base):
                     return None
                 return self.parent.children[i - 1]
 
+
     def nodes_to_execute(self, last_added=False):
         return []
 
-    def get_pre_comment(self):
-        """
-        returns comment before this leaf, excluding #, or None if no comment
-        """
-        match = re.match(r"\s*#(.*)$", self.prefix)
-        if match:
-            return match.group(1)
-        return None
-
     @utf8_repr
     def __repr__(self):
         return "<%s: %s>" % (type(self).__name__, self.value)
@@ -497,6 +512,30 @@ class BaseNode(Base):
         except AttributeError:
             return self.children[0]
 
+    def last_leaf(self):
+        try:
+            return self.children[-1].first_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:
+            whitespace = self.last_leaf().get_next().prefix
+        except AttributeError:
+            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.py

@@ -157,27 +157,3 @@ Y = int
 def just_because_we_can(x: "flo" + "at"):
     #? float()
     x
-
-# python >= 2.6
-
-x = 3  # type: str
-#? str()
-x
-
-y = 3  # type: str but I write more
-#? int()
-y
-
-z = 3  # type: str # I comment more
-#? str()
-z
-
-class BB: pass
-
-def test(a, b):
-    a = a  # type: BB
-    c = a  # type: str
-    #? BB()
-    a
-    #? str()
-    c

test/completion/pep0484_comments.py

@@ -0,0 +1,39 @@
+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