282 lines
8.1 KiB
Python
282 lines
8.1 KiB
Python
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"
|
|
]
|