Skip to main content

python_ast/ast/tree/
await_kw.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, PyAny, PyResult, FromPyObject, prelude::PyAnyMethods};
3use quote::quote;
4
5use crate::{extraction_failure, CodeGen, CodeGenContext, ExprType, PythonOptions, SymbolTableScopes};
6
7use serde::{Deserialize, Serialize};
8
9#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
10//#[pyo3(transparent)]
11pub struct Await {
12    pub value: Box<ExprType>,
13}
14
15impl<'a, 'py> FromPyObject<'a, 'py> for Await {
16    type Error = pyo3::PyErr;
17    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
18        let value = ob.getattr("value").map_err(|e| extraction_failure("Await.value", &ob, e))?;
19        Ok(Await {
20            value: Box::new(value.extract().map_err(|e| extraction_failure("Await.value", &ob, e))?),
21        })
22    }
23}
24
25impl<'a> CodeGen for Await {
26    type Context = CodeGenContext;
27    type Options = PythonOptions;
28    type SymbolTable = SymbolTableScopes;
29
30    fn to_rust(
31        self,
32        _ctx: Self::Context,
33        _options: Self::Options,
34        _symbols: Self::SymbolTable,
35    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
36        let value = self.value.to_rust(_ctx, _options, _symbols)?;
37        // A call to a user async function lowers to `f(...)?` — but the `?`
38        // must apply to the awaited Result, not the future: reorder to
39        // `f(...).await?`.
40        let rendered = value.to_string();
41        if let Some(inner) = rendered.trim_end().strip_suffix('?') {
42            let inner: TokenStream = inner
43                .parse()
44                .map_err(|e| format!("re-parsing awaited call: {:?}", e))?;
45            return Ok(quote!(#inner.await?));
46        }
47        Ok(quote!(#value.await))
48    }
49}