wandb 0.21.2__py3-none-macosx_12_0_x86_64.whl

wandb/automations/_filters/expressions.py

@@ -0,0 +1,181 @@
+"""Pydantic-compatible representations of MongoDB expressions."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any, Union
+
+from pydantic import ConfigDict, model_serializer
+from typing_extensions import Self, TypeAlias, get_args
+
+from wandb._pydantic import CompatBaseModel, model_validator
+
+from .operators import (
+    Contains,
+    Eq,
+    Exists,
+    Gt,
+    Gte,
+    In,
+    Lt,
+    Lte,
+    Ne,
+    NotIn,
+    Op,
+    Regex,
+    RichReprResult,
+    Scalar,
+    ScalarTypes,
+    SupportsLogicalOpSyntax,
+)
+
+
+class FilterableField:
+    """A descriptor that can be used to define a "filterable" field on a class.
+
+    Internal helper to support syntactic sugar for defining event filters.
+    """
+
+    _python_name: str  #: The name of the field this descriptor was assigned to in the Python class.
+    _server_name: str | None  #: If set, the actual server-side field name to filter on.
+
+    def __init__(self, server_name: str | None = None):
+        self._server_name = server_name
+
+    def __set_name__(self, owner: type, name: str) -> None:
+        self._python_name = name
+
+    def __get__(self, obj: Any, objtype: type) -> Self:
+        # By default, if we didn't explicitly provide a backend name for
+        # filtering, assume the field has the same name in the backend as
+        # the python attribute.
+        return self
+
+    @property
+    def _name(self) -> str:
+        return self._server_name or self._python_name
+
+    def __str__(self) -> str:
+        return self._name
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self._name!r})"
+
+    # Methods to define filter expressions through chaining
+    def matches_regex(self, pattern: str) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Regex(regex_=pattern))
+
+    def contains(self, text: str) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Contains(contains_=text))
+
+    def exists(self, exists: bool = True) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Exists(exists_=exists))
+
+    def lt(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Lt(lt_=value))
+
+    def gt(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Gt(gt_=value))
+
+    def lte(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Lte(lte_=value))
+
+    def gte(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Gte(gte_=value))
+
+    def ne(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Ne(ne_=value))
+
+    def eq(self, value: Scalar) -> FilterExpr:
+        return FilterExpr(field=self._name, op=Eq(eq_=value))
+
+    def in_(self, values: Iterable[Scalar]) -> FilterExpr:
+        return FilterExpr(field=self._name, op=In(in_=values))
+
+    def not_in(self, values: Iterable[Scalar]) -> FilterExpr:
+        return FilterExpr(field=self._name, op=NotIn(nin_=values))
+
+    # Override the default behavior of comparison operators: <, >=, ==, etc
+    def __lt__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.lt(other)  # type: ignore[arg-type]
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __gt__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.gt(other)  # type: ignore[arg-type]
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __le__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.lte(other)  # type: ignore[arg-type]
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __ge__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.gte(other)  # type: ignore[arg-type]
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    # Operator behavior is intentionally overridden to allow defining
+    # filter expressions like `field == "value"`.  See similar overrides
+    # of built-in dunder methods in sqlalchemy, polars, pandas, numpy, etc.
+    #
+    # sqlalchemy example for illustrative purposes:
+    # https://github.com/sqlalchemy/sqlalchemy/blob/f21ae633486380a26dc0b67b70ae1c0efc6b4dc4/lib/sqlalchemy/orm/descriptor_props.py#L808-L812
+    def __eq__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.eq(other)  # type: ignore[arg-type]
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+    def __ne__(self, other: Any) -> FilterExpr:
+        if isinstance(other, ScalarTypes):
+            return self.ne(other)  # type: ignore[arg-type]
+        raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}")
+
+
+# ------------------------------------------------------------------------------
+class FilterExpr(CompatBaseModel, SupportsLogicalOpSyntax):
+    """A MongoDB filter expression on a specific field."""
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+    field: str
+    op: Op
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self.field!s}: {self.op!r})"
+
+    def __rich_repr__(self) -> RichReprResult:
+        # https://rich.readthedocs.io/en/stable/pretty.html
+        yield self.field, self.op
+
+    @model_validator(mode="before")
+    @classmethod
+    def _validate(cls, data: Any) -> Any:
+        """Parse a MongoDB dict representation of the filter expression."""
+        if (
+            isinstance(data, dict)
+            and len(data) == 1
+            and not any(key.startswith("$") for key in data)
+        ):
+            # This looks like a MongoDB filter dict.  E.g.:
+            # - in:  `{"display_name": {"$contains": "my-run"}}`
+            # - out: `FilterExpr(field="display_name", op=Contains(contains_="my-run"))`
+            ((field, op),) = data.items()
+            return {"field": field, "op": op}
+        return data
+
+    @model_serializer(mode="plain")
+    def _serialize(self) -> dict[str, Any]:
+        """Return a MongoDB dict representation of the expression."""
+        from pydantic_core import to_jsonable_python  # Only valid in pydantic v2
+
+        return {self.field: to_jsonable_python(self.op, by_alias=True, round_trip=True)}
+
+
+# for type annotations
+MongoLikeFilter: TypeAlias = Union[Op, FilterExpr]
+# for runtime type checks
+MongoLikeFilterTypes: tuple[type, ...] = get_args(MongoLikeFilter)

wandb/automations/_filters/operators.py

@@ -0,0 +1,258 @@
+"""Types that represent operators in MongoDB filter expressions."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, Tuple, TypeVar, Union
+
+from pydantic import ConfigDict, Field, StrictBool, StrictFloat, StrictInt, StrictStr
+from typing_extensions import TypeAlias, get_args
+
+from wandb._pydantic import GQLBase
+
+# for type annotations
+Scalar = Union[StrictStr, StrictInt, StrictFloat, StrictBool]
+# for runtime type checks
+ScalarTypes: tuple[type, ...] = tuple(t.__origin__ for t in get_args(Scalar))
+
+# See: https://rich.readthedocs.io/en/stable/pretty.html#rich-repr-protocol
+RichReprResult: TypeAlias = Iterable[
+    Union[
+        Any,
+        Tuple[Any],
+        Tuple[str, Any],
+        Tuple[str, Any, Any],
+    ]
+]
+
+T = TypeVar("T")
+TupleOf: TypeAlias = Tuple[T, ...]
+
+
+# NOTE: Wherever class descriptions that are not docstrings, this is deliberate.
+# This is done to ensure the descriptions are omitted from generated API docs.
+
+
+# Mixin to support syntactic sugar for MongoDB expressions with (bitwise) logical operators,
+# e.g. `a | b` -> `{"$or": [a, b]}` or `~a` -> `{"$not": a}`.
+class SupportsLogicalOpSyntax:
+    def __or__(self, other: Any) -> Or:
+        """Syntactic sugar for: `a | b` -> `Or(a, b)`."""
+        return Or(or_=[self, other])
+
+    def __and__(self, other: Any) -> And:
+        """Syntactic sugar for: `a & b` -> `And(a, b)`."""
+        from .expressions import FilterExpr
+
+        if isinstance(other, (BaseOp, FilterExpr)):
+            return And(and_=[self, other])
+        return NotImplemented
+
+    def __invert__(self) -> Not:
+        """Syntactic sugar for: `~a` -> `Not(a)`."""
+        return Not(not_=self)
+
+
+# Base class for parsed MongoDB filter/query operators, e.g. `{"$and": [...]}`.
+class BaseOp(GQLBase, SupportsLogicalOpSyntax):
+    model_config = ConfigDict(
+        extra="forbid",
+        frozen=True,  # Make pseudo-immutable for easier comparison and hashing
+    )
+
+    def __repr__(self) -> str:
+        # Display operand as a positional arg
+        values_repr = ", ".join(map(repr, self.model_dump().values()))
+        return f"{type(self).__name__}({values_repr})"
+
+    def __rich_repr__(self) -> RichReprResult:
+        # Display field values as positional args:
+        # https://rich.readthedocs.io/en/stable/pretty.html
+        yield from ((None, v) for v in self.model_dump().values())
+
+
+# Logical operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/and/
+# https://www.mongodb.com/docs/manual/reference/operator/query/or/
+# https://www.mongodb.com/docs/manual/reference/operator/query/nor/
+# https://www.mongodb.com/docs/manual/reference/operator/query/not/
+class And(BaseOp):
+    and_: TupleOf[Any] = Field(default=(), alias="$and")
+
+
+class Or(BaseOp):
+    or_: TupleOf[Any] = Field(default=(), alias="$or")
+
+    def __invert__(self) -> Nor:
+        """Syntactic sugar for: `~Or(a, b)` -> `Nor(a, b)`."""
+        return Nor(nor_=self.or_)
+
+
+class Nor(BaseOp):
+    nor_: TupleOf[Any] = Field(default=(), alias="$nor")
+
+    def __invert__(self) -> Or:
+        """Syntactic sugar for: `~Nor(a, b)` -> `Or(a, b)`."""
+        return Or(or_=self.nor_)
+
+
+class Not(BaseOp):
+    not_: Any = Field(alias="$not")
+
+    def __invert__(self) -> Any:
+        """Syntactic sugar for: `~Not(a)` -> `a`."""
+        return self.not_
+
+
+# Comparison operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/lt/
+# https://www.mongodb.com/docs/manual/reference/operator/query/gt/
+# https://www.mongodb.com/docs/manual/reference/operator/query/lte/
+# https://www.mongodb.com/docs/manual/reference/operator/query/gte/
+# https://www.mongodb.com/docs/manual/reference/operator/query/eq/
+# https://www.mongodb.com/docs/manual/reference/operator/query/ne/
+# https://www.mongodb.com/docs/manual/reference/operator/query/in/
+# https://www.mongodb.com/docs/manual/reference/operator/query/nin/
+class Lt(BaseOp):
+    lt_: Scalar = Field(alias="$lt")
+
+    def __invert__(self) -> Gte:
+        """Syntactic sugar for: `~Lt(a)` -> `Gte(a)`."""
+        return Gte(gte_=self.lt_)
+
+
+class Gt(BaseOp):
+    gt_: Scalar = Field(alias="$gt")
+
+    def __invert__(self) -> Lte:
+        """Syntactic sugar for: `~Gt(a)` -> `Lte(a)`."""
+        return Lte(lte_=self.gt_)
+
+
+class Lte(BaseOp):
+    lte_: Scalar = Field(alias="$lte")
+
+    def __invert__(self) -> Gt:
+        """Syntactic sugar for: `~Lte(a)` -> `Gt(a)`."""
+        return Gt(gt_=self.lte_)
+
+
+class Gte(BaseOp):
+    gte_: Scalar = Field(alias="$gte")
+
+    def __invert__(self) -> Lt:
+        """Syntactic sugar for: `~Gte(a)` -> `Lt(a)`."""
+        return Lt(lt_=self.gte_)
+
+
+class Eq(BaseOp):
+    eq_: Scalar = Field(alias="$eq")
+
+    def __invert__(self) -> Ne:
+        """Syntactic sugar for: `~Eq(a)` -> `Ne(a)`."""
+        return Ne(ne_=self.eq_)
+
+
+class Ne(BaseOp):
+    ne_: Scalar = Field(alias="$ne")
+
+    def __invert__(self) -> Eq:
+        """Syntactic sugar for: `~Ne(a)` -> `Eq(a)`."""
+        return Eq(eq_=self.ne_)
+
+
+class In(BaseOp):
+    in_: TupleOf[Scalar] = Field(default=(), alias="$in")
+
+    def __invert__(self) -> NotIn:
+        """Syntactic sugar for: `~In(a)` -> `NotIn(a)`."""
+        return NotIn(nin_=self.in_)
+
+
+class NotIn(BaseOp):
+    nin_: TupleOf[Scalar] = Field(default=(), alias="$nin")
+
+    def __invert__(self) -> In:
+        """Syntactic sugar for: `~NotIn(a)` -> `In(a)`."""
+        return In(in_=self.nin_)
+
+
+# Element operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/exists/
+class Exists(BaseOp):
+    exists_: bool = Field(alias="$exists")
+
+
+# Evaluation operator(s)
+# https://www.mongodb.com/docs/manual/reference/operator/query/regex/
+#
+# Note: "$contains" is NOT a formal MongoDB operator, but the backend recognizes and
+# executes it as a substring-match filter.
+class Regex(BaseOp):
+    regex_: str = Field(alias="$regex")  #: The regex expression to match against.
+
+
+class Contains(BaseOp):
+    contains_: str = Field(alias="$contains")  #: The substring to match against.
+
+
+And.model_rebuild()
+Or.model_rebuild()
+Not.model_rebuild()
+Lt.model_rebuild()
+Gt.model_rebuild()
+Lte.model_rebuild()
+Gte.model_rebuild()
+Eq.model_rebuild()
+Ne.model_rebuild()
+In.model_rebuild()
+NotIn.model_rebuild()
+Exists.model_rebuild()
+Regex.model_rebuild()
+Contains.model_rebuild()
+
+
+# ------------------------------------------------------------------------------
+# Convenience helpers, constants, and utils for supported MongoDB operators
+# ------------------------------------------------------------------------------
+KEY_TO_OP: dict[str, type[BaseOp]] = {
+    "$and": And,
+    "$or": Or,
+    "$nor": Nor,
+    "$not": Not,
+    "$lt": Lt,
+    "$gt": Gt,
+    "$lte": Lte,
+    "$gte": Gte,
+    "$eq": Eq,
+    "$ne": Ne,
+    "$in": In,
+    "$nin": NotIn,
+    "$exists": Exists,
+    "$regex": Regex,
+    "$contains": Contains,
+}
+
+
+KnownOp = Union[
+    And,
+    Or,
+    Nor,
+    Not,
+    Lt,
+    Gt,
+    Lte,
+    Gte,
+    Eq,
+    Ne,
+    In,
+    NotIn,
+    Exists,
+    Regex,
+    Contains,
+]
+UnknownOp = Dict[str, Any]
+
+# for type annotations
+Op = Union[KnownOp, UnknownOp]
+# for runtime type checks
+OpTypes: tuple[type, ...] = (*get_args(KnownOp), dict)