196 lines
6.3 KiB
Python
196 lines
6.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import polars.functions as F
|
|
from polars._utils.parse import (
|
|
parse_into_expression,
|
|
parse_predicates_constraints_into_expression,
|
|
)
|
|
from polars._utils.wrap import wrap_expr
|
|
from polars.expr.expr import Expr
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterable
|
|
|
|
from polars._plr import PyExpr
|
|
from polars._typing import IntoExpr
|
|
|
|
|
|
class When:
|
|
"""
|
|
Utility class for the `when-then-otherwise` expression.
|
|
|
|
Represents the initial state of the expression after `pl.when(...)` is called.
|
|
|
|
In this state, `then` must be called to continue to finish the expression.
|
|
"""
|
|
|
|
def __init__(self, when: Any) -> None:
|
|
self._when = when
|
|
|
|
def then(self, statement: IntoExpr) -> Then:
|
|
"""
|
|
Attach a statement to the corresponding condition.
|
|
|
|
Parameters
|
|
----------
|
|
statement
|
|
The statement to apply if the corresponding condition is true.
|
|
Accepts expression input. Strings are parsed as column names, other
|
|
non-expression inputs are parsed as literals.
|
|
"""
|
|
statement_pyexpr = parse_into_expression(statement)
|
|
return Then(self._when.then(statement_pyexpr))
|
|
|
|
|
|
class Then(Expr):
|
|
"""
|
|
Utility class for the `when-then-otherwise` expression.
|
|
|
|
Represents the state of the expression after `pl.when(...).then(...)` is called.
|
|
"""
|
|
|
|
def __init__(self, then: Any) -> None:
|
|
self._then = then
|
|
|
|
@classmethod
|
|
def _from_pyexpr(cls, pyexpr: PyExpr) -> Expr:
|
|
return wrap_expr(pyexpr)
|
|
|
|
@property
|
|
def _pyexpr(self) -> PyExpr: # type: ignore[override]
|
|
return self._then.otherwise(F.lit(None)._pyexpr)
|
|
|
|
def when(
|
|
self,
|
|
*predicates: IntoExpr | Iterable[IntoExpr],
|
|
**constraints: Any,
|
|
) -> ChainedWhen:
|
|
"""
|
|
Add a condition to the `when-then-otherwise` expression.
|
|
|
|
Parameters
|
|
----------
|
|
predicates
|
|
Condition(s) that must be met in order to apply the subsequent statement.
|
|
Accepts one or more boolean expressions, which are implicitly combined with
|
|
`&`. String input is parsed as a column name.
|
|
constraints
|
|
Apply conditions as `col_name = value` keyword arguments that are treated as
|
|
equality matches, such as `x = 123`. As with the predicates parameter,
|
|
multiple conditions are implicitly combined using `&`.
|
|
|
|
Notes
|
|
-----
|
|
The expression output name is taken from the first `then` statement. It is
|
|
not affected by `predicates`, nor by `constraints`.
|
|
"""
|
|
condition_pyexpr = parse_predicates_constraints_into_expression(
|
|
*predicates, **constraints
|
|
)
|
|
return ChainedWhen(self._then.when(condition_pyexpr))
|
|
|
|
def otherwise(self, statement: IntoExpr) -> Expr:
|
|
"""
|
|
Define a default for the `when-then-otherwise` expression.
|
|
|
|
Parameters
|
|
----------
|
|
statement
|
|
The statement to apply if all conditions are false.
|
|
Accepts expression input. Strings are parsed as column names, other
|
|
non-expression inputs are parsed as literals.
|
|
"""
|
|
statement_pyexpr = parse_into_expression(statement)
|
|
return wrap_expr(self._then.otherwise(statement_pyexpr))
|
|
|
|
|
|
class ChainedWhen:
|
|
"""
|
|
Utility class for the `when-then-otherwise` expression.
|
|
|
|
Represents the state of the expression after an additional `when` is called.
|
|
|
|
In this state, `then` must be called to continue to finish the expression.
|
|
"""
|
|
|
|
def __init__(self, chained_when: Any) -> None:
|
|
self._chained_when = chained_when
|
|
|
|
def then(self, statement: IntoExpr) -> ChainedThen:
|
|
"""
|
|
Attach a statement to the corresponding condition.
|
|
|
|
Parameters
|
|
----------
|
|
statement
|
|
The statement to apply if the corresponding condition is true.
|
|
Accepts expression input. Strings are parsed as column names, other
|
|
non-expression inputs are parsed as literals.
|
|
"""
|
|
statement_pyexpr = parse_into_expression(statement)
|
|
return ChainedThen(self._chained_when.then(statement_pyexpr))
|
|
|
|
|
|
class ChainedThen(Expr):
|
|
"""
|
|
Utility class for the `when-then-otherwise` expression.
|
|
|
|
Represents the state of the expression after an additional `then` is called.
|
|
"""
|
|
|
|
def __init__(self, chained_then: Any) -> None:
|
|
self._chained_then = chained_then
|
|
|
|
@classmethod
|
|
def _from_pyexpr(cls, pyexpr: PyExpr) -> Expr:
|
|
return wrap_expr(pyexpr)
|
|
|
|
@property
|
|
def _pyexpr(self) -> PyExpr: # type: ignore[override]
|
|
return self._chained_then.otherwise(F.lit(None)._pyexpr)
|
|
|
|
def when(
|
|
self,
|
|
*predicates: IntoExpr | Iterable[IntoExpr],
|
|
**constraints: Any,
|
|
) -> ChainedWhen:
|
|
"""
|
|
Add another condition to the `when-then-otherwise` expression.
|
|
|
|
Parameters
|
|
----------
|
|
predicates
|
|
Condition(s) that must be met in order to apply the subsequent statement.
|
|
Accepts one or more boolean expressions, which are implicitly combined with
|
|
`&`. String input is parsed as a column name.
|
|
constraints
|
|
Apply conditions as `col_name = value` keyword arguments that are treated as
|
|
equality matches, such as `x = 123`. As with the predicates parameter,
|
|
multiple conditions are implicitly combined using `&`.
|
|
|
|
Notes
|
|
-----
|
|
The expression output name is taken from the first `then` statement. It is
|
|
not affected by `predicates`, nor by `constraints`.
|
|
"""
|
|
condition_pyexpr = parse_predicates_constraints_into_expression(
|
|
*predicates, **constraints
|
|
)
|
|
return ChainedWhen(self._chained_then.when(condition_pyexpr))
|
|
|
|
def otherwise(self, statement: IntoExpr) -> Expr:
|
|
"""
|
|
Define a default for the `when-then-otherwise` expression.
|
|
|
|
Parameters
|
|
----------
|
|
statement
|
|
The statement to apply if all conditions are false.
|
|
Accepts expression input. Strings are parsed as column names, other
|
|
non-expression inputs are parsed as literals.
|
|
"""
|
|
statement_pyexpr = parse_into_expression(statement)
|
|
return wrap_expr(self._chained_then.otherwise(statement_pyexpr))
|