python_ast/ast/tree/
async_with.rs1use proc_macro2::TokenStream;
2use pyo3::{Bound, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{
7 CodeGen, CodeGenContext, ExprType, Node, PythonOptions, Statement, SymbolTableScopes,
8 extract_list,
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
13pub struct AsyncWith {
14 pub items: Vec<WithItem>,
16 pub body: Vec<Statement>,
18 pub lineno: Option<usize>,
20 pub col_offset: Option<usize>,
21 pub end_lineno: Option<usize>,
22 pub end_col_offset: Option<usize>,
23}
24
25#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
27pub struct WithItem {
28 pub context_expr: ExprType,
30 pub optional_vars: Option<ExprType>,
32}
33
34impl<'a> FromPyObject<'a> for AsyncWith {
35 fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
36 let items: Vec<WithItem> = extract_list(ob, "items", "async with items")?;
38
39 let body: Vec<Statement> = extract_list(ob, "body", "async with body")?;
41
42 Ok(AsyncWith {
43 items,
44 body,
45 lineno: ob.lineno(),
46 col_offset: ob.col_offset(),
47 end_lineno: ob.end_lineno(),
48 end_col_offset: ob.end_col_offset(),
49 })
50 }
51}
52
53impl<'a> FromPyObject<'a> for WithItem {
54 fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
55 let context_expr: ExprType = ob.getattr("context_expr")?.extract()?;
57
58 let optional_vars: Option<ExprType> = if let Ok(vars_attr) = ob.getattr("optional_vars") {
60 if vars_attr.is_none() {
61 None
62 } else {
63 Some(vars_attr.extract()?)
64 }
65 } else {
66 None
67 };
68
69 Ok(WithItem {
70 context_expr,
71 optional_vars,
72 })
73 }
74}
75
76impl Node for AsyncWith {
77 fn lineno(&self) -> Option<usize> { self.lineno }
78 fn col_offset(&self) -> Option<usize> { self.col_offset }
79 fn end_lineno(&self) -> Option<usize> { self.end_lineno }
80 fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
81}
82
83impl CodeGen for AsyncWith {
84 type Context = CodeGenContext;
85 type Options = PythonOptions;
86 type SymbolTable = SymbolTableScopes;
87
88 fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
89 let symbols = self.items.into_iter().fold(symbols, |acc, item| {
91 let acc = item.context_expr.find_symbols(acc);
92 if let Some(vars) = item.optional_vars {
93 vars.find_symbols(acc)
94 } else {
95 acc
96 }
97 });
98 self.body.into_iter().fold(symbols, |acc, stmt| stmt.find_symbols(acc))
99 }
100
101 fn to_rust(
102 self,
103 ctx: Self::Context,
104 options: Self::Options,
105 symbols: Self::SymbolTable,
106 ) -> Result<TokenStream, Box<dyn std::error::Error>> {
107 let body_tokens: Result<Vec<TokenStream>, Box<dyn std::error::Error>> = self.body.into_iter()
109 .map(|stmt| stmt.to_rust(ctx.clone(), options.clone(), symbols.clone()))
110 .collect();
111 let body_tokens = body_tokens?;
112
113 Ok(quote! {
116 {
117 #(#body_tokens)*
120 }
121 })
122 }
123}
124
125#[cfg(test)]
126mod tests {
127 }