Merge pull request RustPython#139 from RustPython/with

With statement implementation

Author: windelbouwman

parser/src/ast.rs

@@ -74,7 +74,7 @@ pub enum Statement {
         orelse: Option<Vec<LocatedStatement>>,
     },
     With {
-        items: Expression,
+        items: Vec<WithItem>,
         body: Vec<LocatedStatement>,
     },
     For {
@@ -106,6 +106,12 @@ pub enum Statement {
     },
 }
 
+#[derive(Debug, PartialEq)]
+pub struct WithItem {
+    pub context_expr: Expression,
+    pub optional_vars: Option<Expression>,
+}
+
 #[derive(Debug, PartialEq, Clone)]
 pub enum Expression {
     BoolOp {

parser/src/python.lalrpop

@@ -299,14 +299,28 @@ ExceptClause: ast::ExceptHandler = {
 };
 
 WithStatement: ast::LocatedStatement = {
-  <loc:@L> "with" <t:Test> "as" <_e:Expression> ":" <s:Suite> => {
+  <loc:@L> "with" <i1:WithItem> <i2:("," WithItem)*> ":" <s:Suite> => {
+    let mut items = vec![i1];
+    for item in i2 {
+        items.push(item.1);
+    }
     ast::LocatedStatement {
       location: loc,
-      node: ast::Statement::With { items: t, body: s },
+      node: ast::Statement::With { items: items, body: s },
     }
   },
 };
 
+WithItem: ast::WithItem = {
+    <t:Test> <n:("as" Expression)?> => {
+        let optional_vars = match n {
+            Some(val) => Some(val.1),
+            None => None,
+        };
+        ast::WithItem { context_expr: t, optional_vars }
+    },
+};
+
 FuncDef: ast::LocatedStatement = {
   <loc:@L> "def" <i:Identifier> <a:Parameters> ":" <s:Suite> => {
     ast::LocatedStatement {

tests/snippets/with.py

@@ -0,0 +1,42 @@
+
+
+class ContextManager:
+    def __enter__(self):
+        print('Entrada')
+        ls.append(1)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        ls.append(2)
+        print('Wiedersehen')
+
+    def __str__(self):
+        ls.append(3)
+        return "c'est moi!"
+
+ls = []
+with ContextManager() as c:
+    print(c)
+assert ls == [1, 3, 2]
+
+class ContextManager2:
+    def __enter__(self):
+        print('Ni hau')
+        ls.append(4)
+        return ls
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        ls.append(5)
+        print('Ajuus')
+
+ls = []
+with ContextManager2() as c:
+    print(c)
+    assert c == [4]
+assert ls == [4, 5]
+
+ls = []
+with ContextManager() as c1, ContextManager2() as c2:
+    print(c1)
+    assert c2 == [1, 4, 3]
+assert ls == [1, 4, 3, 5, 2]

vm/src/bytecode.rs

@@ -118,6 +118,12 @@ pub enum Instruction {
     SetupExcept {
         handler: Label,
     },
+    SetupWith {
+        end: Label,
+    },
+    CleanupWith {
+        end: Label,
+    },
     PopBlock,
     Raise {
         argc: usize,

vm/src/compile.rs

@@ -199,8 +199,26 @@ impl Compiler {
                 });
                 self.set_label(end_label);
             }
-            ast::Statement::With { items: _, body: _ } => {
-                // TODO
+            ast::Statement::With { items, body } => {
+                let end_label = self.new_label();
+                for item in items {
+                    self.compile_expression(&item.context_expr);
+                    self.emit(Instruction::SetupWith { end: end_label });
+                    match &item.optional_vars {
+                        Some(var) => {
+                            self.compile_store(var);
+                        }
+                        None => {
+                            self.emit(Instruction::Pop);
+                        }
+                    }
+                }
+
+                self.compile_statements(body);
+                for _ in 0..items.len() {
+                    self.emit(Instruction::CleanupWith { end: end_label });
+                }
+                self.set_label(end_label);
             }
             ast::Statement::For {
                 target,

vm/src/frame.rs

@@ -16,6 +16,10 @@ pub enum Block {
     #[allow(dead_code)]
     // TODO: Implement try/except blocks
     TryExcept { handler: bytecode::Label },
+    With {
+        end: bytecode::Label,
+        context_manager: PyObjectRef,
+    },
 }
 
 pub struct Frame {

vm/src/vm.rs

@@ -175,12 +175,24 @@ impl VirtualMachine {
         self.current_frame().last_block()
     }
 
+    fn with_exit(&mut self) {
+        // Assume top of stack is __exit__ method:
+        // TODO: do we want to put the exit call on the stack?
+        let exit_method = self.pop_value();
+        let args = PyFuncArgs::default();
+        // TODO: what happens when we got an error during handling exception?
+        self.invoke(exit_method, args).unwrap();
+    }
+
     fn unwind_loop(&mut self) -> Block {
         loop {
             let block = self.pop_block();
             match block {
                 Some(Block::Loop { start: _, end: __ }) => break block.unwrap(),
                 Some(Block::TryExcept { .. }) => {}
+                Some(Block::With { .. }) => {
+                    self.with_exit();
+                }
                 None => panic!("No block to break / continue"),
             }
         }
@@ -196,7 +208,10 @@ impl VirtualMachine {
                     self.jump(handler);
                     return None;
                 }
-                Some(_) => {}
+                Some(Block::With { .. }) => {
+                    self.with_exit();
+                }
+                Some(Block::Loop { .. }) => {}
                 None => break,
             }
         }
@@ -823,6 +838,45 @@ impl VirtualMachine {
                 self.push_block(Block::TryExcept { handler: *handler });
                 None
             }
+            bytecode::Instruction::SetupWith { end } => {
+                let context_manager = self.pop_value();
+                // Call enter:
+                match self.call_method(context_manager.clone(), "__enter__", vec![]) {
+                    Ok(obj) => {
+                        self.push_block(Block::With {
+                            end: *end,
+                            context_manager: context_manager.clone(),
+                        });
+                        self.push_value(obj);
+                        None
+                    }
+                    Err(err) => Some(Err(err)),
+                }
+            }
+            bytecode::Instruction::CleanupWith { end: end1 } => {
+                let block = self.pop_block().unwrap();
+                if let Block::With {
+                    end: end2,
+                    context_manager,
+                } = &block
+                {
+                    assert!(end1 == end2);
+
+                    // call exit now:
+                    // TODO: improve exception handling in context manager.
+                    let exc_type = self.ctx.none();
+                    let exc_val = self.ctx.none();
+                    let exc_tb = self.ctx.none();
+                    self.call_method(
+                        context_manager.clone(),
+                        "__exit__",
+                        vec![exc_type, exc_val, exc_tb],
+                    ).unwrap();
+                    None
+                } else {
+                    panic!("Block stack is incorrect, expected a with block");
+                }
+            }
             bytecode::Instruction::PopBlock => {
                 self.pop_block();
                 None