inline_python/
context.rs

1use crate::run::run_python_code;
2use crate::PythonBlock;
3use pyo3::{
4	prelude::*,
5	types::{PyCFunction, PyDict},
6	FromPyObject, IntoPyObject, Py, PyResult, Python,
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	pub fn new() -> Self {
51		Python::with_gil(Self::new_with_gil)
52	}
53
54	pub(crate) fn new_with_gil(py: Python) -> Self {
55		match Self::try_new(py) {
56			Ok(x) => x,
57			Err(error) => {
58				error.print(py);
59				panic!("failed to create Python context");
60			}
61		}
62	}
63
64	fn try_new(py: Python) -> PyResult<Self> {
65		Ok(Self {
66			globals: py.import("__main__")?.dict().copy()?.into(),
67		})
68	}
69
70	/// Get the globals as dictionary.
71	pub fn globals<'p>(&self) -> &Py<PyDict> {
72		&self.globals
73	}
74
75	/// Retrieve a global variable from the context.
76	///
77	/// This function panics if the variable doesn't exist, or the conversion fails.
78	pub fn get<T: for<'p> FromPyObject<'p>>(&self, name: &str) -> T {
79		Python::with_gil(|py| match self.globals.bind(py).get_item(name) {
80			Err(_) | Ok(None) => panic!("Python context does not contain a variable named `{}`", name),
81			Ok(Some(value)) => match FromPyObject::extract_bound(&value) {
82				Ok(value) => value,
83				Err(e) => {
84					e.print(py);
85					panic!("Unable to convert `{}` to `{}`", name, std::any::type_name::<T>());
86				}
87			},
88		})
89	}
90
91	/// Set a global variable in the context.
92	///
93	/// This function panics if the conversion fails.
94	pub fn set<T: for<'p> IntoPyObject<'p>>(&self, name: &str, value: T) {
95		Python::with_gil(|py| match self.globals().bind(py).set_item(name, value) {
96			Ok(()) => (),
97			Err(e) => {
98				e.print(py);
99				panic!("Unable to set `{}` from a `{}`", name, std::any::type_name::<T>());
100			}
101		})
102	}
103
104	/// Add a wrapped `#[pyfunction]` or `#[pymodule]` using its own `__name__`.
105	///
106	/// Use this with `pyo3::wrap_pyfunction` or `pyo3::wrap_pymodule`.
107	///
108	/// ```ignore
109	/// # use inline_python::{Context, python};
110	/// use pyo3::{prelude::*, wrap_pyfunction};
111	///
112	/// #[pyfunction]
113	/// fn get_five() -> i32 {
114	///     5
115	/// }
116	///
117	/// fn main() {
118	///     let c = Context::new();
119	///
120	///     c.add_wrapped(wrap_pyfunction!(get_five));
121	///
122	///     c.run(python! {
123	///         assert get_five() == 5
124	///     });
125	/// }
126	/// ```
127	pub fn add_wrapped(&self, wrapper: &impl Fn(Python) -> PyResult<Bound<'_, PyCFunction>>) {
128		Python::with_gil(|py| {
129			let obj = wrapper(py).unwrap();
130			let name = obj.getattr("__name__").expect("Missing __name__");
131			if let Err(e) = self.globals().bind(py).set_item(name, obj) {
132				e.print(py);
133				panic!("Unable to add wrapped function");
134			}
135		})
136	}
137
138	/// Run Python code using this context.
139	///
140	/// This function should be called using the `python!{}` macro:
141	///
142	/// ```
143	/// # use inline_python::{Context, python};
144	/// let c = Context::new();
145	///
146	/// c.run(python!{
147	///     print("Hello World")
148	/// });
149	/// ```
150	///
151	/// This function panics if the Python code fails.
152	pub fn run<F: FnOnce(&Bound<PyDict>)>(&self, code: PythonBlock<F>) {
153		Python::with_gil(|py| self.run_with_gil(py, code));
154	}
155
156	pub(crate) fn run_with_gil<F: FnOnce(&Bound<PyDict>)>(&self, py: Python<'_>, code: PythonBlock<F>) {
157		(code.set_variables)(self.globals().bind(py));
158		match run_python_code(py, self, code.bytecode) {
159			Ok(_) => (),
160			Err(e) => {
161				e.print(py);
162				panic!("{}", "python!{...} failed to execute");
163			}
164		}
165	}
166}
inline_python/context.rs

inline_python/
context.rs
