python_ast/codegen/
python_options.rs

1//! Options for Python compilation.
2
3use std::{
4    collections::{BTreeMap, HashSet},
5    default::Default,
6};
7
8use crate::Scope;
9use pyo3::{prelude::*, PyResult};
10use std::ffi::CString;
11
12/// Supported async runtimes for Python async code generation
13#[derive(Clone, Debug, PartialEq)]
14pub enum AsyncRuntime {
15    /// Tokio runtime (default)
16    Tokio,
17    /// async-std runtime
18    AsyncStd,
19    /// smol runtime
20    Smol,
21    /// Custom runtime with specified attribute and import
22    Custom {
23        /// The attribute to use (e.g., "tokio::main", "async_std::main")
24        attribute: String,
25        /// The import to add (e.g., "tokio", "async_std")
26        import: String,
27    },
28}
29
30impl Default for AsyncRuntime {
31    fn default() -> Self {
32        AsyncRuntime::Tokio
33    }
34}
35
36impl AsyncRuntime {
37    /// Get the attribute string for the async main function
38    pub fn main_attribute(&self) -> &str {
39        match self {
40            AsyncRuntime::Tokio => "tokio::main",
41            AsyncRuntime::AsyncStd => "async_std::main",
42            AsyncRuntime::Smol => "smol::main",
43            AsyncRuntime::Custom { attribute, .. } => attribute,
44        }
45    }
46
47    /// Get the import string for the runtime
48    pub fn import(&self) -> &str {
49        match self {
50            AsyncRuntime::Tokio => "tokio",
51            AsyncRuntime::AsyncStd => "async_std",
52            AsyncRuntime::Smol => "smol",
53            AsyncRuntime::Custom { import, .. } => import,
54        }
55    }
56}
57
58pub fn sys_path() -> PyResult<Vec<String>> {
59    let pymodule_code = include_str!("path.py");
60
61    Python::with_gil(|py| -> PyResult<Vec<String>> {
62        let code_cstr = CString::new(pymodule_code)?;
63        let pymodule = PyModule::from_code(py, &code_cstr, c"path.py", c"path")?;
64        let t = pymodule
65            .getattr("path")
66            .expect("Reading path variable from interpretter");
67        assert!(t.is_callable());
68        let args = ();
69        let paths: Vec<String> = t.call1(args)?.extract()?;
70
71        Ok(paths)
72    })
73}
74
75/// The global context for Python compilation.
76#[derive(Clone, Debug)]
77pub struct PythonOptions {
78    /// Python imports are mapped into a given namespace that can be changed.
79    pub python_namespace: String,
80
81    /// The default path we will search for Python modules.
82    pub python_path: Vec<String>,
83
84    /// Collects all of the things we need to compile imports[module][asnames]
85    pub imports: BTreeMap<String, HashSet<String>>,
86
87    pub scope: Scope,
88
89    pub stdpython: String,
90    pub with_std_python: bool,
91
92    pub allow_unsafe: bool,
93
94    /// The async runtime to use for async Python code
95    pub async_runtime: AsyncRuntime,
96}
97
98impl Default for PythonOptions {
99    fn default() -> Self {
100        Self {
101            python_namespace: String::from("__python_namespace__"),
102            // XXX: Remove unwrap.
103            python_path: sys_path().unwrap(),
104            imports: BTreeMap::new(),
105            scope: Scope::default(),
106            stdpython: "stdpython".to_string(),
107            with_std_python: true,
108            allow_unsafe: false,
109            async_runtime: AsyncRuntime::default(),
110        }
111    }
112}
113
114impl PythonOptions {
115    /// Create PythonOptions with tokio runtime (default)
116    pub fn with_tokio() -> Self {
117        let mut options = Self::default();
118        options.async_runtime = AsyncRuntime::Tokio;
119        options
120    }
121
122    /// Create PythonOptions with async-std runtime
123    pub fn with_async_std() -> Self {
124        let mut options = Self::default();
125        options.async_runtime = AsyncRuntime::AsyncStd;
126        options
127    }
128
129    /// Create PythonOptions with smol runtime
130    pub fn with_smol() -> Self {
131        let mut options = Self::default();
132        options.async_runtime = AsyncRuntime::Smol;
133        options
134    }
135
136    /// Create PythonOptions with a custom async runtime
137    pub fn with_custom_runtime(attribute: impl Into<String>, import: impl Into<String>) -> Self {
138        let mut options = Self::default();
139        options.async_runtime = AsyncRuntime::Custom {
140            attribute: attribute.into(),
141            import: import.into(),
142        };
143        options
144    }
145
146    /// Set the async runtime for these options
147    pub fn set_async_runtime(&mut self, runtime: AsyncRuntime) -> &mut Self {
148        self.async_runtime = runtime;
149        self
150    }
151}
python_ast/codegen/python_options.rs

python_ast/codegen/
python_options.rs
