├── .gitignore ├── .rustfmt.toml ├── Cargo.toml ├── LICENSE ├── README.md ├── pyo3-async-macros ├── Cargo.toml └── src │ └── lib.rs └── src ├── allow_threads.rs ├── async_generator.rs ├── asyncio.rs ├── coroutine.rs ├── lib.rs ├── sniffio.rs ├── trio.rs └── utils.rs /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | debug/ 4 | target/ 5 | 6 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 7 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 8 | Cargo.lock 9 | 10 | # These are backup files generated by rustfmt 11 | **/*.rs.bk 12 | 13 | # MSVC Windows builds of rustc generate these, which store debugging information 14 | *.pdb 15 | 16 | # Jetbrains 17 | .idea 18 | 19 | # Maturin tests 20 | venv -------------------------------------------------------------------------------- /.rustfmt.toml: -------------------------------------------------------------------------------- 1 | imports_granularity = "Crate" 2 | group_imports = "StdExternalCrate" 3 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "pyo3-async" 3 | description = "PyO3 bindings to various Python asynchronous frameworks." 4 | readme = "README.md" 5 | version.workspace = true 6 | edition.workspace = true 7 | exclude.workspace = true 8 | homepage.workspace = true 9 | keywords.workspace = true 10 | categories.workspace = true 11 | license.workspace = true 12 | repository.workspace = true 13 | 14 | [workspace.package] 15 | version = "0.3.2" 16 | edition = "2021" 17 | exclude = [".*"] 18 | homepage = "https://github.com/wyfo/pyo3-async" 19 | keywords = [ 20 | "pyo3", 21 | "python", 22 | "cpython", 23 | "ffi", 24 | "async", 25 | ] 26 | categories = [ 27 | "api-bindings", 28 | "asynchronous", 29 | "development-tools::ffi", 30 | ] 31 | license = "MIT" 32 | repository = "https://github.com/wyfo/pyo3-async" 33 | 34 | [features] 35 | default = ["macros", "allow-threads"] 36 | macros = ["dep:pyo3-async-macros"] 37 | allow-threads = ["dep:pin-project"] 38 | 39 | [dependencies] 40 | futures = "0.3" 41 | pin-project = { version = "1", optional = true } 42 | pyo3 = ">=0.18,<0.21" 43 | pyo3-async-macros = { path = "pyo3-async-macros", version = "=0.3.2", optional = true } 44 | 45 | [workspace] 46 | members = ["pyo3-async-macros"] 47 | 48 | [badges] 49 | maintenance = { "status" = "deprecated" } 50 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Joseph Perez 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pyo3-async 2 | 3 | PyO3 bindings to various Python asynchronous frameworks. 4 | 5 | ## *This crate is deprecated* 6 | 7 | *This crate was an experiment about implementing async support in PyO3.* 8 | 9 | *The experiment was successful, because async support is eventually being implemented at this moment, see [the dedicated issue](https://github.com/PyO3/pyo3/issues/1632).* 10 | 11 | *As a consequence, this crate is now deprecated. If you want to try async support before its official release, you should use the branch of the [last PR](https://github.com/PyO3/pyo3/pull/3613), or `master` (where the implementation is still incomplete, but normally stable).* 12 | 13 | ## Documentation 14 | 15 | https://docs.rs/pyo3-async/ 16 | 17 | ## How it works 18 | 19 | Asynchronous implementations are not so different in Rust and Python. Rust uses callbacks (through `std::task::Waker`) to wake up the related executor, while Python `Asyncio.Future` also has a callback registered to wake up the event loop. 20 | 21 | So, why not use Rust callback to wake up Python event loop, and vice versa ? That's all. 22 | 23 | ## Difference with [PyO3 Asyncio](ashttps://github.com/awestlake87/pyo3-asyncio) 24 | 25 | - PyO3 Asyncio requires a running asynchronous runtime on both Python and Rust side, while this crate doesn't; 26 | - PyO3 Asyncio only focus on *asyncio*, while this crate obviously support *asyncio*, but also [*trio*](https://github.com/python-trio/trio) or [*anyio*](https://github.com/agronholm/anyio); 27 | - This crate provides control over the GIL release; 28 | - This crate provides `#[pyfunction]`/`#[pymethods]` macros. 29 | 30 | ## Example 31 | 32 | You can build this module with [Maturin](https://github.com/PyO3/maturin) 33 | 34 | ```rust 35 | #[pymodule] 36 | fn example(_py: Python<'_>, m: &PyModule) -> PyResult<()> { 37 | m.add_function(wrap_pyfunction!(async_sleep_asyncio, m)?)?; 38 | m.add_function(wrap_pyfunction!(async_sleep_trio, m)?)?; 39 | m.add_function(wrap_pyfunction!(sleep_sniffio, m)?)?; 40 | m.add_function(wrap_pyfunction!(spawn_future, m)?)?; 41 | m.add_function(wrap_pyfunction!(count_asyncio, m)?)?; 42 | m.add_function(wrap_pyfunction!(count_trio, m)?)?; 43 | Ok(()) 44 | } 45 | 46 | fn tokio() -> &'static tokio::runtime::Runtime { 47 | use std::sync::OnceLock; 48 | static RT: OnceLock = OnceLock::new(); 49 | RT.get_or_init(|| tokio::runtime::Runtime::new().unwrap()) 50 | } 51 | 52 | async fn sleep(seconds: u64) { 53 | let sleep = async move { tokio::time::sleep(std::time::Duration::from_secs(seconds)).await }; 54 | tokio().spawn(sleep).await.unwrap(); 55 | } 56 | 57 | fn count(until: i32, tick: u64) -> impl Stream> + Send { 58 | futures::stream::unfold(0, move |i| async move { 59 | if i == until { 60 | return None; 61 | } 62 | sleep(tick).await; 63 | Some((PyResult::Ok(i), i + 1)) 64 | }) 65 | } 66 | 67 | // Works with async function 68 | // It generates an `async_sleep_asyncio` function to be exported (see #[pymodule] above) 69 | #[pyfunction] 70 | async fn sleep_asyncio(seconds: u64) { 71 | sleep(seconds).await; 72 | } 73 | 74 | // Specify Python async backend and GIL release 75 | #[pyfunction(trio, allow_threads)] 76 | async fn sleep_trio(seconds: u64) { 77 | sleep(seconds).await; 78 | } 79 | 80 | // Coroutine can be manually instantiated 81 | #[pyfunction] 82 | fn sleep_sniffio(seconds: u64) -> pyo3_async::sniffio::Coroutine { 83 | pyo3_async::sniffio::Coroutine::from_future(async move { 84 | sleep(seconds).await; 85 | PyResult::Ok(()) 86 | }) 87 | } 88 | 89 | #[pyfunction] 90 | fn spawn_future(fut: PyObject) { 91 | tokio().spawn(async move { 92 | pyo3_async::asyncio::FutureWrapper::new(fut, None) 93 | .await 94 | .unwrap(); 95 | println!("task done") 96 | }); 97 | } 98 | 99 | #[pyfunction] 100 | fn count_asyncio(until: i32, tick: u64) -> pyo3_async::asyncio::AsyncGenerator { 101 | pyo3_async::asyncio::AsyncGenerator::from_stream(count(until, tick)) 102 | } 103 | 104 | #[pyfunction] 105 | fn count_trio(until: i32, tick: u64) -> pyo3_async::trio::AsyncGenerator { 106 | pyo3_async::trio::AsyncGenerator::from_stream(count(until, tick)) 107 | } 108 | 109 | ``` 110 | 111 | and execute this Python code 112 | 113 | ```python 114 | import asyncio 115 | import trio 116 | import example # built with maturin 117 | 118 | async def asyncio_main(): 119 | await example.sleep_asyncio(1) 120 | # sleep 1s 121 | await example.sleep_sniffio(1) 122 | # sleep 1s 123 | example.spawn_future(asyncio.create_task(asyncio.sleep(1))) 124 | await asyncio.sleep(2) 125 | # print "done" after 1s 126 | async for i in example.count_asyncio(2, 1): 127 | print(i) 128 | # sleep 1s, print 0 129 | # sleep 1s, print 1 130 | 131 | async def trio_run(): 132 | await example.sleep_trio(1) 133 | # sleep 1s 134 | await example.sleep_sniffio(1) 135 | # sleep 1s 136 | async for i in example.count_trio(2, 1): 137 | print(i) 138 | # sleep 1s, print 0 139 | # sleep 1s, print 1 140 | 141 | asyncio.run(asyncio_main()) 142 | print("======================") 143 | trio.run(trio_run) 144 | ``` 145 | -------------------------------------------------------------------------------- /pyo3-async-macros/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "pyo3-async-macros" 3 | description = "Procedural macros implementation for pyo3-async." 4 | version.workspace = true 5 | edition.workspace = true 6 | exclude.workspace = true 7 | homepage.workspace = true 8 | keywords.workspace = true 9 | categories.workspace = true 10 | license.workspace = true 11 | repository.workspace = true 12 | 13 | [lib] 14 | proc-macro = true 15 | 16 | [dependencies] 17 | quote = "1" 18 | syn = { version = "2", features = ["full", "extra-traits"] } 19 | 20 | [dev-dependencies] 21 | pyo3 = ">=0.18,<0.21" 22 | pyo3-async = { path = ".." } 23 | -------------------------------------------------------------------------------- /pyo3-async-macros/src/lib.rs: -------------------------------------------------------------------------------- 1 | use proc_macro::TokenStream; 2 | use quote::{format_ident, quote, ToTokens}; 3 | use syn::{ 4 | parse::Parser, parse_macro_input, parse_quote, parse_quote_spanned, punctuated::Punctuated, 5 | spanned::Spanned, 6 | }; 7 | 8 | const MODULES: [&str; 3] = ["asyncio", "trio", "sniffio"]; 9 | 10 | macro_rules! unwrap { 11 | ($result:expr) => { 12 | match $result { 13 | Ok(ok) => ok, 14 | Err(err) => return err.into_compile_error().into(), 15 | } 16 | }; 17 | } 18 | 19 | struct Options { 20 | module: syn::Path, 21 | allow_threads: bool, 22 | } 23 | 24 | fn parse_options(attr: TokenStream) -> syn::Result { 25 | let mut allow_threads = false; 26 | let mut module = None; 27 | let module_parser = syn::meta::parser(|meta| { 28 | if meta.path.is_ident("allow_threads") { 29 | allow_threads = true; 30 | } else if MODULES.iter().any(|m| meta.path.is_ident(m)) { 31 | if module.is_some() { 32 | return Err(meta.error("multiple Python async backend specified")); 33 | } 34 | module = Some(meta.path); 35 | } else { 36 | return Err(meta.error("invalid option")); 37 | } 38 | Ok(()) 39 | }); 40 | module_parser.parse(attr)?; 41 | Ok(Options { 42 | module: module.unwrap_or_else(|| parse_quote!(asyncio)), 43 | allow_threads, 44 | }) 45 | } 46 | 47 | fn build_coroutine( 48 | path: impl ToTokens, 49 | attrs: &mut Vec, 50 | sig: &mut syn::Signature, 51 | block: &mut syn::Block, 52 | options: &Options, 53 | ) -> syn::Result<()> { 54 | attrs.retain(|attr| attr.meta.path().is_ident("pyo3")); 55 | let mut has_name = false; 56 | for attr in attrs.iter() { 57 | has_name |= attr 58 | .parse_args_with(Punctuated::::parse_terminated)? 59 | .into_iter() 60 | .any(|meta| matches!(meta, syn::Meta::NameValue(nv) if nv.path.is_ident("name"))); 61 | } 62 | if !has_name { 63 | let name = format!("{}", &sig.ident); 64 | attrs.push(parse_quote!(#[pyo3(name = #name)])); 65 | } 66 | let ident = sig.ident.clone(); 67 | sig.ident = format_ident!("async_{ident}"); 68 | sig.asyncness = None; 69 | let module = &options.module; 70 | let coro_path = quote!(::pyo3_async::#module::Coroutine); 71 | let params = sig.inputs.iter().map(|arg| match arg { 72 | syn::FnArg::Receiver(_) => quote!(self), 73 | syn::FnArg::Typed(syn::PatType { pat, .. }) => quote!(#pat), 74 | }); 75 | let mut future = quote!(#path(#(#params),*)); 76 | if matches!(sig.output, syn::ReturnType::Default) { 77 | future = quote!(async move {#future.await; pyo3::PyResult::Ok(())}) 78 | } 79 | if options.allow_threads { 80 | future = quote!(::pyo3_async::AllowThreads(#future)); 81 | } 82 | // return statement because `parse_quote_spanned` doesn't work otherwise 83 | block.stmts = vec![parse_quote_spanned! { block.span() => 84 | #[allow(clippy::needless_return)] 85 | return #coro_path::from_future(#future); 86 | }]; 87 | sig.output = parse_quote_spanned!(sig.output.span() => -> #coro_path); 88 | Ok(()) 89 | } 90 | 91 | /// [`pyo3::pyfunction`] with async support. 92 | /// 93 | /// Generate a additional function prefixed by `async_`, decorated by [`pyo3::pyfunction`] and 94 | /// `#[pyo3(name = ...)]`. 95 | /// 96 | /// Python async backend can be specified using macro argument (default to `asyncio`). 97 | /// If `allow_threads` is passed in arguments, GIL will be released for future polling (see 98 | /// [`AllowThreads`]) 99 | /// 100 | /// # Example 101 | /// 102 | /// ```rust 103 | /// #[pyo3_async::pyfunction(allow_threads)] 104 | /// pub async fn print(s: String) { 105 | /// println!("{s}"); 106 | /// } 107 | /// ``` 108 | /// generates 109 | /// ```rust 110 | /// pub async fn print(s: String) { 111 | /// println!("{s}"); 112 | /// } 113 | /// #[::pyo3::pyfunction] 114 | /// #[pyo3(name = "print")] 115 | /// pub fn async_print(s: String) -> ::pyo3_async::asyncio::Coroutine { 116 | /// ::pyo3_async::asyncio::Coroutine::from_future(::pyo3_async::AllowThreads( 117 | /// async move { print(s).await; Ok(()) } 118 | /// )) 119 | /// } 120 | /// ``` 121 | /// 122 | /// [`pyo3::pyfunction`]: https://docs.rs/pyo3/latest/pyo3/attr.pyfunction.html 123 | /// [`AllowThreads`]: https://docs.rs/pyo3-async/latest/pyo3_async/struct.AllowThreads.html 124 | #[proc_macro_attribute] 125 | pub fn pyfunction(attr: TokenStream, input: TokenStream) -> TokenStream { 126 | let options = unwrap!(parse_options(attr)); 127 | let mut func = parse_macro_input!(input as syn::ItemFn); 128 | if func.sig.asyncness.is_none() { 129 | return quote!(#[::pyo3::pyfunction] #func).into(); 130 | } 131 | let mut coro = func.clone(); 132 | unwrap!(build_coroutine( 133 | &func.sig.ident, 134 | &mut coro.attrs, 135 | &mut coro.sig, 136 | &mut coro.block, 137 | &options 138 | )); 139 | func.attrs.retain(|attr| !attr.meta.path().is_ident("pyo3")); 140 | let expanded = quote! { 141 | #func 142 | #[::pyo3::pyfunction] 143 | #coro 144 | }; 145 | expanded.into() 146 | } 147 | 148 | /// [`pyo3::pymethods`] with async support. 149 | /// 150 | /// For each async methods, generate a additional function prefixed by `async_`, decorated with 151 | /// `#[pyo3(name = ...)]`. Original async methods are kept in a separate impl, while the original 152 | /// impl is decorated with [`pyo3::pymethods`]. 153 | /// 154 | /// Python async backend can be specified using macro argument (default to `asyncio`). 155 | /// If `allow_threads` is passed in arguments, GIL will be released for future polling (see 156 | /// [`AllowThreads`]) 157 | /// 158 | /// # Example 159 | /// 160 | /// ```rust 161 | /// #[pyo3::pyclass] 162 | /// struct Counter(usize); 163 | /// 164 | /// #[pyo3_async::pymethods(trio)] 165 | /// impl Counter { 166 | /// fn incr_sync(&mut self) -> usize { 167 | /// self.0 += 1; 168 | /// self.0 169 | /// } 170 | /// 171 | /// // Arguments needs to implement `Send + 'static`, so `self` must be passed using `Py` 172 | /// async fn incr_async(self_: pyo3::Py) -> pyo3::PyResult { 173 | /// pyo3::Python::with_gil(|gil| { 174 | /// let mut this = self_.borrow_mut(gil); 175 | /// this.0 += 1; 176 | /// Ok(this.0) 177 | /// }) 178 | /// } 179 | /// } 180 | /// ``` 181 | /// generates 182 | /// ```rust 183 | /// #[pyo3::pyclass] 184 | /// struct Counter(usize); 185 | /// 186 | /// #[::pyo3::pymethods] 187 | /// impl Counter { 188 | /// fn incr_sync(&mut self) -> usize { 189 | /// self.0 += 1; 190 | /// self.0 191 | /// } 192 | /// 193 | /// #[pyo3(name = "incr_async")] 194 | /// fn async_incr_async(self_: pyo3::Py) -> ::pyo3_async::trio::Coroutine { 195 | /// ::pyo3_async::trio::Coroutine::from_future(Counter::incr_async(self_)) 196 | /// } 197 | /// } 198 | /// impl Counter { 199 | /// async fn incr_async(self_: pyo3::Py) -> pyo3::PyResult { 200 | /// pyo3::Python::with_gil(|gil| { 201 | /// let mut this = self_.borrow_mut(gil); 202 | /// this.0 += 1; 203 | /// Ok(this.0) 204 | /// }) 205 | /// } 206 | /// } 207 | /// ``` 208 | /// 209 | /// [`pyo3::pymethods`]: https://docs.rs/pyo3/latest/pyo3/attr.pymethods.html 210 | /// [`AllowThreads`]: https://docs.rs/pyo3-async/latest/pyo3_async/struct.AllowThreads.html 211 | #[proc_macro_attribute] 212 | pub fn pymethods(attr: TokenStream, input: TokenStream) -> TokenStream { 213 | let options = unwrap!(parse_options(attr)); 214 | let mut r#impl = parse_macro_input!(input as syn::ItemImpl); 215 | let (async_methods, items) = r#impl.items.into_iter().partition::, _>( 216 | |item| matches!(item, syn::ImplItem::Fn(func) if func.sig.asyncness.is_some()), 217 | ); 218 | r#impl.items = items; 219 | if async_methods.is_empty() { 220 | return quote!(#[::pyo3::pymethods] #r#impl).into(); 221 | } 222 | let mut async_impl = r#impl.clone(); 223 | async_impl.items = async_methods; 224 | async_impl.attrs.clear(); 225 | for item in &mut async_impl.items { 226 | let syn::ImplItem::Fn(method) = item else { 227 | unreachable!() 228 | }; 229 | let mut coro = method.clone(); 230 | let self_ty = &r#impl.self_ty; 231 | let method_name = &method.sig.ident; 232 | unwrap!(build_coroutine( 233 | quote!(#self_ty::#method_name), 234 | &mut coro.attrs, 235 | &mut coro.sig, 236 | &mut coro.block, 237 | &options 238 | )); 239 | method 240 | .attrs 241 | .retain(|attr| !attr.meta.path().is_ident("pyo3")); 242 | method.attrs.retain(|attr| { 243 | if ["getter", "classmethod", "staticmethod"] 244 | .iter() 245 | .any(|m| attr.meta.path().is_ident(m)) 246 | { 247 | coro.attrs.push(attr.clone()); 248 | return false; 249 | } 250 | true 251 | }); 252 | r#impl.items.push(syn::ImplItem::Fn(coro)); 253 | } 254 | let expanded = quote! { 255 | #[::pyo3::pymethods] 256 | #r#impl 257 | #async_impl 258 | }; 259 | expanded.into() 260 | } 261 | -------------------------------------------------------------------------------- /src/allow_threads.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | future::Future, 3 | pin::Pin, 4 | task::{Context, Poll}, 5 | }; 6 | 7 | use futures::Stream; 8 | use pin_project::pin_project; 9 | use pyo3::Python; 10 | 11 | /// Wrapper for [`Future`]/[`Stream`] that releases GIL while polling in 12 | /// [`PyFuture`](crate::PyFuture)/[`PyStream`](crate::PyStream). 13 | /// 14 | /// Can be instantiated with [`AllowThreadsExt::allow_threads`]. 15 | /// 16 | /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html 17 | #[derive(Debug)] 18 | #[repr(transparent)] 19 | #[pin_project] 20 | pub struct AllowThreads(#[pin] pub T); 21 | 22 | impl Future for AllowThreads 23 | where 24 | F: Future + Send, 25 | F::Output: Send, 26 | { 27 | type Output = F::Output; 28 | 29 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 30 | let this = self.project(); 31 | let waker = cx.waker(); 32 | Python::with_gil(|gil| gil.allow_threads(|| this.0.poll(&mut Context::from_waker(waker)))) 33 | } 34 | } 35 | 36 | impl Stream for AllowThreads 37 | where 38 | S: Stream + Send, 39 | S::Item: Send, 40 | { 41 | type Item = S::Item; 42 | 43 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 44 | let this = self.project(); 45 | let waker = cx.waker(); 46 | Python::with_gil(|gil| { 47 | gil.allow_threads(|| this.0.poll_next(&mut Context::from_waker(waker))) 48 | }) 49 | } 50 | } 51 | 52 | /// Extension trait to allow threads while polling [`Future`] or [`Stream`]. 53 | /// 54 | /// It is implemented for every types. 55 | /// 56 | /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html 57 | pub trait AllowThreadsExt: Sized { 58 | fn allow_threads(self) -> AllowThreads { 59 | AllowThreads(self) 60 | } 61 | } 62 | 63 | impl AllowThreadsExt for T {} 64 | -------------------------------------------------------------------------------- /src/async_generator.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | marker::PhantomData, 3 | pin::Pin, 4 | sync::{Arc, Mutex}, 5 | task::{ready, Context, Poll}, 6 | }; 7 | 8 | use pyo3::{exceptions::PyStopAsyncIteration, prelude::*}; 9 | 10 | use crate::{PyFuture, PyStream, ThrowCallback}; 11 | 12 | type SharedStream = Arc>>>>; 13 | 14 | struct PyStreamNext { 15 | stream: SharedStream, 16 | close: bool, 17 | } 18 | 19 | impl PyFuture for PyStreamNext { 20 | fn poll_py(self: Pin<&mut Self>, py: Python, cx: &mut Context) -> Poll> { 21 | let err = || Err(PyStopAsyncIteration::new_err(py.None())); 22 | let this = Pin::into_inner(self); 23 | let mut guard = this.stream.lock().unwrap(); 24 | let Some(ref mut stream) = *guard else { 25 | return Poll::Ready(err()); 26 | }; 27 | let opt_res = ready!(stream.as_mut().poll_next_py(py, cx)); 28 | if let Some(res) = opt_res { 29 | if this.close { 30 | *guard = None; 31 | } 32 | return Poll::Ready(res); 33 | } 34 | *guard = None; 35 | Poll::Ready(err()) 36 | } 37 | } 38 | 39 | pub(crate) trait CoroutineFactory { 40 | type Coroutine: IntoPy; 41 | fn coroutine(future: impl PyFuture + 'static) -> Self::Coroutine; 42 | } 43 | 44 | pub(crate) struct AsyncGenerator { 45 | stream: SharedStream, 46 | throw: Option, 47 | _phantom: PhantomData, 48 | } 49 | 50 | impl AsyncGenerator { 51 | pub(crate) fn new(stream: Pin>, throw: Option) -> Self { 52 | Self { 53 | stream: Arc::new(Mutex::new(Some(stream))), 54 | throw, 55 | _phantom: PhantomData, 56 | } 57 | } 58 | } 59 | 60 | impl AsyncGenerator { 61 | pub(crate) fn _next(&mut self, py: Python, close: bool) -> PyResult { 62 | let stream = self.stream.clone(); 63 | Ok(C::coroutine(PyStreamNext { stream, close }).into_py(py)) 64 | } 65 | 66 | pub(crate) fn next(&mut self, py: Python) -> PyResult { 67 | self._next(py, false) 68 | } 69 | 70 | pub(crate) fn throw(&mut self, py: Python, exc: PyErr) -> PyResult { 71 | let Some(throw) = &mut self.throw else { 72 | return Ok(C::coroutine(async move { Err::<(), _>(exc) }).into_py(py)); 73 | }; 74 | throw(py, Some(exc)); 75 | self._next(py, false) 76 | } 77 | 78 | pub(crate) fn close(&mut self, py: Python) -> PyResult { 79 | if let Some(throw) = &mut self.throw { 80 | throw(py, None); 81 | } 82 | self._next(py, true) 83 | } 84 | } 85 | -------------------------------------------------------------------------------- /src/asyncio.rs: -------------------------------------------------------------------------------- 1 | //! `asyncio` compatible coroutine and async generator implementation. 2 | use std::{ 3 | future::Future, 4 | pin::Pin, 5 | task::{ready, Context, Poll}, 6 | }; 7 | 8 | use futures::{FutureExt, Stream, StreamExt}; 9 | use pyo3::{ 10 | exceptions::{PyStopAsyncIteration, PyStopIteration}, 11 | intern, 12 | prelude::*, 13 | }; 14 | 15 | use crate::{coroutine, utils}; 16 | 17 | utils::module!(Asyncio, "asyncio", Future); 18 | 19 | fn asyncio_future(py: Python) -> PyResult { 20 | Asyncio::get(py)?.Future.call0(py) 21 | } 22 | 23 | pub(crate) struct Waker { 24 | call_soon_threadsafe: PyObject, 25 | future: PyObject, 26 | } 27 | 28 | impl coroutine::CoroutineWaker for Waker { 29 | fn new(py: Python) -> PyResult { 30 | let future = asyncio_future(py)?; 31 | let call_soon_threadsafe = future 32 | .call_method0(py, intern!(py, "get_loop"))? 33 | .getattr(py, intern!(py, "call_soon_threadsafe"))?; 34 | Ok(Waker { 35 | call_soon_threadsafe, 36 | future, 37 | }) 38 | } 39 | 40 | fn yield_(&self, py: Python) -> PyResult { 41 | self.future 42 | .call_method0(py, intern!(py, "__await__"))? 43 | .call_method0(py, intern!(py, "__next__")) 44 | } 45 | 46 | fn wake(&self, py: Python) { 47 | self.future 48 | .call_method1(py, intern!(py, "set_result"), (py.None(),)) 49 | .expect("error while calling EventLoop.call_soon_threadsafe"); 50 | } 51 | 52 | fn wake_threadsafe(&self, py: Python) { 53 | let set_result = self 54 | .future 55 | .getattr(py, intern!(py, "set_result")) 56 | .expect("error while calling Future.set_result"); 57 | self.call_soon_threadsafe 58 | .call1(py, (set_result, py.None())) 59 | .expect("error while calling EventLoop.call_soon_threadsafe"); 60 | } 61 | 62 | fn update(&mut self, py: Python) -> PyResult<()> { 63 | self.future = Asyncio::get(py)?.Future.call0(py)?; 64 | Ok(()) 65 | } 66 | 67 | fn raise(&self, py: Python) -> PyResult<()> { 68 | self.future.call_method0(py, intern!(py, "result"))?; 69 | Ok(()) 70 | } 71 | } 72 | 73 | utils::generate!(Waker); 74 | 75 | /// [`Future`] wrapper for a Python awaitable (in `asyncio` context). 76 | /// 77 | /// The future should be polled in the thread where the event loop is running. 78 | pub struct AwaitableWrapper { 79 | future_iter: PyObject, 80 | future: Option, 81 | } 82 | 83 | impl AwaitableWrapper { 84 | /// Wrap a Python awaitable. 85 | pub fn new(awaitable: &PyAny) -> PyResult { 86 | Ok(Self { 87 | future_iter: awaitable 88 | .call_method0(intern!(awaitable.py(), "__await__"))? 89 | .extract()?, 90 | future: None, 91 | }) 92 | } 93 | 94 | /// GIL-bound [`Future`] reference. 95 | pub fn as_mut<'a>( 96 | &'a mut self, 97 | py: Python<'a>, 98 | ) -> impl Future> + Unpin + 'a { 99 | utils::WithGil { inner: self, py } 100 | } 101 | } 102 | 103 | impl<'a> Future for utils::WithGil<'_, &'a mut AwaitableWrapper> { 104 | type Output = PyResult; 105 | 106 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 107 | if let Some(fut) = self.inner.future.as_ref() { 108 | fut.call_method0(self.py, intern!(self.py, "result"))?; 109 | } 110 | match self 111 | .inner 112 | .future_iter 113 | .call_method0(self.py, intern!(self.py, "__next__")) 114 | { 115 | Ok(future) => { 116 | let callback = utils::wake_callback(self.py, cx.waker().clone())?; 117 | future.call_method1(self.py, intern!(self.py, "add_done_callback"), (callback,))?; 118 | self.inner.future = Some(future); 119 | Poll::Pending 120 | } 121 | Err(err) if err.is_instance_of::(self.py) => Poll::Ready(Ok(err 122 | .value(self.py) 123 | .getattr(intern!(self.py, "value"))? 124 | .into())), 125 | Err(err) => Poll::Ready(Err(err)), 126 | } 127 | } 128 | } 129 | 130 | impl Future for AwaitableWrapper { 131 | type Output = PyResult; 132 | 133 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 134 | Python::with_gil(|gil| Pin::into_inner(self).as_mut(gil).poll_unpin(cx)) 135 | } 136 | } 137 | 138 | /// [`Future`] wrapper for Python future. 139 | /// 140 | /// Because its duck-typed, it can work either with [`asyncio.Future`](https://docs.python.org/3/library/asyncio-future.html#asyncio.Future) or [`concurrent.futures.Future`](https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.Future). 141 | #[derive(Debug)] 142 | pub struct FutureWrapper { 143 | future: PyObject, 144 | cancel_on_drop: Option, 145 | } 146 | 147 | /// Cancel-on-drop error handling policy (see [`FutureWrapper::new`]). 148 | #[derive(Debug, Copy, Clone)] 149 | pub enum CancelOnDrop { 150 | IgnoreError, 151 | PanicOnError, 152 | } 153 | 154 | impl FutureWrapper { 155 | /// Wrap a Python future. 156 | /// 157 | /// If `cancel_on_drop` is not `None`, the Python future will be cancelled, and error may be 158 | /// handled following the provided policy. 159 | pub fn new(future: impl Into, cancel_on_drop: Option) -> Self { 160 | Self { 161 | future: future.into(), 162 | cancel_on_drop, 163 | } 164 | } 165 | 166 | /// GIL-bound [`Future`] reference. 167 | pub fn as_mut<'a>( 168 | &'a mut self, 169 | py: Python<'a>, 170 | ) -> impl Future> + Unpin + 'a { 171 | utils::WithGil { inner: self, py } 172 | } 173 | } 174 | 175 | impl<'a> Future for utils::WithGil<'_, &'a mut FutureWrapper> { 176 | type Output = PyResult; 177 | 178 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 179 | if self 180 | .inner 181 | .future 182 | .call_method0(self.py, intern!(self.py, "done"))? 183 | .is_true(self.py)? 184 | { 185 | self.inner.cancel_on_drop = None; 186 | return Poll::Ready( 187 | self.inner 188 | .future 189 | .call_method0(self.py, intern!(self.py, "result")), 190 | ); 191 | } 192 | let callback = utils::wake_callback(self.py, cx.waker().clone())?; 193 | self.inner.future.call_method1( 194 | self.py, 195 | intern!(self.py, "add_done_callback"), 196 | (callback,), 197 | )?; 198 | Poll::Pending 199 | } 200 | } 201 | 202 | impl Future for FutureWrapper { 203 | type Output = PyResult; 204 | 205 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 206 | Python::with_gil(|gil| Pin::into_inner(self).as_mut(gil).poll_unpin(cx)) 207 | } 208 | } 209 | 210 | impl Drop for FutureWrapper { 211 | fn drop(&mut self) { 212 | if let Some(cancel) = self.cancel_on_drop { 213 | let res = Python::with_gil(|gil| self.future.call_method0(gil, intern!(gil, "cancel"))); 214 | if let (Err(err), CancelOnDrop::PanicOnError) = (res, cancel) { 215 | panic!("Cancel error while dropping FutureWrapper: {err:?}"); 216 | } 217 | } 218 | } 219 | } 220 | 221 | /// [`Stream`] wrapper for a Python async generator (in `asyncio` context). 222 | /// 223 | /// The stream should be polled in the thread where the event loop is running. 224 | /// 225 | /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html 226 | pub struct AsyncGeneratorWrapper { 227 | async_generator: PyObject, 228 | next: Option, 229 | } 230 | 231 | impl AsyncGeneratorWrapper { 232 | /// Wrap a Python async generator. 233 | pub fn new(async_generator: &PyAny) -> Self { 234 | Self { 235 | async_generator: async_generator.into(), 236 | next: None, 237 | } 238 | } 239 | 240 | /// GIL-bound [`Stream`] reference. 241 | /// 242 | /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html 243 | pub fn as_mut<'a>( 244 | &'a mut self, 245 | py: Python<'a>, 246 | ) -> impl Stream> + Unpin + 'a { 247 | utils::WithGil { inner: self, py } 248 | } 249 | } 250 | 251 | impl<'a> Stream for utils::WithGil<'_, &'a mut AsyncGeneratorWrapper> { 252 | type Item = PyResult; 253 | 254 | fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 255 | if self.inner.next.is_none() { 256 | let next = self 257 | .inner 258 | .async_generator 259 | .as_ref(self.py) 260 | .call_method0(intern!(self.py, "__anext__"))?; 261 | self.inner.next = Some(AwaitableWrapper::new(next)?); 262 | } 263 | let res = ready!(self.inner.next.as_mut().unwrap().poll_unpin(cx)); 264 | self.inner.next = None; 265 | Poll::Ready(match res { 266 | Ok(obj) => Some(Ok(obj)), 267 | Err(err) if err.is_instance_of::(self.py) => None, 268 | Err(err) => Some(Err(err)), 269 | }) 270 | } 271 | } 272 | 273 | impl Stream for AsyncGeneratorWrapper { 274 | type Item = PyResult; 275 | 276 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 277 | Python::with_gil(|gil| Pin::into_inner(self).as_mut(gil).poll_next_unpin(cx)) 278 | } 279 | } 280 | -------------------------------------------------------------------------------- /src/coroutine.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | pin::Pin, 3 | sync::Arc, 4 | task::{Context, Poll}, 5 | }; 6 | 7 | use futures::task::ArcWake; 8 | use pyo3::{exceptions::PyRuntimeError, iter::IterNextOutput, prelude::*}; 9 | 10 | use crate::{ 11 | utils::{current_thread_id, ThreadId}, 12 | PyFuture, ThrowCallback, 13 | }; 14 | 15 | pub(crate) trait CoroutineWaker: Sized { 16 | fn new(py: Python) -> PyResult; 17 | fn yield_(&self, py: Python) -> PyResult; 18 | fn wake(&self, py: Python); 19 | fn wake_threadsafe(&self, py: Python); 20 | fn update(&mut self, _py: Python) -> PyResult<()> { 21 | Ok(()) 22 | } 23 | fn raise(&self, _py: Python) -> PyResult<()> { 24 | Ok(()) 25 | } 26 | } 27 | 28 | pub(crate) struct Waker { 29 | inner: W, 30 | thread_id: ThreadId, 31 | } 32 | 33 | impl ArcWake for Waker { 34 | fn wake_by_ref(arc_self: &Arc) { 35 | if current_thread_id() == arc_self.thread_id { 36 | Python::with_gil(|gil| CoroutineWaker::wake(&arc_self.inner, gil)) 37 | } else { 38 | Python::with_gil(|gil| CoroutineWaker::wake_threadsafe(&arc_self.inner, gil)) 39 | } 40 | } 41 | } 42 | 43 | pub(crate) struct Coroutine { 44 | future: Option>>, 45 | throw: Option, 46 | waker: Option>>, 47 | } 48 | 49 | impl Coroutine { 50 | pub(crate) fn new(future: Pin>, throw: Option) -> Self { 51 | Self { 52 | future: Some(future), 53 | throw, 54 | waker: None, 55 | } 56 | } 57 | 58 | pub(crate) fn close(&mut self, py: Python) -> PyResult<()> { 59 | if let Some(mut future_rs) = self.future.take() { 60 | if let Some(ref mut throw) = self.throw { 61 | throw(py, None); 62 | let waker = futures::task::noop_waker(); 63 | let res = future_rs 64 | .as_mut() 65 | .poll_py(py, &mut Context::from_waker(&waker)); 66 | if let Poll::Ready(Err(err)) = res { 67 | return Err(err); 68 | } 69 | } 70 | } 71 | Ok(()) 72 | } 73 | } 74 | 75 | impl Coroutine { 76 | pub(crate) fn poll( 77 | &mut self, 78 | py: Python, 79 | exc: Option, 80 | ) -> PyResult> { 81 | let Some(ref mut future_rs) = self.future else { 82 | return Err(PyRuntimeError::new_err( 83 | "cannot reuse already awaited coroutine", 84 | )); 85 | }; 86 | let exc = exc.or_else(|| self.waker.as_ref().and_then(|w| w.inner.raise(py).err())); 87 | match (exc, &mut self.throw) { 88 | (Some(exc), Some(throw)) => throw(py, Some(exc)), 89 | (Some(exc), _) => { 90 | self.future.take(); 91 | return Err(exc); 92 | } 93 | _ => {} 94 | } 95 | if let Some(waker) = self.waker.as_mut().and_then(Arc::get_mut) { 96 | waker.inner.update(py)?; 97 | } else { 98 | self.waker = Some(Arc::new(Waker { 99 | inner: W::new(py)?, 100 | thread_id: current_thread_id(), 101 | })); 102 | } 103 | let waker = futures::task::waker(self.waker.clone().unwrap()); 104 | let res = future_rs 105 | .as_mut() 106 | .poll_py(py, &mut Context::from_waker(&waker)); 107 | Ok(match res { 108 | Poll::Ready(res) => { 109 | self.future.take(); 110 | IterNextOutput::Return(res?) 111 | } 112 | Poll::Pending => IterNextOutput::Yield(self.waker.as_ref().unwrap().inner.yield_(py)?), 113 | }) 114 | } 115 | } 116 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! PyO3 bindings to various Python asynchronous frameworks. 2 | use std::{ 3 | future::Future, 4 | pin::Pin, 5 | task::{Context, Poll}, 6 | }; 7 | 8 | use futures::Stream; 9 | use pyo3::prelude::*; 10 | 11 | #[cfg(feature = "allow-threads")] 12 | mod allow_threads; 13 | mod async_generator; 14 | pub mod asyncio; 15 | mod coroutine; 16 | pub mod sniffio; 17 | pub mod trio; 18 | mod utils; 19 | 20 | #[cfg(feature = "allow-threads")] 21 | pub use allow_threads::{AllowThreads, AllowThreadsExt}; 22 | #[cfg(feature = "macros")] 23 | pub use pyo3_async_macros::{pyfunction, pymethods}; 24 | 25 | /// GIL-bound [`Future`]. 26 | /// 27 | /// Provided with a blanket implementation for [`Future`]. GIL is maintained during polling 28 | /// operation. To release the GIL, see [`AllowThreads`]. 29 | pub trait PyFuture: Send { 30 | /// GIL-bound [`Future::poll`]. 31 | fn poll_py(self: Pin<&mut Self>, py: Python, cx: &mut Context) -> Poll>; 32 | } 33 | 34 | impl PyFuture for F 35 | where 36 | F: Future> + Send, 37 | T: IntoPy + Send, 38 | E: Send, 39 | PyErr: From, 40 | { 41 | fn poll_py(self: Pin<&mut Self>, py: Python, cx: &mut Context) -> Poll> { 42 | let poll = self.poll(cx); 43 | poll.map_ok(|ok| ok.into_py(py)).map_err(PyErr::from) 44 | } 45 | } 46 | 47 | /// GIL-bound [`Stream`]. 48 | /// 49 | /// Provided with a blanket implementation for [`Stream`]. GIL is maintained during polling 50 | /// operation. To release the GIL, see [`AllowThreads`]. 51 | /// 52 | /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html 53 | pub trait PyStream: Send { 54 | /// GIL-bound [`Stream::poll_next`]. 55 | fn poll_next_py( 56 | self: Pin<&mut Self>, 57 | py: Python, 58 | cx: &mut Context, 59 | ) -> Poll>>; 60 | } 61 | 62 | impl PyStream for S 63 | where 64 | S: Stream> + Send, 65 | T: IntoPy + Send, 66 | E: Send, 67 | PyErr: From, 68 | { 69 | fn poll_next_py( 70 | self: Pin<&mut Self>, 71 | py: Python, 72 | cx: &mut Context, 73 | ) -> Poll>> { 74 | let poll = self.poll_next(cx); 75 | poll.map_ok(|ok| ok.into_py(py)).map_err(PyErr::from) 76 | } 77 | } 78 | 79 | /// Callback for Python coroutine `throw` method (see [`asyncio::Coroutine::new`]) and 80 | /// async generator `athrow` method (see [`asyncio::AsyncGenerator::new`]). 81 | pub type ThrowCallback = Box) + Send>; 82 | -------------------------------------------------------------------------------- /src/sniffio.rs: -------------------------------------------------------------------------------- 1 | //! `asyncio`/`trio` compatible coroutine and async generator implementation, lazily specialized 2 | //! using `sniffio`. 3 | use pyo3::{exceptions::PyRuntimeError, prelude::*}; 4 | 5 | use crate::{asyncio, coroutine, trio, utils}; 6 | 7 | utils::module!(Sniffio, "sniffio", current_async_library); 8 | 9 | enum Waker { 10 | Asyncio(asyncio::Waker), 11 | Trio(trio::Waker), 12 | } 13 | 14 | impl coroutine::CoroutineWaker for Waker { 15 | fn new(py: Python) -> PyResult { 16 | let sniffed = Sniffio::get(py)?.current_async_library.call0(py)?; 17 | match sniffed.extract(py)? { 18 | "asyncio" => Ok(Self::Asyncio(asyncio::Waker::new(py)?)), 19 | "trio" => Ok(Self::Trio(trio::Waker::new(py)?)), 20 | rt => Err(PyRuntimeError::new_err(format!("unsupported runtime {rt}"))), 21 | } 22 | } 23 | 24 | fn yield_(&self, py: Python) -> PyResult { 25 | match self { 26 | Self::Asyncio(w) => w.yield_(py), 27 | Self::Trio(w) => w.yield_(py), 28 | } 29 | } 30 | 31 | fn wake(&self, py: Python) { 32 | match self { 33 | Self::Asyncio(w) => w.wake(py), 34 | Self::Trio(w) => w.wake(py), 35 | } 36 | } 37 | 38 | fn wake_threadsafe(&self, py: Python) { 39 | match self { 40 | Self::Asyncio(w) => w.wake_threadsafe(py), 41 | Self::Trio(w) => w.wake_threadsafe(py), 42 | } 43 | } 44 | 45 | fn update(&mut self, py: Python) -> PyResult<()> { 46 | match self { 47 | Self::Asyncio(w) => w.update(py), 48 | Self::Trio(w) => w.update(py), 49 | } 50 | } 51 | 52 | fn raise(&self, py: Python) -> PyResult<()> { 53 | match self { 54 | Self::Asyncio(w) => w.raise(py), 55 | Self::Trio(w) => w.raise(py), 56 | } 57 | } 58 | } 59 | 60 | utils::generate!(Waker); 61 | -------------------------------------------------------------------------------- /src/trio.rs: -------------------------------------------------------------------------------- 1 | //! `trio` compatible coroutine and async generator implementation. 2 | use pyo3::{intern, prelude::*}; 3 | 4 | use crate::{coroutine, utils}; 5 | 6 | utils::module!( 7 | Trio, 8 | "trio.lowlevel", 9 | Abort, 10 | current_task, 11 | current_trio_token, 12 | reschedule, 13 | wait_task_rescheduled 14 | ); 15 | 16 | pub(crate) struct Waker { 17 | task: PyObject, 18 | token: PyObject, 19 | } 20 | 21 | impl coroutine::CoroutineWaker for Waker { 22 | fn new(py: Python) -> PyResult { 23 | let trio = Trio::get(py)?; 24 | Ok(Waker { 25 | task: trio.current_task.call0(py)?, 26 | token: trio.current_trio_token.call0(py)?, 27 | }) 28 | } 29 | 30 | fn yield_(&self, py: Python) -> PyResult { 31 | Trio::get(py)? 32 | .wait_task_rescheduled 33 | .call1(py, (wrap_pyfunction!(abort_func, py)?,))? 34 | .call_method0(py, intern!(py, "__await__"))? 35 | .call_method0(py, intern!(py, "__next__")) 36 | } 37 | 38 | fn wake(&self, py: Python) { 39 | let reschedule = &Trio::get(py).unwrap().reschedule; 40 | reschedule 41 | .call1(py, (&self.task,)) 42 | .expect("unexpected error while calling trio.lowlevel.reschedule"); 43 | } 44 | 45 | fn wake_threadsafe(&self, py: Python) { 46 | let reschedule = &Trio::get(py).unwrap().reschedule; 47 | self.token 48 | .call_method1(py, intern!(py, "run_sync_soon"), (reschedule, &self.task)) 49 | .expect("unexpected error while scheduling TrioToken.run_sync_soon"); 50 | } 51 | } 52 | 53 | #[pyfunction] 54 | fn abort_func(py: Python, _arg: PyObject) -> PyResult { 55 | Trio::get(py)?.Abort.getattr(py, intern!(py, "SUCCEEDED")) 56 | } 57 | 58 | utils::generate!(Waker); 59 | -------------------------------------------------------------------------------- /src/utils.rs: -------------------------------------------------------------------------------- 1 | use std::sync::atomic::{AtomicUsize, Ordering}; 2 | 3 | use pyo3::{exceptions::PyStopIteration, prelude::*, pyclass::IterNextOutput, types::PyCFunction}; 4 | 5 | // Don't use `std::thread::current` because of unnecessary Arc clone + drop. 6 | pub(crate) type ThreadId = usize; 7 | pub(crate) fn current_thread_id() -> ThreadId { 8 | static THREAD_COUNTER: AtomicUsize = AtomicUsize::new(0); 9 | thread_local! { 10 | pub(crate) static THREAD_ID: ThreadId = THREAD_COUNTER.fetch_add(1, Ordering::Relaxed); 11 | } 12 | THREAD_ID.with(|id| *id) 13 | } 14 | 15 | pub(crate) struct WithGil<'py, T> { 16 | pub(crate) inner: T, 17 | pub(crate) py: Python<'py>, 18 | } 19 | 20 | pub(crate) fn wake_callback(py: Python, waker: std::task::Waker) -> PyResult<&PyAny> { 21 | let func = PyCFunction::new_closure(py, None, None, move |_, _| waker.wake_by_ref())?; 22 | Ok(func) 23 | } 24 | 25 | macro_rules! module { 26 | ($name:ident ,$path:literal, $($field:ident),* $(,)?) => { 27 | #[allow(non_upper_case_globals)] 28 | static $name: ::pyo3::sync::GILOnceCell<$name> = ::pyo3::sync::GILOnceCell::new(); 29 | 30 | #[allow(non_snake_case)] 31 | struct $name { 32 | $($field: PyObject),* 33 | } 34 | 35 | impl $name { 36 | fn get(py: Python) -> PyResult<&Self> { 37 | $name.get_or_try_init(py, || { 38 | let module = py.import($path)?; 39 | Ok(Self { 40 | $($field: module.getattr(stringify!($field))?.into(),)* 41 | }) 42 | }) 43 | } 44 | } 45 | }; 46 | } 47 | 48 | pub(crate) use module; 49 | 50 | pub(crate) fn poll_result(result: IterNextOutput) -> PyResult { 51 | match result { 52 | IterNextOutput::Yield(ob) => Ok(ob), 53 | IterNextOutput::Return(ob) => Err(PyStopIteration::new_err(ob)), 54 | } 55 | } 56 | 57 | macro_rules! generate { 58 | ($waker:ty) => { 59 | /// Python coroutine wrapping a [`PyFuture`](crate::PyFuture). 60 | #[pyclass] 61 | pub struct Coroutine($crate::coroutine::Coroutine<$waker>); 62 | 63 | impl Coroutine { 64 | /// Wrap a boxed future in to a Python coroutine. 65 | /// 66 | /// If `throw` callback is provided: 67 | /// - coroutine `throw` method will call it with the passed exception before polling; 68 | /// - coroutine `close` method will call it with `None` before polling and dropping 69 | /// the future. 70 | /// If `throw` callback is not provided, the future will dropped without additional 71 | /// poll. 72 | pub fn new( 73 | future: ::std::pin::Pin>, 74 | throw: Option<$crate::ThrowCallback>, 75 | ) -> Self { 76 | Self($crate::coroutine::Coroutine::new(future, throw)) 77 | } 78 | 79 | /// Wrap a generic future into a Python coroutine. 80 | pub fn from_future(future: impl $crate::PyFuture + 'static) -> Self { 81 | Self::new(Box::pin(future), None) 82 | } 83 | } 84 | 85 | #[pymethods] 86 | impl Coroutine { 87 | fn send(&mut self, py: Python, _value: &PyAny) -> PyResult { 88 | $crate::utils::poll_result(self.0.poll(py, None)?) 89 | } 90 | 91 | fn throw(&mut self, py: Python, exc: &PyAny) -> PyResult { 92 | $crate::utils::poll_result(self.0.poll(py, Some(PyErr::from_value(exc)))?) 93 | } 94 | 95 | fn close(&mut self, py: Python) -> PyResult<()> { 96 | self.0.close(py) 97 | } 98 | 99 | fn __await__(self_: &PyCell) -> PyResult<&PyAny> { 100 | Ok(self_) 101 | } 102 | 103 | fn __iter__(self_: &PyCell) -> PyResult<&PyAny> { 104 | Ok(self_) 105 | } 106 | 107 | fn __next__( 108 | &mut self, 109 | py: Python, 110 | ) -> PyResult<::pyo3::pyclass::IterNextOutput> { 111 | self.0.poll(py, None) 112 | } 113 | } 114 | 115 | impl $crate::async_generator::CoroutineFactory for Coroutine { 116 | type Coroutine = Self; 117 | fn coroutine(future: impl $crate::PyFuture + 'static) -> Self::Coroutine { 118 | Self::from_future(future) 119 | } 120 | } 121 | 122 | /// Python async generator wrapping a [`PyStream`](crate::PyStream). 123 | #[pyclass] 124 | pub struct AsyncGenerator($crate::async_generator::AsyncGenerator); 125 | 126 | impl AsyncGenerator { 127 | /// Wrap a boxed stream in to a Python async generator. 128 | /// 129 | /// If `throw` callback is provided: 130 | /// - async generator `athrow` method will call it with the passed exception 131 | /// before polling; 132 | /// - async generator `aclose` method will call it with `None` before polling and 133 | /// dropping the stream. 134 | /// If `throw` callback is not provided, the stream will dropped without additional 135 | /// poll. 136 | pub fn new( 137 | stream: ::std::pin::Pin>, 138 | throw: Option<$crate::ThrowCallback>, 139 | ) -> Self { 140 | Self($crate::async_generator::AsyncGenerator::new(stream, throw)) 141 | } 142 | 143 | /// Wrap a generic stream. 144 | pub fn from_stream(stream: impl $crate::PyStream + 'static) -> Self { 145 | Self::new(Box::pin(stream), None) 146 | } 147 | } 148 | 149 | #[pymethods] 150 | impl AsyncGenerator { 151 | fn asend(&mut self, py: Python, _value: &PyAny) -> PyResult { 152 | self.0.next(py) 153 | } 154 | 155 | fn athrow(&mut self, py: Python, exc: &PyAny) -> PyResult { 156 | self.0.throw(py, PyErr::from_value(exc)) 157 | } 158 | 159 | fn aclose(&mut self, py: Python) -> PyResult { 160 | self.0.close(py) 161 | } 162 | 163 | fn __aiter__(self_: &PyCell) -> PyResult<&PyAny> { 164 | Ok(self_) 165 | } 166 | 167 | // `Option` because https://github.com/PyO3/pyo3/issues/3190 168 | fn __anext__(&mut self, py: Python) -> PyResult> { 169 | self.0.next(py).map(Some) 170 | } 171 | } 172 | }; 173 | } 174 | pub(crate) use generate; 175 | --------------------------------------------------------------------------------