@@ -33,9 +33,9 @@ impl PyAsyncGen {
3333 & self . inner
3434 }
3535
36- pub fn new ( frame : FrameRef , name : PyStrRef ) -> Self {
36+ pub fn new ( frame : FrameRef , name : PyStrRef , qualname : PyStrRef ) -> Self {
3737 Self {
38- inner : Coro :: new ( frame, name) ,
38+ inner : Coro :: new ( frame, name, qualname ) ,
3939 running_async : AtomicCell :: new ( false ) ,
4040 }
4141 }
@@ -50,6 +50,16 @@ impl PyAsyncGen {
5050 self . inner . set_name ( name)
5151 }
5252
53+ #[ pygetset]
54+ fn __qualname__ ( & self ) -> PyStrRef {
55+ self . inner . qualname ( )
56+ }
57+
58+ #[ pygetset( setter) ]
59+ fn set___qualname__ ( & self , qualname : PyStrRef ) {
60+ self . inner . set_qualname ( qualname)
61+ }
62+
5363 #[ pygetset]
5464 fn ag_await ( & self , _vm : & VirtualMachine ) -> Option < PyObjectRef > {
5565 self . inner . frame ( ) . yield_from_target ( )
@@ -424,8 +434,151 @@ impl IterNext for PyAsyncGenAThrow {
424434 }
425435}
426436
437+ /// Awaitable wrapper for anext() builtin with default value.
438+ /// When StopAsyncIteration is raised, it converts it to StopIteration(default).
439+ #[ pyclass( module = false , name = "anext_awaitable" ) ]
440+ #[ derive( Debug ) ]
441+ pub struct PyAnextAwaitable {
442+ wrapped : PyObjectRef ,
443+ default_value : PyObjectRef ,
444+ }
445+
446+ impl PyPayload for PyAnextAwaitable {
447+ #[ inline]
448+ fn class ( ctx : & Context ) -> & ' static Py < PyType > {
449+ ctx. types . anext_awaitable
450+ }
451+ }
452+
453+ #[ pyclass( with( IterNext , Iterable ) ) ]
454+ impl PyAnextAwaitable {
455+ pub fn new ( wrapped : PyObjectRef , default_value : PyObjectRef ) -> Self {
456+ Self {
457+ wrapped,
458+ default_value,
459+ }
460+ }
461+
462+ #[ pymethod( name = "__await__" ) ]
463+ fn r#await ( zelf : PyRef < Self > , _vm : & VirtualMachine ) -> PyRef < Self > {
464+ zelf
465+ }
466+
467+ /// Get the awaitable iterator from wrapped object.
468+ // = anextawaitable_getiter.
469+ fn get_awaitable_iter ( & self , vm : & VirtualMachine ) -> PyResult {
470+ use crate :: builtins:: PyCoroutine ;
471+ use crate :: protocol:: PyIter ;
472+
473+ let wrapped = & self . wrapped ;
474+
475+ // If wrapped is already an async_generator_asend, it's an iterator
476+ if wrapped. class ( ) . is ( vm. ctx . types . async_generator_asend )
477+ || wrapped. class ( ) . is ( vm. ctx . types . async_generator_athrow )
478+ {
479+ return Ok ( wrapped. clone ( ) ) ;
480+ }
481+
482+ // _PyCoro_GetAwaitableIter equivalent
483+ let awaitable = if wrapped. class ( ) . is ( vm. ctx . types . coroutine_type ) {
484+ // Coroutine - get __await__ later
485+ wrapped. clone ( )
486+ } else {
487+ // Try to get __await__ method
488+ if let Some ( await_method) = vm. get_method ( wrapped. clone ( ) , identifier ! ( vm, __await__) ) {
489+ await_method?. call ( ( ) , vm) ?
490+ } else {
491+ return Err ( vm. new_type_error ( format ! (
492+ "object {} can't be used in 'await' expression" ,
493+ wrapped. class( ) . name( )
494+ ) ) ) ;
495+ }
496+ } ;
497+
498+ // If awaitable is a coroutine, get its __await__
499+ if awaitable. class ( ) . is ( vm. ctx . types . coroutine_type ) {
500+ let coro_await = vm. call_method ( & awaitable, "__await__" , ( ) ) ?;
501+ // Check that __await__ returned an iterator
502+ if !PyIter :: check ( & coro_await) {
503+ return Err ( vm. new_type_error ( "__await__ returned a non-iterable" ) ) ;
504+ }
505+ return Ok ( coro_await) ;
506+ }
507+
508+ // Check the result is an iterator, not a coroutine
509+ if awaitable. downcast_ref :: < PyCoroutine > ( ) . is_some ( ) {
510+ return Err ( vm. new_type_error ( "__await__() returned a coroutine" ) ) ;
511+ }
512+
513+ // Check that the result is an iterator
514+ if !PyIter :: check ( & awaitable) {
515+ return Err ( vm. new_type_error ( format ! (
516+ "__await__() returned non-iterator of type '{}'" ,
517+ awaitable. class( ) . name( )
518+ ) ) ) ;
519+ }
520+
521+ Ok ( awaitable)
522+ }
523+
524+ #[ pymethod]
525+ fn send ( & self , val : PyObjectRef , vm : & VirtualMachine ) -> PyResult {
526+ let awaitable = self . get_awaitable_iter ( vm) ?;
527+ let result = vm. call_method ( & awaitable, "send" , ( val, ) ) ;
528+ self . handle_result ( result, vm)
529+ }
530+
531+ #[ pymethod]
532+ fn throw (
533+ & self ,
534+ exc_type : PyObjectRef ,
535+ exc_val : OptionalArg ,
536+ exc_tb : OptionalArg ,
537+ vm : & VirtualMachine ,
538+ ) -> PyResult {
539+ let awaitable = self . get_awaitable_iter ( vm) ?;
540+ let result = vm. call_method (
541+ & awaitable,
542+ "throw" ,
543+ (
544+ exc_type,
545+ exc_val. unwrap_or_none ( vm) ,
546+ exc_tb. unwrap_or_none ( vm) ,
547+ ) ,
548+ ) ;
549+ self . handle_result ( result, vm)
550+ }
551+
552+ #[ pymethod]
553+ fn close ( & self , vm : & VirtualMachine ) -> PyResult < ( ) > {
554+ if let Ok ( awaitable) = self . get_awaitable_iter ( vm) {
555+ let _ = vm. call_method ( & awaitable, "close" , ( ) ) ;
556+ }
557+ Ok ( ( ) )
558+ }
559+
560+ /// Convert StopAsyncIteration to StopIteration(default_value)
561+ fn handle_result ( & self , result : PyResult , vm : & VirtualMachine ) -> PyResult {
562+ match result {
563+ Ok ( value) => Ok ( value) ,
564+ Err ( exc) if exc. fast_isinstance ( vm. ctx . exceptions . stop_async_iteration ) => {
565+ Err ( vm. new_stop_iteration ( Some ( self . default_value . clone ( ) ) ) )
566+ }
567+ Err ( exc) => Err ( exc) ,
568+ }
569+ }
570+ }
571+
572+ impl SelfIter for PyAnextAwaitable { }
573+ impl IterNext for PyAnextAwaitable {
574+ fn next ( zelf : & Py < Self > , vm : & VirtualMachine ) -> PyResult < PyIterReturn > {
575+ PyIterReturn :: from_pyresult ( zelf. send ( vm. ctx . none ( ) , vm) , vm)
576+ }
577+ }
578+
427579pub fn init ( ctx : & Context ) {
428580 PyAsyncGen :: extend_class ( ctx, ctx. types . async_generator ) ;
429581 PyAsyncGenASend :: extend_class ( ctx, ctx. types . async_generator_asend ) ;
430582 PyAsyncGenAThrow :: extend_class ( ctx, ctx. types . async_generator_athrow ) ;
583+ PyAnextAwaitable :: extend_class ( ctx, ctx. types . anext_awaitable ) ;
431584}
0 commit comments