pyo3_stub_gen/
util.rs

1use pyo3::{prelude::*, types::*};
2use std::ffi::CString;
3
4pub fn all_builtin_types(any: &Bound<'_, PyAny>) -> bool {
5    if any.is_instance_of::<PyString>()
6        || any.is_instance_of::<PyBool>()
7        || any.is_instance_of::<PyInt>()
8        || any.is_instance_of::<PyFloat>()
9        || any.is_none()
10    {
11        return true;
12    }
13    if any.is_instance_of::<PyDict>() {
14        return any
15            .downcast::<PyDict>()
16            .map(|dict| {
17                dict.into_iter()
18                    .all(|(k, v)| all_builtin_types(&k) && all_builtin_types(&v))
19            })
20            .unwrap_or(false);
21    }
22    if any.is_instance_of::<PyList>() {
23        return any
24            .downcast::<PyList>()
25            .map(|list| list.into_iter().all(|v| all_builtin_types(&v)))
26            .unwrap_or(false);
27    }
28    if any.is_instance_of::<PyTuple>() {
29        return any
30            .downcast::<PyTuple>()
31            .map(|list| list.into_iter().all(|v| all_builtin_types(&v)))
32            .unwrap_or(false);
33    }
34    false
35}
36
37/// whether eval(repr(any)) == any
38pub fn valid_external_repr(any: &Bound<'_, PyAny>) -> Option<bool> {
39    let globals = get_globals(any).ok()?;
40    let fmt_str = any.repr().ok()?.to_string();
41    let fmt_cstr = CString::new(fmt_str.clone()).ok()?;
42    let new_any = any.py().eval(&fmt_cstr, Some(&globals), None).ok()?;
43    new_any.eq(any).ok()
44}
45
46fn get_globals<'py>(any: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyDict>> {
47    let type_object = any.get_type();
48    let type_name = type_object.getattr("__name__")?;
49    let type_name: &str = type_name.extract()?;
50    let globals = PyDict::new(any.py());
51    globals.set_item(type_name, type_object)?;
52    Ok(globals)
53}
54
55pub fn fmt_py_obj<'py, T: pyo3::IntoPyObjectExt<'py>>(py: Python<'py>, obj: T) -> String {
56    if let Ok(any) = obj.into_bound_py_any(py) {
57        if all_builtin_types(&any) || valid_external_repr(&any).is_some_and(|valid| valid) {
58            if let Ok(py_str) = any.repr() {
59                return py_str.to_string();
60            }
61        }
62    }
63    "...".to_owned()
64}
65
66#[cfg(test)]
67mod test {
68    use super::*;
69    #[pyclass]
70    #[derive(Debug)]
71    struct A {}
72    #[test]
73    fn test_fmt_dict() {
74        pyo3::prepare_freethreaded_python();
75        Python::with_gil(|py| {
76            let dict = PyDict::new(py);
77            _ = dict.set_item("k1", "v1");
78            _ = dict.set_item("k2", 2);
79            assert_eq!("{'k1': 'v1', 'k2': 2}", fmt_py_obj(py, &dict));
80            // class A variable can not be formatted
81            _ = dict.set_item("k3", A {});
82            assert_eq!("...", fmt_py_obj(py, &dict));
83        })
84    }
85    #[test]
86    fn test_fmt_list() {
87        pyo3::prepare_freethreaded_python();
88        Python::with_gil(|py| {
89            let list = PyList::new(py, [1, 2]).unwrap();
90            assert_eq!("[1, 2]", fmt_py_obj(py, &list));
91            // class A variable can not be formatted
92            let list = PyList::new(py, [A {}, A {}]).unwrap();
93            assert_eq!("...", fmt_py_obj(py, &list));
94        })
95    }
96    #[test]
97    fn test_fmt_tuple() {
98        pyo3::prepare_freethreaded_python();
99        Python::with_gil(|py| {
100            let tuple = PyTuple::new(py, [1, 2]).unwrap();
101            assert_eq!("(1, 2)", fmt_py_obj(py, tuple));
102            let tuple = PyTuple::new(py, [1]).unwrap();
103            assert_eq!("(1,)", fmt_py_obj(py, tuple));
104            // class A variable can not be formatted
105            let tuple = PyTuple::new(py, [A {}]).unwrap();
106            assert_eq!("...", fmt_py_obj(py, tuple));
107        })
108    }
109    #[test]
110    fn test_fmt_other() {
111        pyo3::prepare_freethreaded_python();
112        Python::with_gil(|py| {
113            // str
114            assert_eq!("'123'", fmt_py_obj(py, &"123"));
115            assert_eq!("\"don't\"", fmt_py_obj(py, &"don't"));
116            assert_eq!("'str\\\\'", fmt_py_obj(py, &"str\\"));
117            // bool
118            assert_eq!("True", fmt_py_obj(py, true));
119            assert_eq!("False", fmt_py_obj(py, false));
120            // int
121            assert_eq!("123", fmt_py_obj(py, 123));
122            // float
123            assert_eq!("1.23", fmt_py_obj(py, 1.23));
124            // None
125            let none: Option<usize> = None;
126            assert_eq!("None", fmt_py_obj(py, none));
127            // class A variable can not be formatted
128            assert_eq!("...", fmt_py_obj(py, A {}));
129        })
130    }
131    #[test]
132    fn test_fmt_enum() {
133        #[pyclass(eq, eq_int)]
134        #[derive(Debug, Clone, PartialEq, Eq, Hash)]
135        pub enum Number {
136            Float,
137            Integer,
138        }
139        pyo3::prepare_freethreaded_python();
140        Python::with_gil(|py| {
141            assert_eq!("Number.Float", fmt_py_obj(py, Number::Float));
142        });
143    }
144}
pyo3_stub_gen/util.rs

pyo3_stub_gen/
util.rs
