polars_python/expr/
datatype.rs

1use polars::prelude::{DataType, DataTypeExpr, PlSmallStr, Schema};
2use pyo3::{Bound, IntoPyObject, PyAny, PyResult, Python, pyclass};
3
4use super::PyExpr;
5use super::selector::{PySelector, parse_datatype_selector};
6use crate::error::PyPolarsErr;
7use crate::prelude::Wrap;
8
9#[pyclass]
10#[repr(transparent)]
11#[derive(Clone)]
12pub struct PyDataTypeExpr {
13    pub inner: DataTypeExpr,
14}
15
16impl From<DataTypeExpr> for PyDataTypeExpr {
17    fn from(expr: DataTypeExpr) -> Self {
18        PyDataTypeExpr { inner: expr }
19    }
20}
21
22#[cfg(feature = "pymethods")]
23#[pyo3::pymethods]
24impl PyDataTypeExpr {
25    #[staticmethod]
26    pub fn from_dtype(datatype: Wrap<DataType>) -> Self {
27        DataTypeExpr::Literal(datatype.0).into()
28    }
29
30    #[staticmethod]
31    pub fn of_expr(expr: PyExpr) -> Self {
32        DataTypeExpr::OfExpr(Box::new(expr.inner)).into()
33    }
34
35    #[staticmethod]
36    pub fn self_dtype() -> Self {
37        DataTypeExpr::SelfDtype.into()
38    }
39
40    pub fn collect_dtype<'py>(
41        &self,
42        py: Python<'py>,
43        schema: Wrap<Schema>,
44    ) -> PyResult<Bound<'py, PyAny>> {
45        let dtype = self
46            .clone()
47            .inner
48            .into_datatype(&schema.0)
49            .map_err(PyPolarsErr::from)?;
50        Wrap(dtype).into_pyobject(py)
51    }
52
53    pub fn inner_dtype(&self) -> Self {
54        self.inner.clone().inner_dtype().into()
55    }
56
57    pub fn equals(&self, other: &Self) -> PyExpr {
58        self.inner.clone().equals(other.inner.clone()).into()
59    }
60
61    pub fn display(&self) -> PyExpr {
62        self.inner.clone().display().into()
63    }
64
65    pub fn matches(&self, selector: PySelector) -> PyResult<PyExpr> {
66        let dtype_selector = parse_datatype_selector(selector)?;
67        Ok(self.inner.clone().matches(dtype_selector).into())
68    }
69
70    #[staticmethod]
71    pub fn struct_with_fields(fields: Vec<(String, PyDataTypeExpr)>) -> Self {
72        let fields = fields
73            .into_iter()
74            .map(|(name, dt_expr)| (PlSmallStr::from_string(name), dt_expr.inner))
75            .collect();
76        DataTypeExpr::StructWithFields(fields).into()
77    }
78
79    pub fn wrap_in_list(&self) -> Self {
80        self.inner.clone().wrap_in_list().into()
81    }
82
83    pub fn wrap_in_array(&self, width: usize) -> Self {
84        self.inner.clone().wrap_in_array(width).into()
85    }
86
87    pub fn to_unsigned_integer(&self) -> Self {
88        self.inner.clone().int().to_unsigned().into()
89    }
90
91    pub fn to_signed_integer(&self) -> Self {
92        self.inner.clone().int().to_signed().into()
93    }
94
95    pub fn list_inner_dtype(&self) -> Self {
96        self.inner.clone().list().inner_dtype().into()
97    }
98
99    pub fn arr_inner_dtype(&self) -> Self {
100        self.inner.clone().arr().inner_dtype().into()
101    }
102
103    pub fn arr_width(&self) -> PyExpr {
104        self.inner.clone().arr().width().into()
105    }
106
107    pub fn arr_shape(&self) -> PyExpr {
108        self.inner.clone().arr().shape().into()
109    }
110
111    pub fn struct_field_dtype_by_index(&self, index: i64) -> Self {
112        self.inner
113            .clone()
114            .struct_()
115            .field_dtype_by_index(index)
116            .into()
117    }
118
119    pub fn struct_field_dtype_by_name(&self, name: &str) -> Self {
120        self.inner
121            .clone()
122            .struct_()
123            .field_dtype_by_name(name)
124            .into()
125    }
126
127    pub fn struct_field_names(&self) -> PyExpr {
128        self.inner.clone().struct_().field_names().into()
129    }
130}
polars_python/expr/datatype.rs

polars_python/expr/
datatype.rs
