add support for 'with-assignment' hints

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

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)

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:

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