litellm/litellm-rust/crates/host-python/src/marshal.rs

use std::any::Any;
use std::panic::{AssertUnwindSafe, catch_unwind};

use pyo3::exceptions::PyValueError;
use pyo3::panic::PanicException;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use serde::Serialize;
use serde::de::DeserializeOwned;

/// Converts a `#[pyo3(from_py_with = ...)]` argument, reporting failures as `ValueError`
/// so a bad argument reads as a bad argument rather than as whatever the conversion hit.
pub fn from_py_argument<T>(value: &Bound<'_, PyAny>) -> PyResult<T>
where
    T: DeserializeOwned,
{
    pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string()))
}

pub fn from_py<T>(value: &Bound<'_, PyAny>) -> PyResult<T>
where
    T: DeserializeOwned,
{
    pythonize::depythonize(value).map_err(PyErr::from)
}

pub fn to_py<T>(py: Python<'_>, value: &T) -> PyResult<Py<PyAny>>
where
    T: Serialize + ?Sized,
{
    pythonize::pythonize(py, value)
        .map(Bound::unbind)
        .map_err(PyErr::from)
}

pub fn json_object_field(py: Python<'_>, document: &str, name: &str) -> PyResult<Py<PyAny>> {
    py.import("json")?
        .call_method1("loads", (document,))?
        .call_method1("get", (name,))
        .map(Bound::unbind)
}

pub fn json_loads(py: Python<'_>, document: &[u8]) -> PyResult<Py<PyAny>> {
    py.import("json")?
        .call_method1("loads", (PyBytes::new(py, document),))
        .map(Bound::unbind)
}

pub struct Pythonized<T>(pub T);

impl<'py, T> IntoPyObject<'py> for Pythonized<T>
where
    T: Serialize,
{
    type Target = PyAny;
    type Output = Bound<'py, PyAny>;
    type Error = PyErr;

    fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
        catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0)))
            .map_err(panic_to_pyerr)?
            .map_err(PyErr::from)
    }
}

pub fn panic_to_pyerr(payload: Box<dyn Any + Send>) -> PyErr {
    let message = payload
        .downcast_ref::<String>()
        .map(String::as_str)
        .or_else(|| payload.downcast_ref::<&str>().copied())
        .unwrap_or("panic from Rust code");
    PanicException::new_err(message.to_string())
}

#[cfg(test)]
mod tests {
    use serde::Serializer;

    use super::*;

    struct PanickingSerializer;

    impl Serialize for PanickingSerializer {
        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
        where
            S: Serializer,
        {
            panic!("serializer panicked")
        }
    }

    #[test]
    fn pythonized_converts_on_the_attached_thread() {
        crate::initialize_python();
        Python::attach(|py| {
            let value: Vec<i32> = Pythonized(vec![1, 2, 3])
                .into_pyobject(py)
                .and_then(|value| value.extract())
                .expect("value should convert");
            assert_eq!(value, vec![1, 2, 3]);
        });
    }

    #[test]
    fn pythonized_preserves_python_serialization_error_types() {
        crate::initialize_python();
        Python::attach(|py| {
            let value = std::collections::BTreeMap::from([(vec![1], "value")]);
            let direct = to_py(py, &value).unwrap_err();
            let wrapped = Pythonized(value).into_pyobject(py).unwrap_err();
            assert!(direct.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
            assert!(wrapped.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
            assert_eq!(wrapped.to_string(), direct.to_string());
        });
    }

    #[test]
    fn pythonized_maps_serializer_panics_to_a_base_exception() {
        crate::initialize_python();
        Python::attach(|py| {
            let error = Pythonized(PanickingSerializer)
                .into_pyobject(py)
                .expect_err("serializer panic should become a Python exception");
            assert!(error.is_instance_of::<PanicException>(py));
            assert_eq!(error.to_string(), "PanicException: serializer panicked");
        });
    }

    #[test]
    fn depythonize_preserves_python_exception_identity_and_traceback() {
        crate::initialize_python();
        Python::attach(|py| {
            let locals = pyo3::types::PyDict::new(py);
            py.run(
                pyo3::ffi::c_str!(
                    r#"
failure = LookupError('conversion failed')
cause = ValueError('cause')
class Broken:
    def __index__(self):
        raise failure from cause
value = Broken()
"#
                ),
                Some(&locals),
                Some(&locals),
            )
            .unwrap();
            let value = locals.get_item("value").unwrap().unwrap();
            let argument_error = from_py_argument::<i64>(&value).unwrap_err();
            assert!(argument_error.is_instance_of::<PyValueError>(py));
            assert!(
                !argument_error
                    .value(py)
                    .is(locals.get_item("failure").unwrap().unwrap())
            );
            let error = from_py::<i64>(&value).unwrap_err();
            assert!(
                error
                    .value(py)
                    .is(locals.get_item("failure").unwrap().unwrap())
            );
            assert!(
                error
                    .cause(py)
                    .unwrap()
                    .value(py)
                    .is(locals.get_item("cause").unwrap().unwrap())
            );
            assert!(error.traceback(py).is_some());
        });
    }
}