From 641fb8077398100179d8dc8a2711f50eac9566fa Mon Sep 17 00:00:00 2001
From: Claude <github@claude.nl>
Date: Mon, 15 Feb 2016 17:52:21 +0100
Subject: [PATCH] add support for 'with-assignment' hints

---
 jedi/evaluate/finder.py             | 4 ++++
 jedi/evaluate/pep0484.py            | 7 +++++++
 jedi/parser/tree.py                 | 2 ++
 test/completion/pep0484_comments.py | 7 +++++++
 4 files changed, 20 insertions(+)

diff --git a/jedi/evaluate/finder.py b/jedi/evaluate/finder.py
index 1d98d4dd..4114518a 100644
--- a/jedi/evaluate/finder.py
+++ b/jedi/evaluate/finder.py
@@ -308,6 +308,10 @@ def _name_to_types(evaluator, name, scope):
         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])
diff --git a/jedi/evaluate/pep0484.py b/jedi/evaluate/pep0484.py
index 91f5251e..aeb57735 100644
--- a/jedi/evaluate/pep0484.py
+++ b/jedi/evaluate/pep0484.py
@@ -154,6 +154,13 @@ def find_type_from_comment_hint_for(evaluator, node, name):
         _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)
diff --git a/jedi/parser/tree.py b/jedi/parser/tree.py
index cf33faca..133b75ae 100644
--- a/jedi/parser/tree.py
+++ b/jedi/parser/tree.py
@@ -526,6 +526,8 @@ class BaseNode(Base):
         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:
diff --git a/test/completion/pep0484_comments.py b/test/completion/pep0484_comments.py
index 606bb09a..8e907065 100644
--- a/test/completion/pep0484_comments.py
+++ b/test/completion/pep0484_comments.py
@@ -77,3 +77,10 @@ for bar, baz in foo():  # type: int, float
     bar
     #? float()
     baz
+
+with foo():  # type: int
+    ...
+
+with foo() as f:  # type: str
+    #? str()
+    f