polars_python/functions/
whenthen.rs

1use polars::lazy::dsl;
2use pyo3::prelude::*;
3use pyo3::types::PyBytes;
4
5use crate::PyExpr;
6
7#[pyfunction]
8pub fn when(condition: PyExpr) -> PyWhen {
9    PyWhen {
10        inner: dsl::when(condition.inner),
11    }
12}
13
14#[pyclass(frozen, skip_from_py_object)]
15#[derive(Clone)]
16pub struct PyWhen {
17    inner: dsl::When,
18}
19
20#[pyclass(skip_from_py_object)] // Not marked as frozen for pickling, but that's the only &mut self method.
21#[derive(Clone)]
22pub struct PyThen {
23    inner: dsl::Then,
24}
25
26#[pyclass(frozen, skip_from_py_object)]
27#[derive(Clone)]
28pub struct PyChainedWhen {
29    inner: dsl::ChainedWhen,
30}
31
32#[pyclass(skip_from_py_object)] // Not marked as frozen for pickling, but that's the only &mut self method.
33#[derive(Clone)]
34pub struct PyChainedThen {
35    inner: dsl::ChainedThen,
36}
37
38#[pymethods]
39impl PyWhen {
40    fn then(&self, statement: PyExpr) -> PyThen {
41        PyThen {
42            inner: self.inner.clone().then(statement.inner),
43        }
44    }
45}
46
47#[pymethods]
48impl PyThen {
49    fn when(&self, condition: PyExpr) -> PyChainedWhen {
50        PyChainedWhen {
51            inner: self.inner.clone().when(condition.inner),
52        }
53    }
54
55    fn otherwise(&self, statement: PyExpr) -> PyExpr {
56        self.inner.clone().otherwise(statement.inner).into()
57    }
58
59    fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
60        crate::conversion::serde_pickle(&self.inner, py)
61    }
62
63    fn __setstate__(&mut self, state: &Bound<PyAny>) -> PyResult<()> {
64        crate::conversion::serde_unpickle(&mut self.inner, state)
65    }
66}
67
68#[pymethods]
69impl PyChainedWhen {
70    fn then(&self, statement: PyExpr) -> PyChainedThen {
71        PyChainedThen {
72            inner: self.inner.clone().then(statement.inner),
73        }
74    }
75}
76
77#[pymethods]
78impl PyChainedThen {
79    fn when(&self, condition: PyExpr) -> PyChainedWhen {
80        PyChainedWhen {
81            inner: self.inner.clone().when(condition.inner),
82        }
83    }
84
85    fn otherwise(&self, statement: PyExpr) -> PyExpr {
86        self.inner.clone().otherwise(statement.inner).into()
87    }
88
89    fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
90        crate::conversion::serde_pickle(&self.inner, py)
91    }
92
93    fn __setstate__(&mut self, state: &Bound<PyAny>) -> PyResult<()> {
94        crate::conversion::serde_unpickle(&mut self.inner, state)
95    }
96}
polars_python/functions/whenthen.rs

polars_python/functions/
whenthen.rs
