polars_python/functions/
whenthen.rs

1use polars::lazy::dsl;
2use pyo3::prelude::*;
3
4use crate::PyExpr;
5
6#[pyfunction]
7pub fn when(condition: PyExpr) -> PyWhen {
8    PyWhen {
9        inner: dsl::when(condition.inner),
10    }
11}
12
13#[pyclass]
14#[derive(Clone)]
15pub struct PyWhen {
16    inner: dsl::When,
17}
18
19#[pyclass]
20#[derive(Clone)]
21pub struct PyThen {
22    inner: dsl::Then,
23}
24
25#[pyclass]
26#[derive(Clone)]
27pub struct PyChainedWhen {
28    inner: dsl::ChainedWhen,
29}
30
31#[pyclass]
32#[derive(Clone)]
33pub struct PyChainedThen {
34    inner: dsl::ChainedThen,
35}
36
37#[pymethods]
38impl PyWhen {
39    fn then(&self, statement: PyExpr) -> PyThen {
40        PyThen {
41            inner: self.inner.clone().then(statement.inner),
42        }
43    }
44}
45
46#[pymethods]
47impl PyThen {
48    fn when(&self, condition: PyExpr) -> PyChainedWhen {
49        PyChainedWhen {
50            inner: self.inner.clone().when(condition.inner),
51        }
52    }
53
54    fn otherwise(&self, statement: PyExpr) -> PyExpr {
55        self.inner.clone().otherwise(statement.inner).into()
56    }
57}
58
59#[pymethods]
60impl PyChainedWhen {
61    fn then(&self, statement: PyExpr) -> PyChainedThen {
62        PyChainedThen {
63            inner: self.inner.clone().then(statement.inner),
64        }
65    }
66}
67
68#[pymethods]
69impl PyChainedThen {
70    fn when(&self, condition: PyExpr) -> PyChainedWhen {
71        PyChainedWhen {
72            inner: self.inner.clone().when(condition.inner),
73        }
74    }
75
76    fn otherwise(&self, statement: PyExpr) -> PyExpr {
77        self.inner.clone().otherwise(statement.inner).into()
78    }
79}