python_ast/codegen/
python_options.rs1use std::{
4 collections::{BTreeMap, HashSet},
5 default::Default,
6};
7
8use pyo3::{prelude::*, PyResult};
9use std::ffi::CString;
10
11#[derive(Clone, Debug, PartialEq)]
13pub enum AsyncRuntime {
14 Tokio,
16 AsyncStd,
18 Smol,
20 Custom {
22 attribute: String,
24 import: String,
26 },
27}
28
29impl Default for AsyncRuntime {
30 fn default() -> Self {
31 AsyncRuntime::Tokio
32 }
33}
34
35impl AsyncRuntime {
36 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 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#[derive(Clone, Debug)]
74pub struct PythonOptions {
75 pub python_namespace: String,
77
78 pub python_path: Vec<String>,
80
81 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 pub async_runtime: AsyncRuntime,
92
93 pub lossy_warnings: bool,
99
100 pub optional_names: std::rc::Rc<std::collections::HashSet<String>>,
104
105 pub clone_str_attribute_returns: bool,
110
111 pub local_types:
116 std::rc::Rc<std::collections::HashMap<String, String>>,
117
118 pub no_std: bool,
123}
124
125impl Default for PythonOptions {
126 fn default() -> Self {
127 Self {
128 python_namespace: String::from("__python_namespace__"),
129 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 pub fn with_tokio() -> Self {
153 let mut options = Self::default();
154 options.async_runtime = AsyncRuntime::Tokio;
155 options
156 }
157
158 pub fn with_async_std() -> Self {
160 let mut options = Self::default();
161 options.async_runtime = AsyncRuntime::AsyncStd;
162 options
163 }
164
165 pub fn with_smol() -> Self {
167 let mut options = Self::default();
168 options.async_runtime = AsyncRuntime::Smol;
169 options
170 }
171
172 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 pub fn set_async_runtime(&mut self, runtime: AsyncRuntime) -> &mut Self {
184 self.async_runtime = runtime;
185 self
186 }
187}