python_ast/ast/tree/
list.rs

1use proc_macro2::TokenStream;
2use pyo3::{Bound, FromPyObject, PyAny, types::PyAnyMethods};
3use quote::quote;
4
5use crate::{dump, CodeGen, CodeGenContext, PythonOptions, SymbolTableScopes};
6
7// There are two concepts of List in the same place here. There's the "List" type that represents a node from the Python AST,
8// as received by the Rust AST converter, and there's the List representation of the Python List type. For the sake of
9// consistency, we're using the same type as we use to model
10pub type ListContents = crate::pytypes::List<dyn CodeGen>;
11
12#[derive(Clone, Default)]
13pub struct List<'a> {
14    pub elts: Vec<Bound<'a, PyAny>>,
15    pub ctx: Option<String>,
16}
17
18impl std::fmt::Debug for List<'_> {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        f.debug_struct("List")
21            .field("elts", &format!("Vec<Bound<PyAny>> (len: {})", self.elts.len()))
22            .field("ctx", &self.ctx)
23            .finish()
24    }
25}
26
27impl<'a> FromPyObject<'a> for List<'a> {
28    fn extract_bound(ob: &Bound<'a, PyAny>) -> pyo3::PyResult<Self> {
29        let elts: Vec<Bound<'a, PyAny>> = ob.getattr("elts")?.extract()?;
30        let ctx: Option<String> = ob.getattr("ctx").ok().and_then(|v| v.extract().ok());
31        Ok(List { elts, ctx })
32    }
33}
34
35impl<'a> CodeGen for List<'a> {
36    type Context = CodeGenContext;
37    type Options = PythonOptions;
38    type SymbolTable = SymbolTableScopes;
39
40    fn to_rust(
41        self,
42        ctx: Self::Context,
43        options: Self::Options,
44        symbols: Self::SymbolTable,
45    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
46        use crate::ExprType;
47        
48        let mut elements = Vec::new();
49        let mut has_starred = false;
50        
51        log::debug!("================Processing list with {} elements", self.elts.len());
52        for elt in self.elts {
53            log::debug!("elt: {}", dump(&elt, None)?);
54            
55            // Extract the element as ExprType and convert to Rust
56            let expr: ExprType = elt.extract()?;
57            
58            // Check if this is a starred expression
59            match &expr {
60                ExprType::Starred(_) => {
61                    has_starred = true;
62                    let rust_code = expr.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
63                    let rust_str = rust_code.to_string();
64                    
65                    // If it's a starred sys::argv, we need special handling
66                    if rust_str.contains("sys :: argv") {
67                        // Instead of adding individual elements, we'll build the vector differently
68                        elements.push(quote! { /* STARRED_ARGV */ });
69                    } else {
70                        elements.push(rust_code);
71                    }
72                }
73                _ => {
74                    let rust_code = expr.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
75                    elements.push(rust_code);
76                }
77            }
78        }
79        
80        // If we have starred expressions, especially sys::argv, handle it specially
81        if has_starred && elements.iter().any(|e| e.to_string().contains("STARRED_ARGV")) {
82            // Create a special vector construction that handles sys::argv unpacking
83            let mut final_elements = Vec::new();
84            for element in elements {
85                let elem_str = element.to_string();
86                if elem_str.contains("STARRED_ARGV") {
87                    // Skip the placeholder and add the argv unpacking
88                    continue;
89                } else {
90                    final_elements.push(element);
91                }
92            }
93            
94            // Build the vector with sys::argv unpacking
95            Ok(quote! {
96                {
97                    let mut vec = Vec::new();
98                    #(vec.push(#final_elements);)*
99                    vec.extend((*sys::argv).iter().cloned());
100                    vec
101                }
102            })
103        } else {
104            Ok(quote!(/* LIST_GENERATED */ vec![#(#elements),*]))
105        }
106    }
107}
108
109// It's fairly easy to break the automatic parsing of parameter structs, so we need to have fairly sophisticated
110// test coverage for the various types of
111#[cfg(test)]
112mod tests {
113    use crate::ExprType;
114    use crate::StatementType;
115    use std::panic;
116    use test_log::test;
117
118    #[test]
119    fn parse_list() {
120        let module = crate::parse("[1, 2, 3]", "nothing.py").unwrap();
121        let statement = module.raw.body[0].statement.clone();
122        match statement {
123            StatementType::Expr(e) => match e.value {
124                ExprType::List(list) => {
125                    log::debug!("{:#?}", list);
126                    assert_eq!(list.len(), 3);
127                }
128                _ => panic!("Could not find inner expression"),
129            },
130            _ => panic!("Could not find outer expression."),
131        }
132    }
133}