python_ast/ast/tree/
arguments.rs

1//! The module defines Python-syntax arguments and maps them into Rust-syntax versions.
2use proc_macro2::TokenStream;
3use pyo3::{Borrowed, Bound, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods};
4use quote::quote;
5use serde::{Deserialize, Serialize};
6
7use crate::{
8    CodeGen, CodeGenContext, ExprType, Node, PythonOptions, SymbolTableScopes,
9};
10
11/// A complete argument representation that can hold any Python expression.
12/// This replaces the limited Arg enum to support all argument types.
13#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
14pub struct Argument {
15    /// The argument expression (can be any valid Python expression)
16    pub value: ExprType,
17    /// Position information
18    pub lineno: Option<usize>,
19    pub col_offset: Option<usize>,
20    pub end_lineno: Option<usize>,
21    pub end_col_offset: Option<usize>,
22}
23
24/// An argument value that can be any expression.
25/// This replaces the old limited Arg enum.
26pub type Arg = ExprType;
27
28/// A function parameter definition with optional type annotation and default value.
29#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
30pub struct Parameter {
31    /// Parameter name
32    pub arg: String,
33    /// Optional type annotation
34    pub annotation: Option<Box<ExprType>>,
35    /// Optional type comment (deprecated Python feature)
36    pub type_comment: Option<String>,
37    /// Position information
38    pub lineno: Option<usize>,
39    pub col_offset: Option<usize>,
40    pub end_lineno: Option<usize>,
41    pub end_col_offset: Option<usize>,
42}
43
44/// Comprehensive function arguments structure supporting all Python argument types.
45#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
46pub struct Arguments {
47    /// Positional-only parameters (before / in Python 3.8+)
48    pub posonlyargs: Vec<Parameter>,
49    /// Regular positional parameters
50    pub args: Vec<Parameter>,
51    /// Variable positional parameter (*args)
52    pub vararg: Option<Parameter>,
53    /// Keyword-only parameters (after * or *args)
54    pub kwonlyargs: Vec<Parameter>,
55    /// Default values for keyword-only parameters (None = required)
56    pub kw_defaults: Vec<Option<Box<ExprType>>>,
57    /// Variable keyword parameter (**kwargs)
58    pub kwarg: Option<Parameter>,
59    /// Default values for regular positional parameters
60    pub defaults: Vec<Box<ExprType>>,
61}
62
63
64/// Function call arguments supporting all Python call patterns.
65#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
66pub struct CallArguments {
67    /// Positional arguments
68    pub args: Vec<ExprType>,
69    /// Keyword arguments
70    pub keywords: Vec<crate::Keyword>,
71}
72
73// Implementation for new Argument struct
74impl<'a, 'py> FromPyObject<'a, 'py> for Argument {
75    type Error = pyo3::PyErr;
76    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
77        // Extract the expression value
78        let value: ExprType = ob.extract()?;
79        
80        Ok(Self {
81            value,
82            lineno: ob.lineno(),
83            col_offset: ob.col_offset(),
84            end_lineno: ob.end_lineno(),
85            end_col_offset: ob.end_col_offset(),
86        })
87    }
88}
89
90impl CodeGen for Argument {
91    type Context = CodeGenContext;
92    type Options = PythonOptions;
93    type SymbolTable = SymbolTableScopes;
94
95    fn to_rust(
96        self,
97        ctx: Self::Context,
98        options: Self::Options,
99        symbols: Self::SymbolTable,
100    ) -> std::result::Result<TokenStream, Box<dyn std::error::Error>> {
101        self.value.to_rust(ctx, options, symbols)
102    }
103}
104
105// Implementation for Parameter struct
106impl<'a, 'py> FromPyObject<'a, 'py> for Parameter {
107    type Error = pyo3::PyErr;
108    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
109        let arg: String = ob.getattr("arg")?.extract()?;
110        
111        // Extract optional annotation
112        let annotation = if let Ok(ann) = ob.getattr("annotation") {
113            if ann.is_none() {
114                None
115            } else {
116                Some(Box::new(ann.extract()?))
117            }
118        } else {
119            None
120        };
121        
122        // Extract optional type comment
123        let type_comment = if let Ok(tc) = ob.getattr("type_comment") {
124            if tc.is_none() {
125                None
126            } else {
127                Some(tc.extract()?)
128            }
129        } else {
130            None
131        };
132        
133        Ok(Self {
134            arg,
135            annotation,
136            type_comment,
137            lineno: ob.lineno(),
138            col_offset: ob.col_offset(),
139            end_lineno: ob.end_lineno(),
140            end_col_offset: ob.end_col_offset(),
141        })
142    }
143}
144
145/// Whether an annotation means "optional": `Optional[T]` or a union with
146/// None (`T | None`). Optional-annotated names hold an Option, and stores
147/// into them wrap in Some.
148pub(crate) fn is_optional_annotation(ann: &ExprType) -> bool {
149    match ann {
150        ExprType::Subscript(sub) => {
151            matches!(sub.value.as_ref(), ExprType::Name(n) if n.id == "Optional")
152        }
153        ExprType::BinOp(op) if matches!(op.op, crate::BinOps::BitOr) => {
154            crate::is_none_expr(&op.left) || crate::is_none_expr(&op.right)
155        }
156        _ => false,
157    }
158}
159
160/// Map a Python type annotation to a Rust type, when the mapping is known.
161/// `int`/`float`/`str`/`bool`/`bytes` map to concrete Rust types, and
162/// `list[T]`/`dict[K, V]`/`set[T]` map to the corresponding std containers
163/// when their element annotations map too. `Optional[T]` / `T | None` map
164/// to `Option<T>`.
165pub fn python_annotation_to_rust_type(annotation: &ExprType) -> Option<TokenStream> {
166    match annotation {
167        // T | None (and None | T) is Option<T>.
168        ExprType::BinOp(op) if matches!(op.op, crate::BinOps::BitOr) => {
169            let inner = if crate::is_none_expr(&op.left) {
170                op.right.as_ref()
171            } else if crate::is_none_expr(&op.right) {
172                op.left.as_ref()
173            } else {
174                return None;
175            };
176            let inner = python_annotation_to_rust_type(inner)?;
177            return Some(quote!(Option<#inner>));
178        }
179        _ => {}
180    }
181    match annotation {
182        ExprType::Name(name) => match name.id.as_str() {
183            "int" => Some(quote!(i64)),
184            "float" => Some(quote!(f64)),
185            "str" => Some(quote!(String)),
186            "bool" => Some(quote!(bool)),
187            "bytes" => Some(quote!(Vec<u8>)),
188            _ => None,
189        },
190        // Subscripted generics over known element types: list[int] and
191        // friends map to the concrete Rust containers codegen produces for
192        // the corresponding literals.
193        ExprType::Subscript(sub) => {
194            let container = match sub.value.as_ref() {
195                ExprType::Name(n) => n.id.as_str(),
196                _ => return None,
197            };
198            match (&sub.kind, container) {
199                (crate::SubscriptKind::Index(elt), "Optional") => {
200                    let inner = python_annotation_to_rust_type(elt)?;
201                    Some(quote!(Option<#inner>))
202                }
203                (crate::SubscriptKind::Index(elt), "list") => {
204                    let inner = python_annotation_to_rust_type(elt)?;
205                    Some(quote!(Vec<#inner>))
206                }
207                (crate::SubscriptKind::Index(elt), "set" | "frozenset") => {
208                    let inner = python_annotation_to_rust_type(elt)?;
209                    Some(quote!(std::collections::HashSet<#inner>))
210                }
211                (crate::SubscriptKind::Index(kv), "dict") => {
212                    // dict[K, V] parses as a subscript with a tuple index.
213                    // PyDict is the insertion-ordered map dict literals
214                    // lower to.
215                    if let ExprType::Tuple(t) = kv.as_ref() {
216                        if let [k, v] = t.elts.as_slice() {
217                            let k = python_annotation_to_rust_type(k)?;
218                            let v = python_annotation_to_rust_type(v)?;
219                            return Some(quote!(PyDict<#k, #v>));
220                        }
221                    }
222                    None
223                }
224                _ => None,
225            }
226        }
227        _ => None,
228    }
229}
230
231impl CodeGen for Parameter {
232    type Context = CodeGenContext;
233    type Options = PythonOptions;
234    type SymbolTable = SymbolTableScopes;
235
236    fn to_rust(
237        self,
238        ctx: Self::Context,
239        options: Self::Options,
240        symbols: Self::SymbolTable,
241    ) -> std::result::Result<TokenStream, Box<dyn std::error::Error>> {
242
243        let param_name = crate::safe_ident(&self.arg);
244
245        // Generate type annotation if present
246        if let Some(annotation) = self.annotation {
247            // A str parameter accepts anything convertible to String, so
248            // call sites can pass &str literals as well as owned Strings;
249            // the function prologue converts it (`let s: String = s.into()`).
250            if matches!(&*annotation, ExprType::Name(n) if n.id == "str") {
251                return Ok(quote!(#param_name: impl Into<String>));
252            }
253            // Known Python types map to concrete Rust types; anything else
254            // falls back to rendering the annotation expression (e.g. a
255            // user-defined class name).
256            let rust_type = match python_annotation_to_rust_type(&annotation) {
257                Some(mapped) => mapped,
258                None => annotation.to_rust(ctx, options, symbols)?,
259            };
260            Ok(quote!(#param_name: #rust_type))
261        } else {
262            // Default to generic type for untyped parameters
263            Ok(quote!(#param_name: impl Into<PyObject>))
264        }
265    }
266}
267
268// Implementation for Arguments struct
269impl<'a, 'py> FromPyObject<'a, 'py> for Arguments {
270    type Error = pyo3::PyErr;
271    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
272        // Extract each field with proper error handling
273        let posonlyargs: Vec<Parameter> = ob.getattr("posonlyargs")?.extract().unwrap_or_default();
274        let args: Vec<Parameter> = ob.getattr("args")?.extract().unwrap_or_default();
275        
276        let vararg = if let Ok(va) = ob.getattr("vararg") {
277            if va.is_none() { None } else { Some(va.extract()?) }
278        } else { None };
279        
280        let kwonlyargs: Vec<Parameter> = ob.getattr("kwonlyargs")?.extract().unwrap_or_default();
281        
282        // Handle kw_defaults which can contain None values
283        let kw_defaults = if let Ok(kw_def) = ob.getattr("kw_defaults") {
284            let defaults_list: Vec<Bound<PyAny>> = kw_def.extract().unwrap_or_default();
285            let mut processed_defaults = Vec::new();
286            for default in defaults_list {
287                if default.is_none() {
288                    processed_defaults.push(None);
289                } else {
290                    processed_defaults.push(Some(Box::new(default.extract()?)));
291                }
292            }
293            processed_defaults
294        } else {
295            Vec::new()
296        };
297        
298        let kwarg = if let Ok(kw) = ob.getattr("kwarg") {
299            if kw.is_none() { None } else { Some(kw.extract()?) }
300        } else { None };
301        
302        let defaults_raw: Vec<ExprType> = ob.getattr("defaults")?.extract().unwrap_or_default();
303        let defaults = defaults_raw.into_iter().map(Box::new).collect();
304        
305        Ok(Self {
306            posonlyargs,
307            args,
308            vararg,
309            kwonlyargs,
310            kw_defaults,
311            kwarg,
312            defaults,
313        })
314    }
315}
316
317impl CodeGen for Arguments {
318    type Context = CodeGenContext;
319    type Options = PythonOptions;
320    type SymbolTable = SymbolTableScopes;
321
322    fn to_rust(
323        self,
324        ctx: Self::Context,
325        options: Self::Options,
326        symbols: Self::SymbolTable,
327    ) -> std::result::Result<TokenStream, Box<dyn std::error::Error>> {
328        let mut params = Vec::new();
329        
330        // Process positional-only arguments
331        for arg in self.posonlyargs {
332            let param = arg.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
333            params.push(param);
334        }
335        
336        // Process regular positional arguments. Defaulted parameters lower
337        // to plain required parameters: Rust has no default arguments, and
338        // the old Option<T> wrapping neither type-checked against bodies
339        // that use the parameter directly nor matched call sites (which
340        // never wrapped values in Some). Callers that omit the argument
341        // fail to compile either way; callers that pass it now work.
342        for arg in self.args {
343            let param = arg.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
344            params.push(param);
345        }
346        
347        // Process *args
348        if let Some(vararg) = self.vararg {
349            let vararg_name = crate::safe_ident(&vararg.arg);
350            params.push(quote!(#vararg_name: impl IntoIterator<Item = impl Into<PyObject>>));
351        }
352        
353        // Process keyword-only arguments. Like positional defaults above,
354        // these lower to plain required parameters.
355        for arg in self.kwonlyargs {
356            let param = arg.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
357            params.push(param);
358        }
359        
360        // Process **kwargs
361        if let Some(kwarg) = self.kwarg {
362            let kwarg_name = crate::safe_ident(&kwarg.arg);
363            params.push(quote!(#kwarg_name: impl IntoIterator<Item = (impl AsRef<str>, impl Into<PyObject>)>));
364        }
365        
366        Ok(quote!(#(#params),*))
367    }
368}
369
370
371// Implementation for CallArguments
372impl<'a, 'py> FromPyObject<'a, 'py> for CallArguments {
373    type Error = pyo3::PyErr;
374    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
375        let args: Vec<ExprType> = ob.getattr("args")?.extract().unwrap_or_default();
376        let keywords: Vec<crate::Keyword> = ob.getattr("keywords")?.extract().unwrap_or_default();
377        
378        Ok(Self { args, keywords })
379    }
380}
381
382impl CodeGen for CallArguments {
383    type Context = CodeGenContext;
384    type Options = PythonOptions;
385    type SymbolTable = SymbolTableScopes;
386
387    fn to_rust(
388        self,
389        ctx: Self::Context,
390        options: Self::Options,
391        symbols: Self::SymbolTable,
392    ) -> std::result::Result<TokenStream, Box<dyn std::error::Error>> {
393        let mut all_args = Vec::new();
394        
395        // Add positional arguments
396        for arg in self.args {
397            let rust_arg = arg.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
398            all_args.push(rust_arg);
399        }
400        
401        // Add keyword arguments
402        for keyword in self.keywords {
403            let rust_kw = keyword.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
404            all_args.push(rust_kw);
405        }
406        
407        Ok(quote!(#(#all_args),*))
408    }
409}
410
411
412// Node trait implementations for position tracking
413impl Node for Argument {
414    fn lineno(&self) -> Option<usize> { self.lineno }
415    fn col_offset(&self) -> Option<usize> { self.col_offset }
416    fn end_lineno(&self) -> Option<usize> { self.end_lineno }
417    fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
418}
419
420impl Node for Parameter {
421    fn lineno(&self) -> Option<usize> { self.lineno }
422    fn col_offset(&self) -> Option<usize> { self.col_offset }
423    fn end_lineno(&self) -> Option<usize> { self.end_lineno }
424    fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
425}
426
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431    use crate::{parse, CodeGenContext, ExprType, PythonOptions, SymbolTableScopes};
432    use test_log::test;
433
434    #[test]
435    fn test_simple_function_call() {
436        let code = "func(1, 2, 3)";
437        let result = parse(code, "test.py").unwrap();
438        
439        // Generate Rust code
440        let options = PythonOptions::default();
441        let symbols = SymbolTableScopes::new();
442        let _rust_code = result.to_rust(
443            CodeGenContext::Module("test".to_string()),
444            options,
445            symbols,
446        ).unwrap();
447        
448        // Should generate function call with positional arguments
449    }
450
451    #[test]
452    fn test_keyword_arguments() {
453        // Keywords resolve against the callee's signature and land in
454        // parameter order.
455        let code = "def func(a, b):\n    pass\n\nfunc(b=2, a=1)";
456        let result = parse(code, "test.py").unwrap();
457
458        let options = PythonOptions::default();
459        let symbols = result.clone().find_symbols(SymbolTableScopes::new());
460        let rust_code = result.to_rust(
461            CodeGenContext::Module("test".to_string()),
462            options,
463            symbols,
464        ).unwrap().to_string();
465        assert!(rust_code.contains("func (1 , 2)"), "generated: {}", rust_code);
466    }
467
468    #[test]
469    fn test_mixed_arguments() {
470        let code = "def func(a, b, c, d):\n    pass\n\nfunc(1, 2, d=4, c=3)";
471        let result = parse(code, "test.py").unwrap();
472
473        let options = PythonOptions::default();
474        let symbols = result.clone().find_symbols(SymbolTableScopes::new());
475        let rust_code = result.to_rust(
476            CodeGenContext::Module("test".to_string()),
477            options,
478            symbols,
479        ).unwrap().to_string();
480        assert!(rust_code.contains("func (1 , 2 , 3 , 4)"), "generated: {}", rust_code);
481    }
482
483    #[test]
484    fn test_function_with_defaults() {
485        let code = r#"
486def func(a, b=2, c=3):
487    pass
488        "#;
489        let result = parse(code, "test.py").unwrap();
490        
491        let options = PythonOptions::default();
492        let symbols = SymbolTableScopes::new();
493        let _rust_code = result.to_rust(
494            CodeGenContext::Module("test".to_string()),
495            options,
496            symbols,
497        ).unwrap();
498        
499        // Should generate function with optional parameters
500    }
501
502    #[test]
503    fn test_function_with_varargs() {
504        let code = r#"
505def func(a, *args):
506    pass
507        "#;
508        let result = parse(code, "test.py").unwrap();
509        
510        let options = PythonOptions::default();
511        let symbols = SymbolTableScopes::new();
512        let _rust_code = result.to_rust(
513            CodeGenContext::Module("test".to_string()),
514            options,
515            symbols,
516        ).unwrap();
517        
518        // Should generate function with variable arguments
519    }
520
521    #[test]
522    fn test_function_with_kwargs() {
523        let code = r#"
524def func(a, **kwargs):
525    pass
526        "#;
527        let result = parse(code, "test.py").unwrap();
528        
529        let options = PythonOptions::default();
530        let symbols = SymbolTableScopes::new();
531        let _rust_code = result.to_rust(
532            CodeGenContext::Module("test".to_string()),
533            options,
534            symbols,
535        ).unwrap();
536        
537        // Should generate function with keyword arguments dict
538    }
539
540    #[test]
541    fn test_complex_function_signature() {
542        let code = r#"
543def func(a, b=2, *args, c, d=4, **kwargs):
544    pass
545        "#;
546        let result = parse(code, "test.py").unwrap();
547        
548        let options = PythonOptions::default();
549        let symbols = SymbolTableScopes::new();
550        let _rust_code = result.to_rust(
551            CodeGenContext::Module("test".to_string()),
552            options,
553            symbols,
554        ).unwrap();
555        
556        // Should generate function with all argument types
557    }
558
559    #[test]
560    fn test_keyword_only_arguments() {
561        let code = r#"
562def func(a, *, b, c=3):
563    pass
564        "#;
565        let result = parse(code, "test.py").unwrap();
566        
567        let options = PythonOptions::default();
568        let symbols = SymbolTableScopes::new();
569        let _rust_code = result.to_rust(
570            CodeGenContext::Module("test".to_string()),
571            options,
572            symbols,
573        ).unwrap();
574        
575        // Should generate function with keyword-only arguments
576    }
577
578    #[test]
579    fn test_argument_unpacking_call() {
580        // Note: This would require additional AST node support for Starred expressions
581        let code = "func(*args, **kwargs)";
582        let result = parse(code, "test.py");
583        
584        match result {
585            Ok(ast) => {
586                let options = PythonOptions::default();
587                let symbols = SymbolTableScopes::new();
588                let rust_code = ast.to_rust(
589                    CodeGenContext::Module("test".to_string()),
590                    options,
591                    symbols,
592                );
593                
594                match rust_code {
595                    Ok(_code) => { /* Code generation succeeded as expected */ },
596                    Err(_e) => { /* Expected error for unimplemented feature */ },
597                }
598            }
599            Err(_e) => { /* Parse error expected for unimplemented features */ },
600        }
601    }
602
603    #[test]
604    fn test_arg_with_constant() {
605        // Test that Arg (now ExprType) works with constants
606        use litrs::Literal;
607        let literal = Literal::parse("42").unwrap().into_owned();
608        let constant = crate::Constant(Some(literal));
609        let arg: Arg = ExprType::Constant(constant);
610        
611        let options = PythonOptions::default();
612        let symbols = SymbolTableScopes::new();
613        let rust_code = arg.to_rust(
614            CodeGenContext::Module("test".to_string()),
615            options,
616            symbols,
617        ).unwrap();
618        
619        assert!(rust_code.to_string().contains("42"));
620    }
621
622    #[test]
623    fn test_arg_with_name() {
624        // Test that Arg (now ExprType) works with name expressions
625        let name_expr = ExprType::Name(crate::Name {
626            id: "variable".to_string(),
627        });
628        let arg: Arg = name_expr;
629        
630        let options = PythonOptions::default();
631        let symbols = SymbolTableScopes::new();
632        let rust_code = arg.to_rust(
633            CodeGenContext::Module("test".to_string()),
634            options,
635            symbols,
636        ).unwrap();
637        
638        assert!(rust_code.to_string().contains("variable"));
639    }
640}
python_ast/ast/tree/arguments.rs

python_ast/ast/tree/
arguments.rs
