diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs index 4e760bba2..bd166ff8d 100644 --- a/bytecode/src/bytecode.rs +++ b/bytecode/src/bytecode.rs @@ -278,6 +278,8 @@ pub enum Instruction { SetupAsyncWith { end: Label, }, + GetAIter, + GetANext, } use self::Instruction::*; @@ -315,6 +317,7 @@ pub enum ComparisonOperator { NotIn, Is, IsNot, + ExceptionMatch, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -558,6 +561,8 @@ impl Instruction { PopException => w!(PopException), Reverse { amount } => w!(Reverse, amount), GetAwaitable => w!(GetAwaitable), + GetAIter => w!(GetAIter), + GetANext => w!(GetANext), } } } diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs index 345a04d85..e013f597c 100644 --- a/compiler/src/compile.rs +++ b/compiler/src/compile.rs @@ -442,13 +442,7 @@ impl Compiler { iter, body, orelse, - } => { - if *is_async { - unimplemented!("async for"); - } else { - self.compile_for(target, iter, body, orelse)? - } - } + } => self.compile_for(target, iter, body, orelse, *is_async)?, Raise { exception, cause } => match exception { Some(value) => { self.compile_expression(value)?; @@ -736,14 +730,9 @@ impl Compiler { self.emit(Instruction::Duplicate); // Check exception type: - self.emit(Instruction::LoadName { - name: String::from("isinstance"), - scope: bytecode::NameScope::Global, - }); - self.emit(Instruction::Rotate { amount: 2 }); self.compile_expression(exc_type)?; - self.emit(Instruction::CallFunction { - typ: CallType::Positional(2), + self.emit(Instruction::CompareOperation { + op: bytecode::ComparisonOperator::ExceptionMatch, }); // We cannot handle this exception type: @@ -1108,11 +1097,13 @@ impl Compiler { iter: &ast::Expression, body: &[ast::Statement], orelse: &Option>, + is_async: bool, ) -> Result<(), CompileError> { // Start loop let start_label = self.new_label(); let else_label = self.new_label(); let end_label = self.new_label(); + self.emit(Instruction::SetupLoop { start: start_label, end: end_label, @@ -1121,19 +1112,57 @@ impl Compiler { // The thing iterated: self.compile_expression(iter)?; - // Retrieve Iterator - self.emit(Instruction::GetIter); + if is_async { + let check_asynciter_label = self.new_label(); + let body_label = self.new_label(); - self.set_label(start_label); - self.emit(Instruction::ForIter { target: else_label }); + self.emit(Instruction::GetAIter); - // Start of loop iteration, set targets: - self.compile_store(target)?; + self.set_label(start_label); + self.emit(Instruction::SetupExcept { + handler: check_asynciter_label, + }); + self.emit(Instruction::GetANext); + self.emit(Instruction::LoadConst { + value: bytecode::Constant::None, + }); + self.emit(Instruction::YieldFrom); + self.compile_store(target)?; + self.emit(Instruction::PopBlock); + self.emit(Instruction::Jump { target: body_label }); - let was_in_loop = self.in_loop; - self.in_loop = true; - self.compile_statements(body)?; - self.in_loop = was_in_loop; + self.set_label(check_asynciter_label); + self.emit(Instruction::Duplicate); + self.emit(Instruction::LoadName { + name: "StopAsyncIteration".to_string(), + scope: bytecode::NameScope::Global, + }); + self.emit(Instruction::CompareOperation { + op: bytecode::ComparisonOperator::ExceptionMatch, + }); + self.emit(Instruction::JumpIfTrue { target: else_label }); + self.emit(Instruction::Raise { argc: 0 }); + + let was_in_loop = self.in_loop; + self.in_loop = true; + self.set_label(body_label); + self.compile_statements(body)?; + self.in_loop = was_in_loop; + } else { + // Retrieve Iterator + self.emit(Instruction::GetIter); + + self.set_label(start_label); + self.emit(Instruction::ForIter { target: else_label }); + + // Start of loop iteration, set targets: + self.compile_store(target)?; + + let was_in_loop = self.in_loop; + self.in_loop = true; + self.compile_statements(body)?; + self.in_loop = was_in_loop; + } self.emit(Instruction::Jump { target: start_label, @@ -1144,6 +1173,9 @@ impl Compiler { self.compile_statements(orelse)?; } self.set_label(end_label); + if is_async { + self.emit(Instruction::Pop); + } Ok(()) } diff --git a/tests/snippets/async_stuff.py b/tests/snippets/async_stuff.py index b28a065da..1100ea83f 100644 --- a/tests/snippets/async_stuff.py +++ b/tests/snippets/async_stuff.py @@ -19,11 +19,27 @@ class ContextManager: ls = [] +class AIterWrap: + def __init__(self, obj): + self._it = iter(obj) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + value = next(self._it) + except StopIteration: + raise StopAsyncIteration + return value + + async def a(s, m): async with ContextManager() as b: print(f"val = {b}") await asyncio.sleep(s) - for _ in range(0, 2): + async for i in AIterWrap(range(0, 2)): + print(i) ls.append(m) await asyncio.sleep(1) diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 03fa38c25..e330ab462 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -5,7 +5,7 @@ use indexmap::IndexMap; use itertools::Itertools; use crate::bytecode; -use crate::function::PyFuncArgs; +use crate::function::{single_or_tuple_any, PyFuncArgs}; use crate::obj::objbool; use crate::obj::objcode::PyCodeRef; use crate::obj::objcoroutine::PyCoroutine; @@ -278,7 +278,7 @@ impl Frame { trace!("======="); } - match &instruction { + match instruction { bytecode::Instruction::LoadConst { ref value } => { let obj = vm.ctx.unwrap_constant(value); self.push_value(obj); @@ -503,8 +503,7 @@ impl Frame { } bytecode::Instruction::GetAwaitable => { let awaited_obj = self.pop_value(); - let awaitable = if awaited_obj.payload_is::() - { + let awaitable = if awaited_obj.payload_is::() { awaited_obj } else { let await_method = @@ -519,6 +518,23 @@ impl Frame { self.push_value(awaitable); Ok(None) } + bytecode::Instruction::GetAIter => { + let aiterable = self.pop_value(); + let aiter = vm.call_method(&aiterable, "__aiter__", vec![])?; + self.push_value(aiter); + Ok(None) + } + bytecode::Instruction::GetANext => { + let aiter = self.last_value(); + let awaitable = vm.call_method(&aiter, "__anext__", vec![])?; + let awaitable = if awaitable.payload_is::() { + awaitable + } else { + vm.call_method(&awaitable, "__await__", vec![])? + }; + self.push_value(awaitable); + Ok(None) + } bytecode::Instruction::ForIter { target } => self.execute_for_iter(vm, *target), bytecode::Instruction::MakeFunction => self.execute_make_function(vm), bytecode::Instruction::CallFunction { typ } => self.execute_call_function(vm, typ), @@ -1251,6 +1267,25 @@ impl Frame { !a.is(&b) } + fn exc_match( + &self, + vm: &VirtualMachine, + exc: PyObjectRef, + exc_type: PyObjectRef, + ) -> PyResult { + single_or_tuple_any( + exc_type, + |cls: PyClassRef| vm.isinstance(&exc, &cls), + |o| { + format!( + "isinstance() arg 2 must be a type or tuple of types, not {}", + o.class() + ) + }, + vm, + ) + } + #[cfg_attr(feature = "flame-it", flame("Frame"))] fn execute_compare( &self, @@ -1266,10 +1301,11 @@ impl Frame { bytecode::ComparisonOperator::LessOrEqual => vm._le(a, b)?, bytecode::ComparisonOperator::Greater => vm._gt(a, b)?, bytecode::ComparisonOperator::GreaterOrEqual => vm._ge(a, b)?, - bytecode::ComparisonOperator::Is => vm.ctx.new_bool(self._is(a, b)), - bytecode::ComparisonOperator::IsNot => vm.ctx.new_bool(self._is_not(a, b)), - bytecode::ComparisonOperator::In => vm.ctx.new_bool(self._in(vm, a, b)?), - bytecode::ComparisonOperator::NotIn => vm.ctx.new_bool(self._not_in(vm, a, b)?), + bytecode::ComparisonOperator::Is => vm.new_bool(self._is(a, b)), + bytecode::ComparisonOperator::IsNot => vm.new_bool(self._is_not(a, b)), + bytecode::ComparisonOperator::In => vm.new_bool(self._in(vm, a, b)?), + bytecode::ComparisonOperator::NotIn => vm.new_bool(self._not_in(vm, a, b)?), + bytecode::ComparisonOperator::ExceptionMatch => vm.new_bool(self.exc_match(vm, a, b)?), }; self.push_value(value);