inline_python/
context.rs

1use crate::PythonBlock;
2use crate::run::run_python_code;
3use pyo3::{
4	FromPyObject, IntoPyObject, Py, PyResult, Python,
5	prelude::*,
6	types::{PyCFunction, PyDict},
7};
8
9/// An execution context for Python code.
10///
11/// This can be used to keep all global variables and imports intact between macro invocations:
12///
13/// ```
14/// # use inline_python::{Context, python};
15/// let c = Context::new();
16///
17/// c.run(python! {
18///     foo = 5
19/// });
20///
21/// c.run(python! {
22///     assert foo == 5
23/// });
24/// ```
25///
26/// You may also use it to inspect global variables after the execution of the Python code,
27/// or set global variables before running:
28///
29/// ```
30/// # use inline_python::{Context, python};
31/// let c = Context::new();
32///
33/// c.set("x", 13);
34///
35/// c.run(python! {
36///     foo = x + 2
37/// });
38///
39/// assert_eq!(c.get::<i32>("foo"), 15);
40/// ```
41pub struct Context {
42	pub(crate) globals: Py<PyDict>,
43}
44
45impl Context {
46	/// Create a new context for running Python code.
47	///
48	/// This function panics if it fails to create the context.
49	#[allow(clippy::new_without_default)]
50	#[track_caller]
51	pub fn new() -> Self {
52		Python::with_gil(Self::new_with_gil)
53	}
54
55	#[track_caller]
56	pub(crate) fn new_with_gil(py: Python) -> Self {
57		match Self::try_new(py) {
58			Ok(x) => x,
59			Err(err) => panic!("{}", panic_string(py, &err)),
60		}
61	}
62
63	fn try_new(py: Python) -> PyResult<Self> {
64		Ok(Self {
65			globals: py.import("__main__")?.dict().copy()?.into(),
66		})
67	}
68
69	/// Get the globals as dictionary.
70	pub fn globals(&self) -> &Py<PyDict> {
71		&self.globals
72	}
73
74	/// Retrieve a global variable from the context.
75	///
76	/// This function panics if the variable doesn't exist, or the conversion fails.
77	pub fn get<T: for<'p> FromPyObject<'p>>(&self, name: &str) -> T {
78		Python::with_gil(|py| match self.globals.bind(py).get_item(name) {
79			Err(_) | Ok(None) => panic!("Python context does not contain a variable named `{}`", name),
80			Ok(Some(value)) => match FromPyObject::extract_bound(&value) {
81				Ok(value) => value,
82				Err(e) => panic!("Unable to convert `{}` to `{}`: {e}", name, std::any::type_name::<T>()),
83			},
84		})
85	}
86
87	/// Set a global variable in the context.
88	///
89	/// This function panics if the conversion fails.
90	pub fn set<T: for<'p> IntoPyObject<'p>>(&self, name: &str, value: T) {
91		Python::with_gil(|py| {
92			if let Err(e) = self.globals().bind(py).set_item(name, value) {
93				panic!("Unable to set `{}` from a `{}`: {e}", name, std::any::type_name::<T>());
94			}
95		})
96	}
97
98	/// Add a wrapped `#[pyfunction]` or `#[pymodule]` using its own `__name__`.
99	///
100	/// Use this with `pyo3::wrap_pyfunction` or `pyo3::wrap_pymodule`.
101	///
102	/// ```ignore
103	/// # use inline_python::{Context, python};
104	/// use pyo3::{prelude::*, wrap_pyfunction};
105	///
106	/// #[pyfunction]
107	/// fn get_five() -> i32 {
108	///     5
109	/// }
110	///
111	/// fn main() {
112	///     let c = Context::new();
113	///
114	///     c.add_wrapped(wrap_pyfunction!(get_five));
115	///
116	///     c.run(python! {
117	///         assert get_five() == 5
118	///     });
119	/// }
120	/// ```
121	pub fn add_wrapped(&self, wrapper: &impl Fn(Python) -> PyResult<Bound<'_, PyCFunction>>) {
122		Python::with_gil(|py| {
123			let obj = wrapper(py).unwrap();
124			let name = obj.getattr("__name__").expect("wrapped item should have a __name__");
125			if let Err(err) = self.globals().bind(py).set_item(name, obj) {
126				panic!("{}", panic_string(py, &err));
127			}
128		})
129	}
130
131	/// Run Python code using this context.
132	///
133	/// This function should be called using the `python!{}` macro:
134	///
135	/// ```
136	/// # use inline_python::{Context, python};
137	/// let c = Context::new();
138	///
139	/// c.run(python!{
140	///     print("Hello World")
141	/// });
142	/// ```
143	///
144	/// This function panics if the Python code fails.
145	pub fn run(
146		&self,
147		#[cfg(not(doc))] code: PythonBlock<impl FnOnce(&Bound<PyDict>)>,
148		#[cfg(doc)] code: PythonBlock, // Just show 'PythonBlock' in the docs.
149	) {
150		Python::with_gil(|py| self.run_with_gil(py, code));
151	}
152
153	#[cfg(not(doc))]
154	pub(crate) fn run_with_gil<F: FnOnce(&Bound<PyDict>)>(&self, py: Python<'_>, block: PythonBlock<F>) {
155		(block.set_variables)(self.globals().bind(py));
156		if let Err(err) = run_python_code(py, self, block.bytecode) {
157			(block.panic)(panic_string(py, &err));
158		}
159	}
160}
161
162fn panic_string(py: Python, err: &PyErr) -> String {
163	match py_err_to_string(py, &err) {
164		Ok(msg) => msg,
165		Err(_) => err.to_string(),
166	}
167}
168
169/// Print the error while capturing stderr into a String.
170fn py_err_to_string(py: Python, err: &PyErr) -> Result<String, PyErr> {
171	let sys = py.import("sys")?;
172	let stderr = py.import("io")?.getattr("StringIO")?.call0()?;
173	let original_stderr = sys.dict().get_item("stderr")?;
174	sys.dict().set_item("stderr", &stderr)?;
175	err.print(py);
176	sys.dict().set_item("stderr", original_stderr)?;
177	stderr.call_method0("getvalue")?.extract()
178}