Lib/site-packages/altair/expr/core.py

from __future__ import annotations

import datetime as dt
from typing import TYPE_CHECKING, Any, Literal, Union

from altair.utils import SchemaBase

if TYPE_CHECKING:
    import sys

    from altair.vegalite.v5.schema._typing import Map, PrimitiveValue_T

    if sys.version_info >= (3, 10):
        from typing import TypeAlias
    else:
        from typing_extensions import TypeAlias


class DatumType:
    """An object to assist in building Vega-Lite Expressions."""

    def __repr__(self) -> str:
        return "datum"

    def __getattr__(self, attr) -> GetAttrExpression:
        if attr.startswith("__") and attr.endswith("__"):
            raise AttributeError(attr)
        return GetAttrExpression("datum", attr)

    def __getitem__(self, attr) -> GetItemExpression:
        return GetItemExpression("datum", attr)

    def __call__(self, datum, **kwargs) -> dict[str, Any]:
        """Specify a datum for use in an encoding."""
        return dict(datum=datum, **kwargs)


datum = DatumType()


def _js_repr(val) -> str:
    """Return a javascript-safe string representation of val."""
    if val is True:
        return "true"
    elif val is False:
        return "false"
    elif val is None:
        return "null"
    elif isinstance(val, OperatorMixin):
        return val._to_expr()
    elif isinstance(val, dt.date):
        return _from_date_datetime(val)
    else:
        return repr(val)


def _from_date_datetime(obj: dt.date | dt.datetime, /) -> str:
    """
    Parse native `datetime.(date|datetime)` into a `datetime expression`_ string.

    **Month is 0-based**

    .. _datetime expression:
        https://vega.github.io/vega/docs/expressions/#datetime
    """
    fn_name: Literal["datetime", "utc"] = "datetime"
    args: tuple[int, ...] = obj.year, obj.month - 1, obj.day
    if isinstance(obj, dt.datetime):
        if tzinfo := obj.tzinfo:
            if tzinfo is dt.timezone.utc:
                fn_name = "utc"
            else:
                msg = (
                    f"Unsupported timezone {tzinfo!r}.\n"
                    "Only `'UTC'` or naive (local) datetimes are permitted.\n"
                    "See https://altair-viz.github.io/user_guide/generated/core/altair.DateTime.html"
                )
                raise TypeError(msg)
        us = obj.microsecond
        ms = us if us == 0 else us // 1_000
        args = *args, obj.hour, obj.minute, obj.second, ms
    return FunctionExpression(fn_name, args)._to_expr()


# Designed to work with Expression and VariableParameter
class OperatorMixin:
    def _to_expr(self) -> str:
        return repr(self)

    def _from_expr(self, expr) -> Any:
        return expr

    def __add__(self, other):
        comp_value = BinaryExpression("+", self, other)
        return self._from_expr(comp_value)

    def __radd__(self, other):
        comp_value = BinaryExpression("+", other, self)
        return self._from_expr(comp_value)

    def __sub__(self, other):
        comp_value = BinaryExpression("-", self, other)
        return self._from_expr(comp_value)

    def __rsub__(self, other):
        comp_value = BinaryExpression("-", other, self)
        return self._from_expr(comp_value)

    def __mul__(self, other):
        comp_value = BinaryExpression("*", self, other)
        return self._from_expr(comp_value)

    def __rmul__(self, other):
        comp_value = BinaryExpression("*", other, self)
        return self._from_expr(comp_value)

    def __truediv__(self, other):
        comp_value = BinaryExpression("/", self, other)
        return self._from_expr(comp_value)

    def __rtruediv__(self, other):
        comp_value = BinaryExpression("/", other, self)
        return self._from_expr(comp_value)

    __div__ = __truediv__

    __rdiv__ = __rtruediv__

    def __mod__(self, other):
        comp_value = BinaryExpression("%", self, other)
        return self._from_expr(comp_value)

    def __rmod__(self, other):
        comp_value = BinaryExpression("%", other, self)
        return self._from_expr(comp_value)

    def __pow__(self, other):
        # "**" Javascript operator is not supported in all browsers
        comp_value = FunctionExpression("pow", (self, other))
        return self._from_expr(comp_value)

    def __rpow__(self, other):
        # "**" Javascript operator is not supported in all browsers
        comp_value = FunctionExpression("pow", (other, self))
        return self._from_expr(comp_value)

    def __neg__(self):
        comp_value = UnaryExpression("-", self)
        return self._from_expr(comp_value)

    def __pos__(self):
        comp_value = UnaryExpression("+", self)
        return self._from_expr(comp_value)

    # comparison operators

    def __eq__(self, other):
        comp_value = BinaryExpression("===", self, other)
        return self._from_expr(comp_value)

    def __ne__(self, other):
        comp_value = BinaryExpression("!==", self, other)
        return self._from_expr(comp_value)

    def __gt__(self, other):
        comp_value = BinaryExpression(">", self, other)
        return self._from_expr(comp_value)

    def __lt__(self, other):
        comp_value = BinaryExpression("<", self, other)
        return self._from_expr(comp_value)

    def __ge__(self, other):
        comp_value = BinaryExpression(">=", self, other)
        return self._from_expr(comp_value)

    def __le__(self, other):
        comp_value = BinaryExpression("<=", self, other)
        return self._from_expr(comp_value)

    def __abs__(self):
        comp_value = FunctionExpression("abs", (self,))
        return self._from_expr(comp_value)

    # logical operators

    def __and__(self, other):
        comp_value = BinaryExpression("&&", self, other)
        return self._from_expr(comp_value)

    def __rand__(self, other):
        comp_value = BinaryExpression("&&", other, self)
        return self._from_expr(comp_value)

    def __or__(self, other):
        comp_value = BinaryExpression("||", self, other)
        return self._from_expr(comp_value)

    def __ror__(self, other):
        comp_value = BinaryExpression("||", other, self)
        return self._from_expr(comp_value)

    def __invert__(self):
        comp_value = UnaryExpression("!", self)
        return self._from_expr(comp_value)


class Expression(OperatorMixin, SchemaBase):
    """
    Expression.

    Base object for enabling build-up of Javascript expressions using
    a Python syntax. Calling ``repr(obj)`` will return a Javascript
    representation of the object and the operations it encodes.
    """

    _schema = {"type": "string"}

    def to_dict(self, *args, **kwargs):
        return repr(self)

    def __setattr__(self, attr, val) -> None:
        # We don't need the setattr magic defined in SchemaBase
        return object.__setattr__(self, attr, val)

    # item access
    def __getitem__(self, val):
        return GetItemExpression(self, val)


class UnaryExpression(Expression):
    def __init__(self, op, val) -> None:
        super().__init__(op=op, val=val)

    def __repr__(self):
        return f"({self.op}{_js_repr(self.val)})"


class BinaryExpression(Expression):
    def __init__(self, op, lhs, rhs) -> None:
        super().__init__(op=op, lhs=lhs, rhs=rhs)

    def __repr__(self):
        return f"({_js_repr(self.lhs)} {self.op} {_js_repr(self.rhs)})"


class FunctionExpression(Expression):
    def __init__(self, name, args) -> None:
        super().__init__(name=name, args=args)

    def __repr__(self):
        args = ",".join(_js_repr(arg) for arg in self.args)
        return f"{self.name}({args})"


class ConstExpression(Expression):
    def __init__(self, name) -> None:
        super().__init__(name=name)

    def __repr__(self) -> str:
        return str(self.name)


class GetAttrExpression(Expression):
    def __init__(self, group, name) -> None:
        super().__init__(group=group, name=name)

    def __repr__(self):
        return f"{self.group}.{self.name}"


class GetItemExpression(Expression):
    def __init__(self, group, name) -> None:
        super().__init__(group=group, name=name)

    def __repr__(self) -> str:
        return f"{self.group}[{self.name!r}]"


IntoExpression: TypeAlias = Union[
    "PrimitiveValue_T", dt.date, dt.datetime, OperatorMixin, "Map"
]