polars_python/functions/
whenthen.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
use polars::lazy::dsl;
use pyo3::prelude::*;

use crate::PyExpr;

#[pyfunction]
pub fn when(condition: PyExpr) -> PyWhen {
    PyWhen {
        inner: dsl::when(condition.inner),
    }
}

#[pyclass]
#[derive(Clone)]
pub struct PyWhen {
    inner: dsl::When,
}

#[pyclass]
#[derive(Clone)]
pub struct PyThen {
    inner: dsl::Then,
}

#[pyclass]
#[derive(Clone)]
pub struct PyChainedWhen {
    inner: dsl::ChainedWhen,
}

#[pyclass]
#[derive(Clone)]
pub struct PyChainedThen {
    inner: dsl::ChainedThen,
}

#[pymethods]
impl PyWhen {
    fn then(&self, statement: PyExpr) -> PyThen {
        PyThen {
            inner: self.inner.clone().then(statement.inner),
        }
    }
}

#[pymethods]
impl PyThen {
    fn when(&self, condition: PyExpr) -> PyChainedWhen {
        PyChainedWhen {
            inner: self.inner.clone().when(condition.inner),
        }
    }

    fn otherwise(&self, statement: PyExpr) -> PyExpr {
        self.inner.clone().otherwise(statement.inner).into()
    }
}

#[pymethods]
impl PyChainedWhen {
    fn then(&self, statement: PyExpr) -> PyChainedThen {
        PyChainedThen {
            inner: self.inner.clone().then(statement.inner),
        }
    }
}

#[pymethods]
impl PyChainedThen {
    fn when(&self, condition: PyExpr) -> PyChainedWhen {
        PyChainedWhen {
            inner: self.inner.clone().when(condition.inner),
        }
    }

    fn otherwise(&self, statement: PyExpr) -> PyExpr {
        self.inner.clone().otherwise(statement.inner).into()
    }
}