diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py index 7d21a7a2f..ea17d8d36 100644 --- a/Lib/test/test_coroutines.py +++ b/Lib/test/test_coroutines.py @@ -1668,7 +1668,6 @@ class CoroutineTest(unittest.TestCase): self.assertEqual(sys.getrefcount(aiter), refs_before) - @unittest.expectedFailure # TODO: RUSTPYTHON; refcount leak in async for/with def test_for_6(self): I = 0 diff --git a/crates/vm/src/builtins/asyncgenerator.rs b/crates/vm/src/builtins/asyncgenerator.rs index 9903c9280..4aafa267a 100644 --- a/crates/vm/src/builtins/asyncgenerator.rs +++ b/crates/vm/src/builtins/asyncgenerator.rs @@ -180,6 +180,7 @@ impl PyRef { exc_tb: OptionalArg, vm: &VirtualMachine, ) -> PyResult { + warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?; PyAsyncGen::init_hooks(&self, vm)?; Ok(PyAsyncGenAThrow { ag: self, @@ -328,7 +329,7 @@ impl PyAsyncGenASend { let res = self.ag.inner.send(self.ag.as_object(), val, vm); let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm); if res.is_err() { - self.close(); + self.set_closed(); } res } @@ -341,8 +342,23 @@ impl PyAsyncGenASend { exc_tb: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - if let AwaitableState::Closed = self.state.load() { - return Err(vm.new_runtime_error("cannot reuse already awaited __anext__()/asend()")); + match self.state.load() { + AwaitableState::Closed => { + return Err( + vm.new_runtime_error("cannot reuse already awaited __anext__()/asend()") + ); + } + AwaitableState::Init => { + if self.ag.running_async.load() { + self.state.store(AwaitableState::Closed); + return Err( + vm.new_runtime_error("anext(): asynchronous generator is already running") + ); + } + self.ag.running_async.store(true); + self.state.store(AwaitableState::Iter); + } + AwaitableState::Iter => {} } warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?; @@ -355,13 +371,36 @@ impl PyAsyncGenASend { ); let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm); if res.is_err() { - self.close(); + self.set_closed(); } res } #[pymethod] - fn close(&self) { + fn close(&self, vm: &VirtualMachine) -> PyResult<()> { + if matches!(self.state.load(), AwaitableState::Closed) { + return Ok(()); + } + let result = self.throw( + vm.ctx.exceptions.generator_exit.to_owned().into(), + OptionalArg::Missing, + OptionalArg::Missing, + vm, + ); + match result { + Ok(_) => Err(vm.new_runtime_error("coroutine ignored GeneratorExit")), + Err(e) + if e.fast_isinstance(vm.ctx.exceptions.stop_iteration) + || e.fast_isinstance(vm.ctx.exceptions.stop_async_iteration) + || e.fast_isinstance(vm.ctx.exceptions.generator_exit) => + { + Ok(()) + } + Err(e) => Err(e), + } + } + + fn set_closed(&self) { self.state.store(AwaitableState::Closed); } } @@ -472,6 +511,32 @@ impl PyAsyncGenAThrow { exc_tb: OptionalArg, vm: &VirtualMachine, ) -> PyResult { + match self.state.load() { + AwaitableState::Closed => { + return Err( + vm.new_runtime_error("cannot reuse already awaited aclose()/athrow()") + ); + } + AwaitableState::Init => { + if self.ag.running_async.load() { + self.state.store(AwaitableState::Closed); + let msg = if self.aclose { + "aclose(): asynchronous generator is already running" + } else { + "athrow(): asynchronous generator is already running" + }; + return Err(vm.new_runtime_error(msg.to_owned())); + } + if self.ag.inner.closed() { + self.state.store(AwaitableState::Closed); + return Err(vm.new_stop_iteration(None)); + } + self.ag.running_async.store(true); + self.state.store(AwaitableState::Iter); + } + AwaitableState::Iter => {} + } + warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?; let ret = self.ag.inner.throw( self.ag.as_object(), @@ -493,8 +558,27 @@ impl PyAsyncGenAThrow { } #[pymethod] - fn close(&self) { - self.state.store(AwaitableState::Closed); + fn close(&self, vm: &VirtualMachine) -> PyResult<()> { + if matches!(self.state.load(), AwaitableState::Closed) { + return Ok(()); + } + let result = self.throw( + vm.ctx.exceptions.generator_exit.to_owned().into(), + OptionalArg::Missing, + OptionalArg::Missing, + vm, + ); + match result { + Ok(_) => Err(vm.new_runtime_error("coroutine ignored GeneratorExit")), + Err(e) + if e.fast_isinstance(vm.ctx.exceptions.stop_iteration) + || e.fast_isinstance(vm.ctx.exceptions.stop_async_iteration) + || e.fast_isinstance(vm.ctx.exceptions.generator_exit) => + { + Ok(()) + } + Err(e) => Err(e), + } } fn ignored_close(&self, res: &PyResult) -> bool {