reflex_copy/reflex/vars/datetime.py

"""Immutable datetime and date vars."""

from __future__ import annotations

import dataclasses
from datetime import date, datetime
from typing import TypeVar, Union

from .base import (
    CustomVarOperationReturn,
    LiteralVar,
    Var,
    VarData,
    var_operation,
    var_operation_return,
)

DATETIME_T = TypeVar("DATETIME_T", datetime, date)

datetime_types = Union[datetime, date]


def date_compare_operation(
    lhs: Var[datetime_types],
    rhs: Var[datetime_types],
    strict: bool = False,
) -> CustomVarOperationReturn[bool]:
    """Check if the value is less than the other value.

    Args:
        lhs: The left-hand side of the operation.
        rhs: The right-hand side of the operation.
        strict: Whether to use strict comparison.

    Returns:
        The result of the operation.
    """
    return var_operation_return(
        f"({lhs} {'<' if strict else '<='} {rhs})",
        bool,
    )


@var_operation
def date_gt_operation(
    lhs: Var[datetime_types],
    rhs: Var[datetime_types],
) -> CustomVarOperationReturn:
    """Greater than comparison.

    Args:
        lhs: The left-hand side of the operation.
        rhs: The right-hand side of the operation.

    Returns:
        The result of the operation.
    """
    return date_compare_operation(rhs, lhs, strict=True)


@var_operation
def date_lt_operation(
    lhs: Var[datetime_types],
    rhs: Var[datetime_types],
) -> CustomVarOperationReturn:
    """Less than comparison.

    Args:
        lhs: The left-hand side of the operation.
        rhs: The right-hand side of the operation.

    Returns:
        The result of the operation.
    """
    return date_compare_operation(lhs, rhs, strict=True)


@var_operation
def date_le_operation(
    lhs: Var[datetime_types], rhs: Var[datetime_types]
) -> CustomVarOperationReturn:
    """Less than or equal comparison.

    Args:
        lhs: The left-hand side of the operation.
        rhs: The right-hand side of the operation.

    Returns:
        The result of the operation.
    """
    return date_compare_operation(lhs, rhs)


@var_operation
def date_ge_operation(
    lhs: Var[datetime_types], rhs: Var[datetime_types]
) -> CustomVarOperationReturn:
    """Greater than or equal comparison.

    Args:
        lhs: The left-hand side of the operation.
        rhs: The right-hand side of the operation.

    Returns:
        The result of the operation.
    """
    return date_compare_operation(rhs, lhs)


class DateTimeVar(Var[DATETIME_T], python_types=(datetime, date)):
    """A variable that holds a datetime or date object."""

    __lt__ = date_lt_operation

    __le__ = date_le_operation

    __gt__ = date_gt_operation

    __ge__ = date_ge_operation


@dataclasses.dataclass(
    eq=False,
    frozen=True,
    slots=True,
)
class LiteralDatetimeVar(LiteralVar, DateTimeVar):
    """Base class for immutable datetime and date vars."""

    _var_value: datetime | date = dataclasses.field(default=datetime.now())

    @classmethod
    def create(cls, value: datetime | date, _var_data: VarData | None = None):
        """Create a new instance of the class.

        Args:
            value: The value to set.

        Returns:
            LiteralDatetimeVar: The new instance of the class.
        """
        js_expr = f'"{value!s}"'
        return cls(
            _js_expr=js_expr,
            _var_type=type(value),
            _var_value=value,
            _var_data=_var_data,
        )


DATETIME_TYPES = (datetime, date, DateTimeVar)