typehints for variables in comments

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()