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(frozen)]
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 default_value(&self, n: usize, numeric_to_one: bool, num_list_values: usize) -> PyExpr {
96        self.inner
97            .clone()
98            .default_value(n, numeric_to_one, num_list_values)
99            .into()
100    }
101
102    pub fn list_inner_dtype(&self) -> Self {
103        self.inner.clone().list().inner_dtype().into()
104    }
105
106    pub fn arr_inner_dtype(&self) -> Self {
107        self.inner.clone().arr().inner_dtype().into()
108    }
109
110    pub fn arr_width(&self) -> PyExpr {
111        self.inner.clone().arr().width().into()
112    }
113
114    pub fn arr_shape(&self) -> PyExpr {
115        self.inner.clone().arr().shape().into()
116    }
117
118    pub fn struct_field_dtype_by_index(&self, index: i64) -> Self {
119        self.inner
120            .clone()
121            .struct_()
122            .field_dtype_by_index(index)
123            .into()
124    }
125
126    pub fn struct_field_dtype_by_name(&self, name: &str) -> Self {
127        self.inner
128            .clone()
129            .struct_()
130            .field_dtype_by_name(name)
131            .into()
132    }
133
134    pub fn struct_field_names(&self) -> PyExpr {
135        self.inner.clone().struct_().field_names().into()
136    }
137}
polars_python/expr/datatype.rs

polars_python/expr/
datatype.rs
