monarch_types/
python.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::collections::HashMap;
10
11use monarch_gil::GilSite;
12use monarch_gil::monarch_with_gil_blocking;
13use pyo3::Bound;
14use pyo3::IntoPyObject;
15use pyo3::IntoPyObjectExt;
16use pyo3::PyAny;
17use pyo3::PyResult;
18use pyo3::Python;
19use pyo3::prelude::*;
20use pyo3::types::PyDict;
21use pyo3::types::PyNone;
22use pyo3::types::PyTuple;
23use serde::Deserialize;
24use serde::Serialize;
25
26/// A variant of `pyo3::IntoPyObject` used to wrap unsafe impls and propagates the
27/// unsafety to the caller.
28pub trait TryIntoPyObjectUnsafe<'py, P> {
29    /// # Safety
30    ///
31    /// The caller must ensure self is valid for the duration of the call
32    /// and that the type-erased pointer invariants of the trait are upheld.
33    unsafe fn try_to_object_unsafe(self, py: Python<'py>) -> PyResult<Bound<'py, P>>;
34}
35
36/// Helper impl for casting into args for python functions calls.
37impl<'a, 'py, T> TryIntoPyObjectUnsafe<'py, PyTuple> for &'a Vec<T>
38where
39    &'a T: TryIntoPyObjectUnsafe<'py, PyAny>,
40    T: 'a,
41{
42    unsafe fn try_to_object_unsafe(self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
43        PyTuple::new(
44            py,
45            self.iter()
46                // SAFETY: Safety requirements are propagated via the `unsafe`
47                // tag on this method.
48                .map(|v| unsafe { v.try_to_object_unsafe(py) })
49                .collect::<Result<Vec<_>, _>>()?,
50        )
51    }
52}
53
54/// Helper impl for casting into kwargs for python functions calls.
55impl<'a, 'py, K, V> TryIntoPyObjectUnsafe<'py, PyDict> for &'a HashMap<K, V>
56where
57    &'a K: IntoPyObject<'py> + std::cmp::Eq + std::hash::Hash,
58    &'a V: TryIntoPyObjectUnsafe<'py, PyAny>,
59    K: 'a,
60    V: 'a,
61{
62    unsafe fn try_to_object_unsafe(self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
63        let dict = PyDict::new(py);
64        for (key, val) in self {
65            // SAFETY: Safety requirements are propagated via the `unsafe`
66            // tag on this method.
67            dict.set_item(key, unsafe { val.try_to_object_unsafe(py) }?)?;
68        }
69        Ok(dict)
70    }
71}
72
73/// A wrapper around `PyErr` that contains a serialized traceback.
74#[derive(Debug, Clone, Serialize, Deserialize, derive_more::Error)]
75pub struct SerializablePyErr {
76    pub message: String,
77}
78
79impl SerializablePyErr {
80    pub fn from(py: Python, err: &PyErr) -> Self {
81        // first construct the full traceback including any python frames that were used
82        // to invoke where we currently are. This is pre-pended to the traceback of the
83        // currently unwinded frames (err.traceback())
84        let inspect = py.import("inspect").unwrap();
85        let types = py.import("types").unwrap();
86        let traceback_type = types.getattr("TracebackType").unwrap();
87        let traceback = py.import("traceback").unwrap();
88
89        let mut f = inspect
90            .call_method0("currentframe")
91            .unwrap_or(PyNone::get(py).to_owned().into_any());
92        let mut tb: Bound<'_, PyAny> = err.traceback(py).into_bound_py_any(py).unwrap();
93        while !f.is_none() {
94            let lasti = f.getattr("f_lasti").unwrap();
95            let lineno = f.getattr("f_lineno").unwrap();
96            let back = f.getattr("f_back").unwrap();
97            tb = traceback_type.call1((tb, f, lasti, lineno)).unwrap();
98            f = back;
99        }
100
101        let traceback_exception = traceback.getattr("TracebackException").unwrap();
102
103        let tb = traceback_exception
104            .call1((err.get_type(py), err.value(py), tb))
105            .unwrap();
106
107        let message: String = tb
108            .getattr("format")
109            .unwrap()
110            .call0()
111            .unwrap()
112            .try_iter()
113            .unwrap()
114            .map(|x| -> String { x.unwrap().extract().unwrap() })
115            .collect::<Vec<String>>()
116            .join("");
117
118        Self { message }
119    }
120
121    pub fn from_fn<'py>(py: Python<'py>) -> impl Fn(PyErr) -> Self + 'py {
122        move |err| Self::from(py, &err)
123    }
124}
125
126impl std::fmt::Display for SerializablePyErr {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        write!(f, "{}", self.message)
129    }
130}
131
132impl<T> From<T> for SerializablePyErr
133where
134    T: Into<PyErr>,
135{
136    fn from(value: T) -> Self {
137        monarch_with_gil_blocking(GilSite::Traceback, |py| {
138            SerializablePyErr::from(py, &value.into())
139        })
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use monarch_gil::GilSite;
146    use monarch_gil::monarch_with_gil_blocking;
147    use pyo3::Python;
148    use pyo3::ffi::c_str;
149    use pyo3::indoc::indoc;
150    use pyo3::prelude::*;
151    use timed_test::async_timed_test;
152
153    use crate::SerializablePyErr;
154
155    #[async_timed_test(timeout_secs = 60)]
156    async fn test_serializable_py_err() {
157        Python::initialize();
158        let _unused = monarch_with_gil_blocking(GilSite::Test, |py| {
159            let module = PyModule::from_code(
160                py,
161                c_str!(indoc! {r#"
162                        def func1():
163                            raise Exception("test")
164
165                        def func2():
166                            func1()
167
168                        def func3():
169                            func2()
170                    "#}),
171                c_str!("test_helpers.py"),
172                c_str!("test_helpers"),
173            )?;
174
175            let err = SerializablePyErr::from(py, &module.call_method0("func3").unwrap_err());
176            assert_eq!(
177                err.message.as_str(),
178                indoc! {r#"
179                    Traceback (most recent call last):
180                      File "test_helpers.py", line 8, in func3
181                      File "test_helpers.py", line 5, in func2
182                      File "test_helpers.py", line 2, in func1
183                    Exception: test
184                "#}
185            );
186
187            PyResult::Ok(())
188        });
189    }
190}
monarch_types/python.rs

monarch_types/
python.rs
