diff --git a/parser/src/ast.rs b/parser/src/ast.rs index 2f811322e..4f821d2fe 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 b7238e8b7..3135d7d15 100644 --- a/parser/src/python.lalrpop +++ b/parser/src/python.lalrpop @@ -291,14 +291,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 000000000..c26a1c724 --- /dev/null +++ b/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] diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs index c0d5d286d..f01f6ae4a 100644 --- a/vm/src/bytecode.rs +++ b/vm/src/bytecode.rs @@ -118,6 +118,9 @@ pub enum Instruction { SetupExcept { handler: Label, }, + SetupWith { + end: Label, + }, PopBlock, Raise { argc: usize, diff --git a/vm/src/compile.rs b/vm/src/compile.rs index 3208f9700..c9494a0e7 100644 --- a/vm/src/compile.rs +++ b/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, diff --git a/vm/src/frame.rs b/vm/src/frame.rs index fb76f48a4..4e4120944 100644 --- a/vm/src/frame.rs +++ b/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 { diff --git a/vm/src/vm.rs b/vm/src/vm.rs index a50b2eb4b..b165678a1 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: + 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