python_ast/ast/tree/
import.rs

1use log::debug;
2use proc_macro2::TokenStream;
3use pyo3::FromPyObject;
4use quote::{format_ident, quote};
5use serde::{Deserialize, Serialize};
6
7use crate::{CodeGen, CodeGenContext, PythonOptions, SymbolTableNode, SymbolTableScopes};
8
9#[derive(Clone, Debug, FromPyObject, Serialize, Deserialize, PartialEq)]
10pub struct Alias {
11    pub name: String,
12    pub asname: Option<String>,
13}
14
15#[derive(Clone, Debug, FromPyObject, Serialize, Deserialize, PartialEq)]
16pub struct Import {
17    pub names: Vec<Alias>,
18}
19
20/// An Import (or FromImport) statement causes 2 things to occur:
21/// 1. Declares the imported object within the existing scope.
22/// 2. Causes the referenced module to be compiled into the program (only once).
23
24impl CodeGen for Import {
25    type Context = CodeGenContext;
26    type Options = PythonOptions;
27    type SymbolTable = SymbolTableScopes;
28
29    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
30        let mut symbols = symbols;
31        for alias in self.names.iter() {
32            symbols.insert(alias.name.clone(), SymbolTableNode::Import(self.clone()));
33            if let Some(a) = alias.asname.clone() {
34                symbols.insert(a, SymbolTableNode::Alias(alias.name.clone()))
35            }
36        }
37        symbols
38    }
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        let mut tokens = TokenStream::new();
47        for alias in self.names.iter() {
48            // Check if this is a Python standard library module that needs special handling
49            let rust_import = match alias.name.as_str() {
50                // Python stdlib modules that don't have direct Rust equivalents
51                "os" | "sys" | "subprocess" | "json" | "urllib" | "xml" | "asyncio" => {
52                    // These will be provided by the stdpython runtime
53                    // Generate a comment instead of a use statement
54                    quote! {
55                        // Python module '{}' will be provided by stdpython runtime
56                    }
57                }
58                "os.path" => {
59                    quote! {
60                        // Python os.path module will be provided by stdpython runtime
61                    }
62                }
63                _ => {
64                    // Handle other imports normally
65                    let names = if alias.name.contains('.') {
66                        let parts: Vec<&str> = alias.name.split('.').collect();
67                        let idents: Vec<_> = parts.iter().map(|part| format_ident!("{}", part)).collect();
68                        quote!(#(#idents)::*)
69                    } else {
70                        let single_name = format_ident!("{}", alias.name);
71                        quote!(#single_name)
72                    };
73                    
74                    match &alias.asname {
75                        None => {
76                            quote! {use #names;}
77                        }
78                        Some(n) => {
79                            let name = format_ident!("{}", n);
80                            quote! {use #names as #name;}
81                        }
82                    }
83                }
84            };
85            
86            tokens.extend(rust_import);
87        }
88        debug!("context: {:?}", ctx);
89        debug!("options: {:?}", options);
90        debug!("tokens: {}", tokens);
91        Ok(tokens)
92    }
93}
94
95#[derive(Clone, Debug, FromPyObject, Serialize, Deserialize, PartialEq)]
96pub struct ImportFrom {
97    pub module: String,
98    pub names: Vec<Alias>,
99    pub level: usize,
100}
101
102impl CodeGen for ImportFrom {
103    type Context = CodeGenContext;
104    type Options = PythonOptions;
105    type SymbolTable = SymbolTableScopes;
106
107    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
108        let mut symbols = symbols;
109        for alias in self.names.iter() {
110            symbols.insert(
111                alias.name.clone(),
112                SymbolTableNode::ImportFrom(self.clone()),
113            );
114        }
115        symbols
116    }
117
118    fn to_rust(
119        self,
120        ctx: Self::Context,
121        _options: Self::Options,
122        _symbols: Self::SymbolTable,
123    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
124        debug!("ctx: {:?}", ctx);
125        Ok(quote! {})
126    }
127}
python_ast/ast/tree/import.rs

python_ast/ast/tree/
import.rs
