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
9pub struct Context {
42 pub(crate) globals: Py<PyDict>,
43}
44
45impl Context {
46 #[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 pub fn globals(&self) -> &Py<PyDict> {
71 &self.globals
72 }
73
74 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 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 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 pub fn run(
146 &self,
147 #[cfg(not(doc))] code: PythonBlock<impl FnOnce(&Bound<PyDict>)>,
148 #[cfg(doc)] code: PythonBlock, ) {
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
169fn 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}