diff --git a/parser/src/ast.rs b/parser/src/ast.rs index 2f811322eb..4f821d2fe5 100644 --- a/parser/src/ast.rs +++ b/parser/src/ast.rs @@ -74,7 +74,7 @@ pub enum Statement { orelse: Option>, }, With { - items: Expression, + items: Vec, body: Vec, }, For { @@ -106,6 +106,12 @@ pub enum Statement { }, } +#[derive(Debug, PartialEq)] +pub struct WithItem { + pub context_expr: Expression, + pub optional_vars: Option, +} + #[derive(Debug, PartialEq, Clone)] pub enum Expression { BoolOp { diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop index 9596f3d139..9f988ce8b8 100644 --- a/parser/src/python.lalrpop +++ b/parser/src/python.lalrpop @@ -299,14 +299,28 @@ ExceptClause: ast::ExceptHandler = { }; WithStatement: ast::LocatedStatement = { - "with" "as" <_e:Expression> ":" => { + "with" ":" => { + 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 = { + => { + let optional_vars = match n { + Some(val) => Some(val.1), + None => None, + }; + ast::WithItem { context_expr: t, optional_vars } + }, +}; + FuncDef: ast::LocatedStatement = { "def" ":" => { ast::LocatedStatement { diff --git a/tests/snippets/with.py b/tests/snippets/with.py new file mode 100644 index 0000000000..f4efa86e79 --- /dev/null +++ b/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] diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs index 22f0fb2a80..debd63c371 100644 --- a/vm/src/bytecode.rs +++ b/vm/src/bytecode.rs @@ -118,6 +118,12 @@ pub enum Instruction { SetupExcept { handler: Label, }, + SetupWith { + end: Label, + }, + CleanupWith { + end: Label, + }, PopBlock, Raise { argc: usize, diff --git a/vm/src/compile.rs b/vm/src/compile.rs index c378246159..f3922d03cc 100644 --- a/vm/src/compile.rs +++ b/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, diff --git a/vm/src/frame.rs b/vm/src/frame.rs index fb76f48a42..9fc0e8d34b 100644 --- a/vm/src/frame.rs +++ b/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 { diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 6ab5ca91aa..2984ec7fa4 100644 --- a/vm/src/vm.rs +++ b/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