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)