From 3a1b2e7104da044433b2822526d923d17e67cca7 Mon Sep 17 00:00:00 2001
From: Claude <github@claude.nl>
Date: Mon, 15 Feb 2016 17:37:03 +0100
Subject: [PATCH] add support for 'for-assignment' hints

---
 jedi/evaluate/finder.py             |  7 ++++++-
 jedi/evaluate/pep0484.py            | 23 ++++++++++++++++-------
 jedi/parser/tree.py                 |  5 ++++-
 test/completion/pep0484_comments.py | 11 +++++++++++
 4 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/jedi/evaluate/finder.py b/jedi/evaluate/finder.py
index 5cf45434..1d98d4dd 100644
--- a/jedi/evaluate/finder.py
+++ b/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)
diff --git a/jedi/evaluate/pep0484.py b/jedi/evaluate/pep0484.py
index 69162178..91f5251e 100644
--- a/jedi/evaluate/pep0484.py
+++ b/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)
-    annotation.parent = stmt.parent
+        node.start_pos)
+    annotation.parent = node.parent
     return _evaluate_for_annotation(evaluator, annotation, index)
diff --git a/jedi/parser/tree.py b/jedi/parser/tree.py
index 0625a98a..cf33faca 100644
--- a/jedi/parser/tree.py
+++ b/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:
diff --git a/test/completion/pep0484_comments.py b/test/completion/pep0484_comments.py
index e49717e6..606bb09a 100644
--- a/test/completion/pep0484_comments.py
+++ b/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