polars_python/expr/
datatype.rs

1use polars::prelude::{DataType, DataTypeExpr, Schema};
2use pyo3::{Bound, IntoPyObject, PyAny, PyResult, Python, pyclass};
3
4use super::PyExpr;
5use crate::error::PyPolarsErr;
6use crate::prelude::Wrap;
7
8#[pyclass]
9#[repr(transparent)]
10#[derive(Clone)]
11pub struct PyDataTypeExpr {
12    pub inner: DataTypeExpr,
13}
14
15impl From<DataTypeExpr> for PyDataTypeExpr {
16    fn from(expr: DataTypeExpr) -> Self {
17        PyDataTypeExpr { inner: expr }
18    }
19}
20
21#[cfg(feature = "pymethods")]
22#[pyo3::pymethods]
23impl PyDataTypeExpr {
24    #[staticmethod]
25    pub fn from_dtype(datatype: Wrap<DataType>) -> Self {
26        DataTypeExpr::Literal(datatype.0).into()
27    }
28
29    #[staticmethod]
30    pub fn of_expr(expr: PyExpr) -> Self {
31        DataTypeExpr::OfExpr(Box::new(expr.inner)).into()
32    }
33
34    pub fn collect_dtype<'py>(
35        &self,
36        py: Python<'py>,
37        schema: Wrap<Schema>,
38    ) -> PyResult<Bound<'py, PyAny>> {
39        let dtype = self
40            .clone()
41            .inner
42            .into_datatype(&schema.0)
43            .map_err(PyPolarsErr::from)?;
44        Wrap(dtype).into_pyobject(py)
45    }
46}