diff --git a/vm/src/builtins/asyncgenerator.rs b/vm/src/builtins/asyncgenerator.rs index b9b8493a1..5b4f6265a 100644 --- a/vm/src/builtins/asyncgenerator.rs +++ b/vm/src/builtins/asyncgenerator.rs @@ -26,7 +26,7 @@ impl PyPayload for PyAsyncGen { } } -#[pyclass(with(Constructor, Representable))] +#[pyclass(with(PyRef, Constructor, Representable))] impl PyAsyncGen { pub fn as_coro(&self) -> &Coro { &self.inner @@ -49,59 +49,6 @@ impl PyAsyncGen { self.inner.set_name(name) } - #[pymethod(magic)] - fn aiter(zelf: PyRef, _vm: &VirtualMachine) -> PyRef { - zelf - } - - #[pymethod(magic)] - fn anext(zelf: PyRef, vm: &VirtualMachine) -> PyAsyncGenASend { - Self::asend(zelf, vm.ctx.none(), vm) - } - - #[pymethod] - fn asend(zelf: PyRef, value: PyObjectRef, _vm: &VirtualMachine) -> PyAsyncGenASend { - PyAsyncGenASend { - ag: zelf, - state: AtomicCell::new(AwaitableState::Init), - value, - } - } - - #[pymethod] - fn athrow( - zelf: PyRef, - exc_type: PyObjectRef, - exc_val: OptionalArg, - exc_tb: OptionalArg, - vm: &VirtualMachine, - ) -> PyAsyncGenAThrow { - PyAsyncGenAThrow { - ag: zelf, - aclose: false, - state: AtomicCell::new(AwaitableState::Init), - value: ( - exc_type, - exc_val.unwrap_or_none(vm), - exc_tb.unwrap_or_none(vm), - ), - } - } - - #[pymethod] - fn aclose(zelf: PyRef, vm: &VirtualMachine) -> PyAsyncGenAThrow { - PyAsyncGenAThrow { - ag: zelf, - aclose: true, - state: AtomicCell::new(AwaitableState::Init), - value: ( - vm.ctx.exceptions.generator_exit.to_owned().into(), - vm.ctx.none(), - vm.ctx.none(), - ), - } - } - #[pygetset] fn ag_await(&self, _vm: &VirtualMachine) -> Option { self.inner.frame().yield_from_target() @@ -125,6 +72,62 @@ impl PyAsyncGen { } } +#[pyclass] +impl PyRef { + #[pymethod(magic)] + fn aiter(self, _vm: &VirtualMachine) -> PyRef { + self + } + + #[pymethod(magic)] + fn anext(self, vm: &VirtualMachine) -> PyAsyncGenASend { + Self::asend(self, vm.ctx.none(), vm) + } + + #[pymethod] + fn asend(self, value: PyObjectRef, _vm: &VirtualMachine) -> PyAsyncGenASend { + PyAsyncGenASend { + ag: self, + state: AtomicCell::new(AwaitableState::Init), + value, + } + } + + #[pymethod] + fn athrow( + self, + exc_type: PyObjectRef, + exc_val: OptionalArg, + exc_tb: OptionalArg, + vm: &VirtualMachine, + ) -> PyAsyncGenAThrow { + PyAsyncGenAThrow { + ag: self, + aclose: false, + state: AtomicCell::new(AwaitableState::Init), + value: ( + exc_type, + exc_val.unwrap_or_none(vm), + exc_tb.unwrap_or_none(vm), + ), + } + } + + #[pymethod] + fn aclose(self, vm: &VirtualMachine) -> PyAsyncGenAThrow { + PyAsyncGenAThrow { + ag: self, + aclose: true, + state: AtomicCell::new(AwaitableState::Init), + value: ( + vm.ctx.exceptions.generator_exit.to_owned().into(), + vm.ctx.none(), + vm.ctx.none(), + ), + } + } +} + impl Representable for PyAsyncGen { #[inline] fn repr_str(zelf: &Py, vm: &VirtualMachine) -> PyResult {