add support for 'for-assignment' hints

jedi/evaluate/finder.py

@@ -304,6 +304,10 @@ 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.ForStmt, tree.CompFor):
         container_types = evaluator.eval_element(typ.children[3])
         for_types = iterable.py__iter__types(evaluator, container_types, typ.children[3])
@@ -355,7 +359,8 @@ def _remove_statements(evaluator, stmt, name):
         check_instance = stmt.instance
         stmt = stmt.var
 
-    pep0484types = pep0484.find_type_from_comment_hint(evaluator, stmt, name)
+    pep0484types = \
+        pep0484.find_type_from_comment_hint_assign(evaluator, stmt, name)
     if pep0484types:
         return pep0484types
     types |= evaluator.eval_statement(stmt, seek_name=name)

jedi/evaluate/pep0484.py

@@ -149,13 +149,22 @@ def get_types_for_typing_module(evaluator, typ, node):
     return result
 
 
-def find_type_from_comment_hint(evaluator, stmt, name):
+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_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 stmt.children[0].type == "testlist_star_expr":
+    if varlist.type in ("testlist_star_expr", "exprlist"):
         # something like "a, b = 1, 2"
-        leftside = stmt.children[0]
         index = 0
-        for child in leftside.children:
+        for child in varlist.children:
             if child == name:
                 break
             if child.type == "operator":
@@ -164,7 +173,7 @@ def find_type_from_comment_hint(evaluator, stmt, name):
         else:
             return []
 
-    comment = stmt.get_following_comment_same_line()
+    comment = node.get_following_comment_same_line()
     if comment is None:
         return []
     match = re.match(r"^#\s*type:\s*([^#]*)", comment)
@@ -173,6 +182,6 @@ def find_type_from_comment_hint(evaluator, stmt, name):
     annotation = tree.String(
         tree.zero_position_modifier,
         repr(str(match.group(1).strip())),
-        stmt.start_pos)
+        node.start_pos)
-    annotation.parent = stmt.parent
+    annotation.parent = node.parent
     return _evaluate_for_annotation(evaluator, annotation, index)

jedi/parser/tree.py

@@ -524,7 +524,10 @@ class BaseNode(Base):
         after the node, including the #
         """
         try:
-            whitespace = self.last_leaf().get_next().prefix
+            if self.isinstance(ForStmt):
+                whitespace = self.children[5].first_leaf().prefix
+            else:
+                whitespace = self.last_leaf().get_next().prefix
         except AttributeError:
             return None
         except ValueError:

test/completion/pep0484_comments.py

@@ -66,3 +66,14 @@ x = [
 
 #? str()
 x[1]
+
+
+for bar in foo():  # type: str
+    #? str()
+    bar
+
+for bar, baz in foo():  # type: int, float
+    #? int()
+    bar
+    #? float()
+    baz