Initial work on 'with' context manager logic

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

@@ -291,14 +291,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,22 @@
+
+ls = []
+
+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!"
+
+
+with ContextManager() as c:
+    print(c)
+
+assert ls == [1, 3, 2]

vm/src/bytecode.rs

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

vm/src/compile.rs

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

vm/src/frame.rs

@@ -16,6 +16,7 @@ pub enum Block {
     #[allow(dead_code)]
     // TODO: Implement try/except blocks
     TryExcept { handler: bytecode::Label },
+    With { end: bytecode::Label },
 }
 
 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:
+        let exit_method = self.pop_value();
+        let args = PyFuncArgs {
+            args: vec![],
+            kwargs: vec![],
+        };
+        // 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,8 @@ impl VirtualMachine {
                     self.jump(&handler);
                     return None;
                 }
-                Some(_) => {}
+                Some(Block::With { .. }) => { self.with_exit(); }
+                Some(Block::Loop { .. }) => {}
                 None => break,
             }
         }
@@ -823,6 +836,10 @@ impl VirtualMachine {
                 self.push_block(Block::TryExcept { handler: *handler });
                 None
             }
+            bytecode::Instruction::SetupWith { end } => {
+                self.push_block(Block::With { end: *end });
+                None
+            }
             bytecode::Instruction::PopBlock => {
                 self.pop_block();
                 None