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 file_provider_result: FromPython,
14    pub series: FromPython,
15    pub df: FromPython,
16    pub dsl_plan: FromPython,
17    pub schema: FromPython,
18}
19
20#[derive(Clone)]
21pub struct ToPythonConvertRegistry {
22    pub df: ToPython,
23    pub series: ToPython,
24    pub dsl_plan: ToPython,
25    pub schema: ToPython,
26}
27
28impl ToPythonConvertRegistry {
29    /// Convert a Rust `DataFrame` to a Python `pl.DataFrame` object.
30    pub fn df_to_wrapped_pydf(&self, df: &dyn Any) -> PyResult<Py<PyAny>> {
31        static WRAP_DF: LazyLock<Py<PyAny>> = LazyLock::new(|| {
32            Python::attach(|py| {
33                py.import("polars._utils.wrap")
34                    .unwrap()
35                    .getattr("wrap_df")
36                    .unwrap()
37                    .unbind()
38            })
39        });
40
41        let pydf = (self.df)(df)?;
42
43        Python::attach(|py| WRAP_DF.call1(py, (pydf,)))
44    }
45}
46
47#[derive(Clone)]
48pub struct PythonConvertRegistry {
49    pub from_py: FromPythonConvertRegistry,
50    pub to_py: ToPythonConvertRegistry,
51}
52
53impl PythonConvertRegistry {
54    pub fn py_file_provider_args_dataclass(&self) -> &'static Py<PyAny> {
55        static CLS: LazyLock<Py<PyAny>> = LazyLock::new(|| {
56            Python::attach(|py| {
57                py.import("polars.io.partition")
58                    .unwrap()
59                    .getattr("FileProviderArgs")
60                    .unwrap()
61                    .unbind()
62            })
63        });
64
65        &CLS
66    }
67
68    pub fn py_sinked_paths_callback_args_dataclass(&self) -> &'static Py<PyAny> {
69        static CLS: LazyLock<Py<PyAny>> = LazyLock::new(|| {
70            Python::attach(|py| {
71                py.import("polars.io.partition")
72                    .unwrap()
73                    .getattr("SinkedPathsCallbackArgs")
74                    .unwrap()
75                    .unbind()
76            })
77        });
78
79        &CLS
80    }
81}
82
83static PYTHON_CONVERT_REGISTRY: LazyLock<RwLock<Option<PythonConvertRegistry>>> =
84    LazyLock::new(Default::default);
85
86pub fn get_python_convert_registry() -> PythonConvertRegistry {
87    PYTHON_CONVERT_REGISTRY
88        .deref()
89        .read()
90        .unwrap()
91        .as_ref()
92        .unwrap()
93        .clone()
94}
95
96pub fn register_converters(registry: PythonConvertRegistry) {
97    *PYTHON_CONVERT_REGISTRY.deref().write().unwrap() = Some(registry);
98}
polars_utils/python_convert_registry.rs

polars_utils/
python_convert_registry.rs
