polars_utils/
python_convert_registry.rs

1use std::any::Any;
2use std::ops::Deref;
3use std::sync::{Arc, LazyLock, RwLock};
4
5use pyo3::types::PyAnyMethods;
6use pyo3::{Py, PyAny, PyResult, Python};
7
8pub type FromPython = Arc<dyn Fn(Py<PyAny>) -> PyResult<Box<dyn Any>> + Send + Sync>;
9pub type ToPython = Arc<dyn for<'a> Fn(&'a dyn Any) -> PyResult<Py<PyAny>> + Send + Sync>;
10
11#[derive(Clone)]
12pub struct FromPythonConvertRegistry {
13    pub partition_target_cb_result: FromPython,
14    pub file_provider_result: FromPython,
15    pub series: FromPython,
16    pub df: FromPython,
17    pub dsl_plan: FromPython,
18    pub schema: FromPython,
19}
20
21#[derive(Clone)]
22pub struct ToPythonConvertRegistry {
23    pub df: ToPython,
24    pub series: ToPython,
25    pub dsl_plan: ToPython,
26    pub schema: ToPython,
27}
28
29impl ToPythonConvertRegistry {
30    /// Convert a Rust `DataFrame` to a Python `pl.DataFrame` object.
31    pub fn df_to_wrapped_pydf(&self, df: &dyn Any) -> PyResult<Py<PyAny>> {
32        static WRAP_DF: LazyLock<Py<PyAny>> = LazyLock::new(|| {
33            Python::attach(|py| {
34                py.import("polars._utils.wrap")
35                    .unwrap()
36                    .getattr("wrap_df")
37                    .unwrap()
38                    .unbind()
39            })
40        });
41
42        let pydf = (self.df)(df)?;
43
44        Python::attach(|py| WRAP_DF.call1(py, (pydf,)))
45    }
46}
47
48#[derive(Clone)]
49pub struct PythonConvertRegistry {
50    pub from_py: FromPythonConvertRegistry,
51    pub to_py: ToPythonConvertRegistry,
52}
53
54impl PythonConvertRegistry {
55    pub fn py_file_provider_args_dataclass(&self) -> &'static Py<PyAny> {
56        static CLS: LazyLock<Py<PyAny>> = LazyLock::new(|| {
57            Python::attach(|py| {
58                py.import("polars.io.partition")
59                    .unwrap()
60                    .getattr("FileProviderArgs")
61                    .unwrap()
62                    .unbind()
63            })
64        });
65
66        &CLS
67    }
68}
69
70static PYTHON_CONVERT_REGISTRY: LazyLock<RwLock<Option<PythonConvertRegistry>>> =
71    LazyLock::new(Default::default);
72
73pub fn get_python_convert_registry() -> PythonConvertRegistry {
74    PYTHON_CONVERT_REGISTRY
75        .deref()
76        .read()
77        .unwrap()
78        .as_ref()
79        .unwrap()
80        .clone()
81}
82
83pub fn register_converters(registry: PythonConvertRegistry) {
84    *PYTHON_CONVERT_REGISTRY.deref().write().unwrap() = Some(registry);
85}
polars_utils/python_convert_registry.rs

polars_utils/
python_convert_registry.rs
