python_ast/ast/tree/
list.rs1use proc_macro2::TokenStream;
2use pyo3::{Bound, FromPyObject, PyAny, types::PyAnyMethods};
3use quote::quote;
4
5use crate::{dump, CodeGen, CodeGenContext, PythonOptions, SymbolTableScopes};
6
7pub 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 let expr: ExprType = elt.extract()?;
57
58 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 rust_str.contains("sys :: argv") {
67 elements.push(quote! { });
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 has_starred && elements.iter().any(|e| e.to_string().contains("STARRED_ARGV")) {
82 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 continue;
89 } else {
90 final_elements.push(element);
91 }
92 }
93
94 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!(vec![#(#elements),*]))
105 }
106 }
107}
108
109#[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}