portalis_transpiler/
python_ast.rs

1//! Python AST types for translation to Rust
2//!
3//! Defines a simplified Python AST that can be easily translated to Rust.
4//! Based on the 527 Python language features cataloged in PYTHON_LANGUAGE_FEATURES.md
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9/// Python literal types
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub enum PyLiteral {
12    /// Integer literal: 42, 0x2A, 0o52, 0b101010
13    Int(i64),
14    /// Float literal: 3.14, 1.0e10
15    Float(f64),
16    /// String literal: "hello", 'world', r"raw", f"formatted"
17    String(String),
18    /// Boolean literal: True, False
19    Bool(bool),
20    /// None literal
21    None,
22    /// Bytes literal: b"data"
23    Bytes(Vec<u8>),
24}
25
26/// Python expression types
27#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28pub enum PyExpr {
29    /// Literal value
30    Literal(PyLiteral),
31    /// Variable name
32    Name(String),
33    /// Binary operation: a + b, x * y
34    BinOp {
35        left: Box<PyExpr>,
36        op: BinOp,
37        right: Box<PyExpr>,
38    },
39    /// Unary operation: -x, not y
40    UnaryOp {
41        op: UnaryOp,
42        operand: Box<PyExpr>,
43    },
44    /// Function call: func(args)
45    Call {
46        func: Box<PyExpr>,
47        args: Vec<PyExpr>,
48        kwargs: HashMap<String, PyExpr>,
49    },
50    /// Attribute access: obj.attr
51    Attribute {
52        value: Box<PyExpr>,
53        attr: String,
54    },
55    /// Subscript: list[0], dict["key"]
56    Subscript {
57        value: Box<PyExpr>,
58        index: Box<PyExpr>,
59    },
60    /// Slice: list[1:3], list[:5], list[2:], list[1:10:2]
61    Slice {
62        value: Box<PyExpr>,
63        lower: Option<Box<PyExpr>>,
64        upper: Option<Box<PyExpr>>,
65        step: Option<Box<PyExpr>>,
66    },
67    /// List literal: [1, 2, 3]
68    List(Vec<PyExpr>),
69    /// Tuple literal: (1, 2, 3)
70    Tuple(Vec<PyExpr>),
71    /// Dict literal: {"a": 1, "b": 2}
72    Dict {
73        keys: Vec<PyExpr>,
74        values: Vec<PyExpr>,
75    },
76    /// Set literal: {1, 2, 3}
77    Set(Vec<PyExpr>),
78    /// List comprehension: [x*2 for x in range(10)]
79    ListComp {
80        element: Box<PyExpr>,
81        generators: Vec<Comprehension>,
82    },
83    /// Conditional expression: x if condition else y
84    IfExp {
85        test: Box<PyExpr>,
86        body: Box<PyExpr>,
87        orelse: Box<PyExpr>,
88    },
89    /// Lambda: lambda x: x + 1
90    Lambda {
91        args: Vec<String>,
92        body: Box<PyExpr>,
93    },
94    /// Comparison: x == y, a < b
95    Compare {
96        left: Box<PyExpr>,
97        op: CmpOp,
98        right: Box<PyExpr>,
99    },
100    /// Boolean operation: and, or
101    BoolOp {
102        op: BoolOp,
103        left: Box<PyExpr>,
104        right: Box<PyExpr>,
105    },
106    /// Await expression: await expr
107    Await(Box<PyExpr>),
108    /// Yield expression: yield expr
109    Yield(Option<Box<PyExpr>>),
110}
111
112/// Binary operators
113#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
114pub enum BinOp {
115    Add,      // +
116    Sub,      // -
117    Mult,     // *
118    Div,      // /
119    FloorDiv, // //
120    Mod,      // %
121    Pow,      // **
122    LShift,   // <<
123    RShift,   // >>
124    BitOr,    // |
125    BitXor,   // ^
126    BitAnd,   // &
127    MatMult,  // @ (Python 3.5+)
128}
129
130/// Unary operators
131#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
132pub enum UnaryOp {
133    Invert, // ~
134    Not,    // not
135    UAdd,   // +
136    USub,   // -
137}
138
139/// Comparison operators
140#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
141pub enum CmpOp {
142    Eq,    // ==
143    NotEq, // !=
144    Lt,    // <
145    LtE,   // <=
146    Gt,    // >
147    GtE,   // >=
148    Is,    // is
149    IsNot, // is not
150    In,    // in
151    NotIn, // not in
152}
153
154/// Boolean operators
155#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
156pub enum BoolOp {
157    And, // and
158    Or,  // or
159}
160
161/// Comprehension clause: for x in iterable if condition
162#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
163pub struct Comprehension {
164    pub target: PyExpr,
165    pub iter: PyExpr,
166    pub ifs: Vec<PyExpr>,
167}
168
169/// Function parameter with type annotation
170#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
171pub struct FunctionParam {
172    pub name: String,
173    pub type_annotation: Option<TypeAnnotation>,
174    pub default_value: Option<PyExpr>,
175}
176
177/// Type annotation for variables and function parameters
178#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
179pub enum TypeAnnotation {
180    /// Simple type name: int, str, bool
181    Name(String),
182    /// Generic type: List[int], Dict[str, int]
183    Generic {
184        base: Box<TypeAnnotation>,
185        args: Vec<TypeAnnotation>,
186    },
187}
188
189/// Python statement types
190#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
191pub enum PyStmt {
192    /// Expression statement
193    Expr(PyExpr),
194    /// Assignment: x = 42
195    Assign {
196        target: PyExpr,
197        value: PyExpr,
198    },
199    /// Augmented assignment: x += 1
200    AugAssign {
201        target: PyExpr,
202        op: BinOp,
203        value: PyExpr,
204    },
205    /// Annotated assignment: x: int = 42
206    AnnAssign {
207        target: PyExpr,
208        annotation: TypeAnnotation,
209        value: Option<PyExpr>,
210    },
211    /// Function definition
212    FunctionDef {
213        name: String,
214        params: Vec<FunctionParam>,
215        body: Vec<PyStmt>,
216        return_type: Option<TypeAnnotation>,
217        decorators: Vec<PyExpr>,
218        is_async: bool,
219    },
220    /// Return statement
221    Return {
222        value: Option<PyExpr>,
223    },
224    /// If statement
225    If {
226        test: PyExpr,
227        body: Vec<PyStmt>,
228        orelse: Vec<PyStmt>,
229    },
230    /// While loop
231    While {
232        test: PyExpr,
233        body: Vec<PyStmt>,
234        orelse: Vec<PyStmt>,
235    },
236    /// For loop
237    For {
238        target: PyExpr,
239        iter: PyExpr,
240        body: Vec<PyStmt>,
241        orelse: Vec<PyStmt>,
242    },
243    /// Pass statement
244    Pass,
245    /// Break statement
246    Break,
247    /// Continue statement
248    Continue,
249    /// Class definition
250    ClassDef {
251        name: String,
252        bases: Vec<PyExpr>,
253        body: Vec<PyStmt>,
254        decorators: Vec<PyExpr>,
255    },
256    /// Import statement: import module [as alias]
257    Import {
258        modules: Vec<(String, Option<String>)>,
259    },
260    /// From import: from module import name [as alias]
261    ImportFrom {
262        module: Option<String>,
263        names: Vec<(String, Option<String>)>,
264        level: usize,
265    },
266    /// Assert statement: assert condition, message
267    Assert {
268        test: PyExpr,
269        msg: Option<PyExpr>,
270    },
271    /// Try-except statement
272    Try {
273        body: Vec<PyStmt>,
274        handlers: Vec<ExceptHandler>,
275        orelse: Vec<PyStmt>,
276        finalbody: Vec<PyStmt>,
277    },
278    /// Raise statement
279    Raise {
280        exception: Option<PyExpr>,
281    },
282    /// With statement (context manager)
283    With {
284        items: Vec<WithItem>,
285        body: Vec<PyStmt>,
286    },
287    /// Delete statement: del x
288    Delete {
289        targets: Vec<PyExpr>,
290    },
291    /// Global declaration: global x, y
292    Global {
293        names: Vec<String>,
294    },
295    /// Nonlocal declaration: nonlocal x, y
296    Nonlocal {
297        names: Vec<String>,
298    },
299}
300
301/// With item for context managers
302#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
303pub struct WithItem {
304    pub context_expr: PyExpr,
305    pub optional_vars: Option<PyExpr>,
306}
307
308/// Exception handler for try-except
309#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
310pub struct ExceptHandler {
311    pub exception_type: Option<PyExpr>,
312    pub name: Option<String>,
313    pub body: Vec<PyStmt>,
314}
315
316/// Function argument
317#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
318pub struct Arg {
319    pub name: String,
320    pub type_hint: Option<String>,
321    pub default: Option<PyExpr>,
322}
323
324/// Python module (top-level)
325#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
326pub struct PyModule {
327    pub statements: Vec<PyStmt>,
328}
329
330impl PyModule {
331    pub fn new() -> Self {
332        Self {
333            statements: vec![],
334        }
335    }
336
337    pub fn add_stmt(&mut self, stmt: PyStmt) {
338        self.statements.push(stmt);
339    }
340
341    // Alias for compatibility
342    pub fn body(&self) -> &[PyStmt] {
343        &self.statements
344    }
345}
346
347impl Default for PyModule {
348    fn default() -> Self {
349        Self::new()
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356
357    #[test]
358    fn test_literal_creation() {
359        let int_lit = PyLiteral::Int(42);
360        assert_eq!(int_lit, PyLiteral::Int(42));
361
362        let str_lit = PyLiteral::String("hello".to_string());
363        assert_eq!(str_lit, PyLiteral::String("hello".to_string()));
364
365        let bool_lit = PyLiteral::Bool(true);
366        assert_eq!(bool_lit, PyLiteral::Bool(true));
367    }
368
369    #[test]
370    fn test_simple_assignment() {
371        let module = PyModule {
372            statements: vec![PyStmt::Assign {
373                target: PyExpr::Name("x".to_string()),
374                value: PyExpr::Literal(PyLiteral::Int(42)),
375            }],
376        };
377
378        assert_eq!(module.statements.len(), 1);
379    }
380
381    #[test]
382    fn test_binary_operation() {
383        let expr = PyExpr::BinOp {
384            left: Box::new(PyExpr::Literal(PyLiteral::Int(2))),
385            op: BinOp::Add,
386            right: Box::new(PyExpr::Literal(PyLiteral::Int(3))),
387        };
388
389        match expr {
390            PyExpr::BinOp { left, op, right } => {
391                assert_eq!(*left, PyExpr::Literal(PyLiteral::Int(2)));
392                assert_eq!(op, BinOp::Add);
393                assert_eq!(*right, PyExpr::Literal(PyLiteral::Int(3)));
394            }
395            _ => panic!("Expected BinOp"),
396        }
397    }
398
399    #[test]
400    fn test_function_definition() {
401        let func = PyStmt::FunctionDef {
402            name: "add".to_string(),
403            params: vec![
404                FunctionParam {
405                    name: "a".to_string(),
406                    type_annotation: Some(TypeAnnotation::Name("int".to_string())),
407                    default_value: None,
408                },
409                FunctionParam {
410                    name: "b".to_string(),
411                    type_annotation: Some(TypeAnnotation::Name("int".to_string())),
412                    default_value: None,
413                },
414            ],
415            body: vec![PyStmt::Return {
416                value: Some(PyExpr::BinOp {
417                    left: Box::new(PyExpr::Name("a".to_string())),
418                    op: BinOp::Add,
419                    right: Box::new(PyExpr::Name("b".to_string())),
420                }),
421            }],
422            return_type: Some(TypeAnnotation::Name("int".to_string())),
423            decorators: vec![],
424            is_async: false,
425        };
426
427        match func {
428            PyStmt::FunctionDef { name, params, .. } => {
429                assert_eq!(name, "add");
430                assert_eq!(params.len(), 2);
431            }
432            _ => panic!("Expected FunctionDef"),
433        }
434    }
435}
portalis_transpiler/python_ast.rs

portalis_transpiler/
python_ast.rs
