Support 'with' statements with names (close #123)

Author: JukkaL

mypy/checker.py

@@ -769,6 +769,8 @@ def is_definition(self, s):
             ctx = self.accept(expr)
             enter = echk.analyse_external_member_access('__enter__', ctx, expr)
             obj = echk.check_call(enter, [], [], expr)[0]
+            if name:
+                self.check_assignments([name], self.temp_node(obj, expr))
             exit = echk.analyse_external_member_access('__exit__', ctx, expr)
             arg = self.temp_node(Any(), expr)
             echk.check_call(exit, [arg] * 3, [nodes.ARG_POS] * 3, expr)

mypy/nodes.py

@@ -581,10 +581,10 @@ class TryStmt(Node):
 
 class WithStmt(Node):
     Node[] expr
-    Var[] name
+    NameExpr[] name
     Block body
 
-    void __init__(self, Node[] expr, Var[] name, Block body):
+    void __init__(self, Node[] expr, NameExpr[] name, Block body):
         self.expr = expr
         self.name = name
         self.body = body

mypy/parse.py

@@ -974,17 +974,17 @@ def parse_type_def(self, is_interface):
         Token[] as_toks = []
         Token[] commas = []
         Node[] expr = []
-        Var[] name = []
+        NameExpr[] name = []
         while True:
             e = self.parse_expression(precedence[','])
-            Var v = None
             if self.current_str() == 'as':
                 as_toks.append(self.expect('as'))
-                v = self.parse_var()
+                n = self.parse_name_expr()
             else:
                 as_toks.append(none)
+                n = None
             expr.append(e)
-            name.append(v)
+            name.append(n)
             if self.current_str() != ',':
                 break
             commas.append(self.expect(','))

mypy/semanal.py

@@ -106,6 +106,8 @@ class SemanticAnalyzer(NodeVisitor):
                 self.anal_var_def((VarDef)d)
             elif isinstance(d, ForStmt):
                 self.anal_for_stmt((ForStmt)d)
+            elif isinstance(d, WithStmt):
+                self.anal_with_stmt((WithStmt)d)
             elif isinstance(d, Decorator):
                 self.anal_decorator((Decorator)d)
         # Add implicit definition of 'None' to builtins, as we cannot define a
@@ -147,6 +149,11 @@ class SemanticAnalyzer(NodeVisitor):
         for n in s.index:
             self.analyse_lvalue(n, False, True)
 
+    void anal_with_stmt(self, WithStmt s):
+        for n in s.name:
+            if n:
+                self.analyse_lvalue(n, False, True)
+
     void anal_decorator(self, Decorator d):
         d.var._fullname = self.qualified_name(d.var.name())
         self.add_symbol(d.var.name(), SymbolTableNode(GDEF, d.var), d)
@@ -559,7 +566,7 @@ class SemanticAnalyzer(NodeVisitor):
             e.accept(self)
         for n in s.name:
             if n:
-                self.add_var(n, s)
+                self.analyse_lvalue(n)
         self.visit_block(s.body)
 
     void visit_del_stmt(self, DelStmt s):

test/data/check-statements.test

@@ -491,3 +491,35 @@ class A:
     void __exit__(self, x, y, z): pass
 with A(): # E: "A" has no member "__enter__"
     pass
+
+[case testWithStmtAndMultipleExprs]
+class A:
+    void __enter__(self): pass
+    void __exit__(self, x, y, z): pass
+class B:
+    void __enter__(self): pass
+with A(), B(): # E: "B" has no member "__exit__"
+    pass
+with B(), A(): # E: "B" has no member "__exit__"
+    pass
+
+[case testWithStmtAndResult]
+class B: pass
+class A:
+    B __enter__(self): pass
+    def __exit__(self, x, y, z): pass
+with A() as b:
+    b = B()
+    b = A() # E: Incompatible types in assignment
+
+[case testWithStmtAndMultipleResults]
+class B: pass
+class C: pass
+class A<t>:
+    t __enter__(self): pass
+    def __exit__(self, x, y, z): pass
+with A<B>() as b, A<C>() as c:
+    b = B()
+    c = C()
+    b = c # E: Incompatible types in assignment
+    c = b # E: Incompatible types in assignment

test/data/semanal-statements.test

@@ -575,12 +575,12 @@ MypyFile:1(
     Expr(
       NameExpr(object [builtins.object]))
     Name(
-      Var(x))
+      NameExpr(x* [__main__.x]))
     Block:1(
       ExpressionStmt:2(
         NameExpr(x [__main__.x])))))
 
-[case testInFunction]
+[case testWithInFunction]
 def f():
   with f as x:
     x
@@ -593,7 +593,7 @@ MypyFile:1(
         Expr(
           NameExpr(f [__main__.f]))
         Name(
-          Var(x))
+          NameExpr(x* [l]))
         Block:2(
           ExpressionStmt:3(
             NameExpr(x [l])))))))
@@ -616,11 +616,11 @@ MypyFile:1(
     Expr(
       NameExpr(object [builtins.object]))
     Name(
-      Var(a))
+      NameExpr(a* [__main__.a]))
     Expr(
       NameExpr(object [builtins.object]))
     Name(
-      Var(b))
+      NameExpr(b* [__main__.b]))
     Block:3(
       PassStmt:4())))