Skip to main content

polars_io/external_reader/
python.rs

1use std::sync::Arc;
2
3use polars_error::PolarsResult;
4use polars_utils::aliases::PlIndexMap;
5use polars_utils::python_function::PythonObject;
6use pyo3::types::{PyAnyMethods, PyDict};
7use pyo3::{
8    Borrowed, Bound, FromPyObject, IntoPyObject, IntoPyObjectExt, Py, PyAny, PyErr, PyResult,
9    Python, intern,
10};
11
12#[derive(Debug, Clone, PartialEq, Eq, Hash)]
13#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
14#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
15pub struct PythonFileReaderBuilder {
16    builder: Arc<PythonObject>,
17}
18
19impl PythonFileReaderBuilder {
20    pub fn new(builder: Py<PyAny>) -> Self {
21        Self {
22            builder: Arc::new(PythonObject(builder)),
23        }
24    }
25
26    pub fn builder(&self) -> &Arc<PythonObject> {
27        &self.builder
28    }
29
30    pub fn explain_properties(&self) -> PolarsResult<PlIndexMap<String, String>> {
31        Python::attach(|py| {
32            let properties: Py<PyDict> = self
33                .builder
34                .call_method0(py, intern!(py, "explain_properties"))?
35                .extract(py)?;
36
37            let properties = properties.bind(py);
38
39            let mut ret = PlIndexMap::default();
40
41            for (k, v) in properties
42                .try_iter()?
43                .zip(properties.call_method0(intern!(py, "values"))?.try_iter()?)
44            {
45                let k: String = k?.extract()?;
46                let v: String = v?.extract()?;
47
48                ret.insert(k, v);
49            }
50
51            Ok(ret)
52        })
53    }
54
55    pub fn init_multi_scan_context(&self, py: Python<'_>, cx: &Py<PyDict>) -> PolarsResult<()> {
56        self.builder
57            .call_method1(py, intern!(py, "init_multi_scan_context"), (cx,))?;
58        Ok(())
59    }
60}
61
62impl<'a, 'py> FromPyObject<'a, 'py> for PythonFileReaderBuilder {
63    type Error = PyErr;
64
65    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
66        Ok(Self {
67            builder: Arc::new(PythonObject(ob.into_py_any(ob.py())?)),
68        })
69    }
70}
71
72impl<'py> IntoPyObject<'py> for PythonFileReaderBuilder {
73    type Target = PyAny;
74    type Output = Bound<'py, Self::Target>;
75    type Error = PyErr;
76
77    fn into_pyobject(self, py: pyo3::prelude::Python<'py>) -> Result<Self::Output, Self::Error> {
78        Ok(self.builder.0.as_ref().clone_ref(py).into_bound(py))
79    }
80}
81
82impl<'py> IntoPyObject<'py> for &PythonFileReaderBuilder {
83    type Target = PyAny;
84    type Output = Bound<'py, Self::Target>;
85    type Error = PyErr;
86
87    fn into_pyobject(self, py: pyo3::prelude::Python<'py>) -> Result<Self::Output, Self::Error> {
88        Ok(self.builder.0.clone_ref(py).into_bound(py))
89    }
90}