python_ast/ast/tree/
module.rs

1use std::{collections::HashMap, default::Default};
2
3use log::info;
4use proc_macro2::TokenStream;
5use pyo3::{FromPyObject, PyAny, PyResult};
6use quote::{format_ident, quote};
7use serde::{Deserialize, Serialize};
8
9use crate::{CodeGen, CodeGenContext, Name, Object, PythonOptions, Statement, SymbolTableScopes};
10
11#[derive(Clone, Debug, Serialize, Deserialize)]
12pub enum Type {
13    Unimplemented,
14}
15
16impl<'a> FromPyObject<'a> for Type {
17    fn extract(ob: &'a PyAny) -> PyResult<Self> {
18        info!("Type: {:?}", ob);
19        Ok(Type::Unimplemented)
20    }
21}
22
23/// Represents a module as imported from an ast. See the Module struct for the processed module.
24#[derive(Clone, Debug, Default, FromPyObject, Serialize, Deserialize)]
25pub struct RawModule {
26    pub body: Vec<Statement>,
27    pub type_ignores: Vec<Type>,
28}
29
30/// Represents a module as imported from an ast.
31#[derive(Clone, Debug, Default, Serialize, Deserialize)]
32pub struct Module {
33    pub raw: RawModule,
34    pub name: Option<Name>,
35    pub doc: Option<String>,
36    pub filename: Option<String>,
37    pub attributes: HashMap<Name, String>,
38}
39
40impl<'a> FromPyObject<'a> for Module {
41    fn extract(ob: &'a PyAny) -> PyResult<Self> {
42        let raw_module = RawModule::extract(ob).expect("Failed parsing module.");
43
44        Ok(Self {
45            raw: raw_module,
46            ..Default::default()
47        })
48    }
49}
50
51impl CodeGen for Module {
52    type Context = CodeGenContext;
53    type Options = PythonOptions;
54    type SymbolTable = SymbolTableScopes;
55
56    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
57        let mut symbols = symbols;
58        symbols.new_scope();
59        for s in self.raw.body {
60            symbols = s.clone().find_symbols(symbols);
61        }
62        symbols
63    }
64
65    fn to_rust(
66        self,
67        ctx: Self::Context,
68        options: Self::Options,
69        symbols: Self::SymbolTable,
70    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
71        let mut stream = TokenStream::new();
72        let stdpython = format_ident!("{}", options.stdpython);
73        if options.with_std_python {
74            stream.extend(quote!(use #stdpython::*;));
75        }
76        for s in self.raw.body {
77            let statement = s
78                .clone()
79                .to_rust(ctx.clone(), options.clone(), symbols.clone())
80                .expect(format!("parsing statement {:?} in module", s).as_str());
81            if statement.to_string() != "" {
82                stream.extend(statement);
83            }
84        }
85        Ok(stream)
86    }
87}
88
89impl Object for Module {
90    /// __dir__ is called to list the attributes of the object.
91    fn __dir__(&self) -> Vec<impl AsRef<str>> {
92        // XXX - Make this meaningful.
93        vec![
94            "__class__",
95            "__class_getitem__",
96            "__contains__",
97            "__delattr__",
98            "__delitem__",
99            "__dir__",
100            "__doc__",
101            "__eq__",
102            "__format__",
103            "__ge__",
104            "__getattribute__",
105            "__getitem__",
106            "__getstate__",
107            "__gt__",
108            "__hash__",
109            "__init__",
110            "__init_subclass__",
111            "__ior__",
112            "__iter__",
113            "__le__",
114            "__len__",
115            "__lt__",
116            "__ne__",
117            "__new__",
118            "__or__",
119            "__reduce__",
120            "__reduce_ex__",
121            "__repr__",
122            "__reversed__",
123            "__ror__",
124            "__setattr__",
125            "__setitem__",
126            "__sizeof__",
127            "__str__",
128            "__subclasshook__",
129            "clear",
130            "copy",
131            "fromkeys",
132            "get",
133            "items",
134            "keys",
135            "pop",
136            "popitem",
137            "setdefault",
138            "update",
139            "values",
140        ]
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn can_we_print() {
150        let options = PythonOptions::default();
151        let result = crate::parse(
152            "#test comment
153def foo():
154    print(\"Test print.\")
155",
156            "test_case.py",
157        )
158        .unwrap();
159        info!("Python tree: {:?}", result);
160        //info!("{}", result);
161
162        let code = result.to_rust(
163            CodeGenContext::Module("test_case".to_string()),
164            options,
165            SymbolTableScopes::new(),
166        );
167        info!("module: {:?}", code);
168    }
169
170    #[test]
171    fn can_we_import() {
172        let result = crate::parse("import ast", "ast.py").unwrap();
173        let options = PythonOptions::default();
174        info!("{:?}", result);
175
176        let code = result.to_rust(
177            CodeGenContext::Module("test_case".to_string()),
178            options,
179            SymbolTableScopes::new(),
180        );
181        info!("module: {:?}", code);
182    }
183
184    #[test]
185    fn can_we_import2() {
186        let result = crate::parse("import ast.test as test", "ast.py").unwrap();
187        let options = PythonOptions::default();
188        info!("{:?}", result);
189
190        let code = result.to_rust(
191            CodeGenContext::Module("test_case".to_string()),
192            options,
193            SymbolTableScopes::new(),
194        );
195        info!("module: {:?}", code);
196    }
197}