# -*- coding: utf-8 -*-
import typing as T
import dataclasses
from datetime import timedelta
import polars as pl
from ..arg import REQ, NA, rm_na, T_KWARGS
from ..base_expr import ExprEnum, BaseExpr, expr_enum_to_klass_mapping, parse_expr
from ..utils_expr import to_jsonpolars_other_expr
if T.TYPE_CHECKING: # pragma: no cover
from .api import T_EXPR
def ensure_datetime(expr: "T_EXPR") -> pl.Expr:
if isinstance(expr, Datetime):
return expr.to_polars()
else:
return expr.to_polars().dt
[docs]@dataclasses.dataclass
class Datetime(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/temporal.html
"""
type: str = dataclasses.field(default=ExprEnum.dt.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr)
expr_enum_to_klass_mapping[ExprEnum.dt.value] = Datetime
[docs]@dataclasses.dataclass
class DtToString(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.to_string.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_to_string.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
format: str = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(
expr=parse_expr(dct["expr"]),
format=dct["format"],
)
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).to_string(format=self.format)
expr_enum_to_klass_mapping[ExprEnum.dt_to_string.value] = DtToString
[docs]@dataclasses.dataclass
class DtYear(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.year.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_year.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).year()
expr_enum_to_klass_mapping[ExprEnum.dt_year.value] = DtYear
[docs]@dataclasses.dataclass
class DtQuarter(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.quarter.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_quarter.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).quarter()
expr_enum_to_klass_mapping[ExprEnum.dt_quarter.value] = DtQuarter
[docs]@dataclasses.dataclass
class DtMonth(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.month.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_month.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: REQ):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).month()
expr_enum_to_klass_mapping[ExprEnum.dt_month.value] = DtMonth
[docs]@dataclasses.dataclass
class DtDay(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.day.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_day.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: REQ):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).day()
expr_enum_to_klass_mapping[ExprEnum.dt_day.value] = DtDay
[docs]@dataclasses.dataclass
class DtHour(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.hour.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_hour.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: REQ):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).hour()
expr_enum_to_klass_mapping[ExprEnum.dt_hour.value] = DtHour
[docs]@dataclasses.dataclass
class DtMinute(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.minute.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_minute.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: REQ):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).minute()
expr_enum_to_klass_mapping[ExprEnum.dt_minute.value] = DtMinute
[docs]@dataclasses.dataclass
class DtSecond(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.second.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_second.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: REQ):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).second()
expr_enum_to_klass_mapping[ExprEnum.dt_second.value] = DtSecond
[docs]@dataclasses.dataclass
class DtNanoSecond(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.nanosecond.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_nanosecond.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).nanosecond()
expr_enum_to_klass_mapping[ExprEnum.dt_nanosecond.value] = DtNanoSecond
[docs]@dataclasses.dataclass
class DtEpoch(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.epoch.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_epoch.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
time_unit: str = dataclasses.field(default=NA)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
req_kwargs, opt_kwargs = cls._split_req_opt(dct)
req_kwargs["expr"] = parse_expr(req_kwargs["expr"])
return cls(**req_kwargs, **rm_na(**opt_kwargs))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).epoch(**rm_na(time_unit=self.time_unit))
expr_enum_to_klass_mapping[ExprEnum.dt_epoch.value] = DtEpoch
[docs]@dataclasses.dataclass
class DtTimestamp(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.timestamp.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_timestamp.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
time_unit: str = dataclasses.field(default=NA)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
req_kwargs, opt_kwargs = cls._split_req_opt(dct)
req_kwargs["expr"] = parse_expr(req_kwargs["expr"])
return cls(**req_kwargs, **rm_na(**opt_kwargs))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).timestamp(**rm_na(time_unit=self.time_unit))
expr_enum_to_klass_mapping[ExprEnum.dt_timestamp.value] = DtTimestamp
[docs]@dataclasses.dataclass
class DtTotalDays(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.total_days.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_total_days.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).total_days()
expr_enum_to_klass_mapping[ExprEnum.dt_total_days.value] = DtTotalDays
[docs]@dataclasses.dataclass
class DtTotalHours(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.total_hours.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_total_hours.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).total_hours()
expr_enum_to_klass_mapping[ExprEnum.dt_total_hours.value] = DtTotalHours
[docs]@dataclasses.dataclass
class DtTotalMinutes(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.total_minutes.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_total_minutes.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).total_minutes()
expr_enum_to_klass_mapping[ExprEnum.dt_total_minutes.value] = DtTotalMinutes
[docs]@dataclasses.dataclass
class DtTotalSeconds(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.total_seconds.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_total_seconds.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).total_seconds()
expr_enum_to_klass_mapping[ExprEnum.dt_total_seconds.value] = DtTotalSeconds
[docs]@dataclasses.dataclass
class DtTotalMilliSeconds(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.total_milliseconds.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_total_milliseconds.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).total_milliseconds()
expr_enum_to_klass_mapping[ExprEnum.dt_total_milliseconds.value] = DtTotalMilliSeconds
[docs]@dataclasses.dataclass
class DtTotalMicroSeconds(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.total_microseconds.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_total_microseconds.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).total_microseconds()
expr_enum_to_klass_mapping[ExprEnum.dt_total_microseconds.value] = DtTotalMicroSeconds
[docs]@dataclasses.dataclass
class DtTotalNanoSeconds(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.total_nanoseconds.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_total_nanoseconds.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(expr=parse_expr(dct["expr"]))
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).total_nanoseconds()
expr_enum_to_klass_mapping[ExprEnum.dt_total_nanoseconds.value] = DtTotalNanoSeconds
[docs]@dataclasses.dataclass
class DtTruncate(BaseExpr):
"""
Ref: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.truncate.html
"""
type: str = dataclasses.field(default=ExprEnum.dt_truncate.value)
expr: "T_EXPR" = dataclasses.field(default=REQ)
every: T.Union[str, timedelta, "T_EXPR"] = dataclasses.field(default=REQ)
[docs] @classmethod
def from_dict(cls, dct: T_KWARGS):
return cls(
expr=parse_expr(dct["expr"]),
every=to_jsonpolars_other_expr(dct["every"]),
)
def to_polars(self) -> pl.Expr:
return ensure_datetime(self.expr).truncate(every=self.every)
expr_enum_to_klass_mapping[ExprEnum.dt_truncate.value] = DtTruncate