Skip to main content

python_ast/codegen/
python_options.rs

1//! Options for Python compilation.
2
3use std::{
4    collections::{BTreeMap, HashSet},
5    default::Default,
6};
7
8use pyo3::{prelude::*, PyResult};
9use std::ffi::CString;
10
11/// Supported async runtimes for Python async code generation
12#[derive(Clone, Debug, PartialEq)]
13pub enum AsyncRuntime {
14    /// Tokio runtime (default)
15    Tokio,
16    /// async-std runtime
17    AsyncStd,
18    /// smol runtime
19    Smol,
20    /// Custom runtime with specified attribute and import
21    Custom {
22        /// The attribute to use (e.g., "tokio::main", "async_std::main")
23        attribute: String,
24        /// The import to add (e.g., "tokio", "async_std")
25        import: String,
26    },
27}
28
29impl Default for AsyncRuntime {
30    fn default() -> Self {
31        AsyncRuntime::Tokio
32    }
33}
34
35impl AsyncRuntime {
36    /// Get the attribute string for the async main function
37    pub fn main_attribute(&self) -> &str {
38        match self {
39            AsyncRuntime::Tokio => "tokio::main",
40            AsyncRuntime::AsyncStd => "async_std::main",
41            AsyncRuntime::Smol => "smol::main",
42            AsyncRuntime::Custom { attribute, .. } => attribute,
43        }
44    }
45
46    /// Get the import string for the runtime
47    pub fn import(&self) -> &str {
48        match self {
49            AsyncRuntime::Tokio => "tokio",
50            AsyncRuntime::AsyncStd => "async_std",
51            AsyncRuntime::Smol => "smol",
52            AsyncRuntime::Custom { import, .. } => import,
53        }
54    }
55}
56
57pub fn sys_path() -> PyResult<Vec<String>> {
58    let pymodule_code = include_str!("path.py");
59
60    Python::attach(|py| -> PyResult<Vec<String>> {
61        let code_cstr = CString::new(pymodule_code)?;
62        let pymodule = PyModule::from_code(py, &code_cstr, c"path.py", c"path")?;
63        let t = pymodule.getattr("path")?;
64        assert!(t.is_callable());
65        let args = ();
66        let paths: Vec<String> = t.call1(args)?.extract()?;
67
68        Ok(paths)
69    })
70}
71
72/// The global context for Python compilation.
73#[derive(Clone, Debug)]
74pub struct PythonOptions {
75    /// Python imports are mapped into a given namespace that can be changed.
76    pub python_namespace: String,
77
78    /// The default path we will search for Python modules.
79    pub python_path: Vec<String>,
80
81    /// Collects all of the things we need to compile imports[module][asnames]
82    pub imports: BTreeMap<String, HashSet<String>>,
83
84
85    pub stdpython: String,
86    pub with_std_python: bool,
87
88    pub allow_unsafe: bool,
89
90    /// The async runtime to use for async Python code
91    pub async_runtime: AsyncRuntime,
92
93    /// Emit #[deprecated] notes on generated items whose conversion was
94    /// lossy (dropped parameter defaults, ignored return annotations, ...).
95    /// On by default: silent semantic divergence from the Python source is
96    /// exactly what these warnings exist to surface. Tools may disable this
97    /// to suppress the warnings at the user's explicit request.
98    pub lossy_warnings: bool,
99
100    /// Names in the CURRENT scope that hold an Option (assigned None on
101    /// some path, or annotated Optional): non-None stores into them wrap
102    /// in Some. Set per scope by the function/module generators.
103    pub optional_names: std::rc::Rc<std::collections::HashSet<String>>,
104
105    /// Whether the CURRENT function's return annotation is `str`: returning
106    /// an attribute chain then clones the String field out of the shared
107    /// receiver. Python strings are immutable, so the clone reproduces
108    /// Python's aliasing semantics exactly. Set per function.
109    pub clone_str_attribute_returns: bool,
110
111    /// Statically-known types of names in the CURRENT scope (parameter
112    /// annotations and literal assignments), as canonical Python type
113    /// names ("int", "float", "str", "bool"). Set per function; consumed
114    /// by isinstance(), which lowers to a constant.
115    pub local_types:
116        std::rc::Rc<std::collections::HashMap<String, String>>,
117
118    /// Target stdpython's no_std (alloc) tier. Python constructs that need
119    /// the OS — print/input/open, imports of os/datetime/random/…, and
120    /// `__main__` entry points — fail loudly at conversion time instead of
121    /// surfacing as resolution errors when the generated crate builds.
122    pub no_std: bool,
123}
124
125impl Default for PythonOptions {
126    fn default() -> Self {
127        Self {
128            python_namespace: String::from("__python_namespace__"),
129            // Default must not panic: fall back to an empty search path if
130            // the embedded interpreter can't report sys.path.
131            python_path: sys_path().unwrap_or_else(|e| {
132                tracing::warn!("could not read Python sys.path: {}; using empty path", e);
133                Vec::new()
134            }),
135            imports: BTreeMap::new(),
136
137            stdpython: "stdpython".to_string(),
138            with_std_python: true,
139            allow_unsafe: false,
140            async_runtime: AsyncRuntime::default(),
141            lossy_warnings: true,
142            optional_names: std::rc::Rc::new(std::collections::HashSet::new()),
143            clone_str_attribute_returns: false,
144            local_types: std::rc::Rc::new(std::collections::HashMap::new()),
145            no_std: false,
146        }
147    }
148}
149
150impl PythonOptions {
151    /// Create PythonOptions with tokio runtime (default)
152    pub fn with_tokio() -> Self {
153        let mut options = Self::default();
154        options.async_runtime = AsyncRuntime::Tokio;
155        options
156    }
157
158    /// Create PythonOptions with async-std runtime
159    pub fn with_async_std() -> Self {
160        let mut options = Self::default();
161        options.async_runtime = AsyncRuntime::AsyncStd;
162        options
163    }
164
165    /// Create PythonOptions with smol runtime
166    pub fn with_smol() -> Self {
167        let mut options = Self::default();
168        options.async_runtime = AsyncRuntime::Smol;
169        options
170    }
171
172    /// Create PythonOptions with a custom async runtime
173    pub fn with_custom_runtime(attribute: impl Into<String>, import: impl Into<String>) -> Self {
174        let mut options = Self::default();
175        options.async_runtime = AsyncRuntime::Custom {
176            attribute: attribute.into(),
177            import: import.into(),
178        };
179        options
180    }
181
182    /// Set the async runtime for these options
183    pub fn set_async_runtime(&mut self, runtime: AsyncRuntime) -> &mut Self {
184        self.async_runtime = runtime;
185        self
186    }
187}