1
1
//! Python coroutine implementation, used notably when wrapping `async fn`
2
2
//! with `#[pyfunction]`/`#[pymethods]`.
3
3
use std:: {
4
- any:: Any ,
5
4
future:: Future ,
6
5
panic,
7
6
pin:: Pin ,
8
7
sync:: Arc ,
9
- task:: { Context , Poll } ,
8
+ task:: { Context , Poll , Waker } ,
10
9
} ;
11
10
12
- use futures_util:: FutureExt ;
13
11
use pyo3_macros:: { pyclass, pymethods} ;
14
12
15
13
use crate :: {
16
- coroutine:: waker:: AsyncioWaker ,
14
+ coroutine:: { cancel :: ThrowCallback , waker:: AsyncioWaker } ,
17
15
exceptions:: { PyAttributeError , PyRuntimeError , PyStopIteration } ,
18
16
panic:: PanicException ,
19
17
pyclass:: IterNextOutput ,
@@ -24,20 +22,17 @@ use crate::{
24
22
pub ( crate ) mod cancel;
25
23
mod waker;
26
24
27
- use crate :: coroutine:: cancel:: ThrowCallback ;
28
25
pub use cancel:: CancelHandle ;
29
26
30
27
const COROUTINE_REUSED_ERROR : & str = "cannot reuse already awaited coroutine" ;
31
28
32
- type FutureOutput = Result < PyResult < PyObject > , Box < dyn Any + Send > > ;
33
-
34
29
/// Python coroutine wrapping a [`Future`].
35
30
#[ pyclass( crate = "crate" ) ]
36
31
pub struct Coroutine {
37
32
name : Option < Py < PyString > > ,
38
33
qualname_prefix : Option < & ' static str > ,
39
34
throw_callback : Option < ThrowCallback > ,
40
- future : Option < Pin < Box < dyn Future < Output = FutureOutput > + Send > > > ,
35
+ future : Option < Pin < Box < dyn Future < Output = PyResult < PyObject > > + Send > > > ,
41
36
waker : Option < Arc < AsyncioWaker > > ,
42
37
}
43
38
@@ -68,7 +63,7 @@ impl Coroutine {
68
63
name,
69
64
qualname_prefix,
70
65
throw_callback,
71
- future : Some ( Box :: pin ( panic :: AssertUnwindSafe ( wrap) . catch_unwind ( ) ) ) ,
66
+ future : Some ( Box :: pin ( wrap) ) ,
72
67
waker : None ,
73
68
}
74
69
}
@@ -98,22 +93,28 @@ impl Coroutine {
98
93
} else {
99
94
self . waker = Some ( Arc :: new ( AsyncioWaker :: new ( ) ) ) ;
100
95
}
101
- let waker = futures_util :: task :: waker ( self . waker . clone ( ) . unwrap ( ) ) ;
96
+ let waker = Waker :: from ( self . waker . clone ( ) . unwrap ( ) ) ;
102
97
// poll the Rust future and forward its results if ready
103
- if let Poll :: Ready ( res) = future_rs. as_mut ( ) . poll ( & mut Context :: from_waker ( & waker) ) {
104
- self . close ( ) ;
105
- return match res {
106
- Ok ( res) => Ok ( IterNextOutput :: Return ( res?) ) ,
107
- Err ( err) => Err ( PanicException :: from_panic_payload ( err) ) ,
108
- } ;
98
+ // polling is UnwindSafe because the future is dropped in case of panic
99
+ let poll = || future_rs. as_mut ( ) . poll ( & mut Context :: from_waker ( & waker) ) ;
100
+ match panic:: catch_unwind ( panic:: AssertUnwindSafe ( poll) ) {
101
+ Ok ( Poll :: Ready ( res) ) => {
102
+ self . close ( ) ;
103
+ return Ok ( IterNextOutput :: Return ( res?) ) ;
104
+ }
105
+ Err ( err) => {
106
+ self . close ( ) ;
107
+ return Err ( PanicException :: from_panic_payload ( err) ) ;
108
+ }
109
+ _ => { }
109
110
}
110
111
// otherwise, initialize the waker `asyncio.Future`
111
112
if let Some ( future) = self . waker . as_ref ( ) . unwrap ( ) . initialize_future ( py) ? {
112
113
// `asyncio.Future` must be awaited; fortunately, it implements `__iter__ = __await__`
113
114
// and will yield itself if its result has not been set in polling above
114
115
if let Some ( future) = PyIterator :: from_object ( future) . unwrap ( ) . next ( ) {
115
116
// future has not been leaked into Python for now, and Rust code can only call
116
- // `set_result(None)` in `ArcWake ` implementation, so it's safe to unwrap
117
+ // `set_result(None)` in `Wake ` implementation, so it's safe to unwrap
117
118
return Ok ( IterNextOutput :: Yield ( future. unwrap ( ) . into ( ) ) ) ;
118
119
}
119
120
}
0 commit comments