wandb 0.19.9__py3-none-musllinux_1_2_aarch64.whl → 0.19.11__py3-none-musllinux_1_2_aarch64.whl

wandb/apis/public/utils.py

@@ -1,8 +1,11 @@
 import re
 from enum import Enum
-from typing import Optional
+from typing import Any, Dict, Iterable, Mapping, Optional, Set
 from urllib.parse import urlparse
 
+from wandb_gql import gql
+from wandb_graphql.language import ast, visitor
+
 from wandb._iterutils import one
 from wandb.sdk.artifacts._validators import is_artifact_registry_project
 from wandb.sdk.internal.internal_api import Api as InternalApi
@@ -102,3 +105,106 @@ def fetch_org_from_settings_or_entity(
         )
         organization = entity_org.display_name
     return organization
+
+
+class _GQLCompatRewriter(visitor.Visitor):
+    """GraphQL AST visitor to rewrite queries/mutations to be compatible with older server versions."""
+
+    omit_variables: Set[str]
+    omit_fragments: Set[str]
+    omit_fields: Set[str]
+    rename_fields: Dict[str, str]
+
+    def __init__(
+        self,
+        omit_variables: Optional[Iterable[str]] = None,
+        omit_fragments: Optional[Iterable[str]] = None,
+        omit_fields: Optional[Iterable[str]] = None,
+        rename_fields: Optional[Mapping[str, str]] = None,
+    ):
+        self.omit_variables = set(omit_variables or ())
+        self.omit_fragments = set(omit_fragments or ())
+        self.omit_fields = set(omit_fields or ())
+        self.rename_fields = dict(rename_fields or {})
+
+    def enter_VariableDefinition(self, node: ast.VariableDefinition, *_, **__) -> Any:  # noqa: N802
+        if node.variable.name.value in self.omit_variables:
+            return visitor.REMOVE
+        # return node
+
+    def enter_ObjectField(self, node: ast.ObjectField, *_, **__) -> Any:  # noqa: N802
+        # For context, note that e.g.:
+        #
+        #   {description: $description
+        #   ...}
+        #
+        # Is parsed as:
+        #
+        #   ObjectValue(fields=[
+        #     ObjectField(name=Name(value='description'), value=Variable(name=Name(value='description'))),
+        #   ...])
+        if (
+            isinstance(var := node.value, ast.Variable)
+            and var.name.value in self.omit_variables
+        ):
+            return visitor.REMOVE
+
+    def enter_Argument(self, node: ast.Argument, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_variables:
+            return visitor.REMOVE
+
+    def enter_FragmentDefinition(self, node: ast.FragmentDefinition, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_fragments:
+            return visitor.REMOVE
+
+    def enter_FragmentSpread(self, node: ast.FragmentSpread, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_fragments:
+            return visitor.REMOVE
+
+    def enter_Field(self, node: ast.Field, *_, **__) -> Any:  # noqa: N802
+        if node.name.value in self.omit_fields:
+            return visitor.REMOVE
+        if new_name := self.rename_fields.get(node.name.value):
+            node.name.value = new_name
+            return node
+
+    def leave_Field(self, node: ast.Field, *_, **__) -> Any:  # noqa: N802
+        # If the field had a selection set, but now it's empty, remove the field entirely
+        if (node.selection_set is not None) and (not node.selection_set.selections):
+            return visitor.REMOVE
+
+
+def gql_compat(
+    request_string: str,
+    omit_variables: Optional[Iterable[str]] = None,
+    omit_fragments: Optional[Iterable[str]] = None,
+    omit_fields: Optional[Iterable[str]] = None,
+    rename_fields: Optional[Mapping[str, str]] = None,
+) -> ast.Document:
+    """Rewrite a GraphQL request string to ensure compatibility with older server versions.
+
+    Args:
+        request_string (str): The GraphQL request string to rewrite.
+        omit_variables (Iterable[str] | None): Names of variables to remove from the request string.
+        omit_fragments (Iterable[str] | None): Names of fragments to remove from the request string.
+        omit_fields (Iterable[str] | None): Names of fields to remove from the request string.
+        rename_fields (Mapping[str, str] | None):
+            A mapping of fields to rename in the request string, given as `{old_name -> new_name}`.
+
+    Returns:
+        str: Modified GraphQL request string with fragments on omitted types removed.
+    """
+    # Parse the request into a GraphQL AST
+    doc = gql(request_string)
+
+    if not (omit_variables or omit_fragments or omit_fields or rename_fields):
+        return doc
+
+    # Visit the AST with our visitor to filter out unwanted fragments
+    rewriter = _GQLCompatRewriter(
+        omit_variables=omit_variables,
+        omit_fragments=omit_fragments,
+        omit_fields=omit_fields,
+        rename_fields=rename_fields,
+    )
+    return visitor.visit(doc, rewriter)

wandb/automations/__init__.py

@@ -0,0 +1,73 @@
+import wandb
+from wandb._pydantic import IS_PYDANTIC_V2
+
+from .actions import ActionType, DoNothing, SendNotification, SendWebhook
+from .automations import Automation, NewAutomation
+from .events import (
+    ArtifactEvent,
+    EventType,
+    MetricChangeFilter,
+    MetricThresholdFilter,
+    OnAddArtifactAlias,
+    OnCreateArtifact,
+    OnLinkArtifact,
+    OnRunMetric,
+    RunEvent,
+)
+from .integrations import Integration, SlackIntegration, WebhookIntegration
+from .scopes import ArtifactCollectionScope, ProjectScope, ScopeType
+
+# ----------------------------------------------------------------------------
+# WARNINGS on import
+if not IS_PYDANTIC_V2:
+    # Raises an error in Pydantic v1 environments, where the Automations API
+    # has not been tested and is unlikely to work as expected.
+    #
+    # Remove this when we either:
+    # - Drop support for Pydantic v1
+    # - Are able to implement (limited) Pydantic v1 support
+    raise ImportError(
+        "The W&B Automations API requires Pydantic v2. "
+        "We recommend upgrading `pydantic` to use this feature."
+    )
+
+else:
+    # If Pydantic v2 is available, we can use the full Automations API
+    # but communicate to users that the API is still experimental and
+    # may change rapidly.
+    wandb.termwarn(
+        "The W&B Automations API is experimental and the implementation is subject to change."
+        "Review the release notes before upgrading. We recommend pinning your "
+        f"package version to `{wandb.__package__}=={wandb.__version__}` to reduce the risk of disruption.",
+        repeat=False,
+    )
+# ----------------------------------------------------------------------------
+
+__all__ = [
+    # Scopes
+    "ScopeType",
+    "ArtifactCollectionScope",
+    "ProjectScope",
+    # Events
+    "EventType",
+    "OnAddArtifactAlias",
+    "OnCreateArtifact",
+    "OnLinkArtifact",
+    "OnRunMetric",
+    "ArtifactEvent",
+    "RunEvent",
+    "MetricThresholdFilter",
+    "MetricChangeFilter",
+    # Actions
+    "ActionType",
+    "SendNotification",
+    "SendWebhook",
+    "DoNothing",
+    # Automations
+    "Automation",
+    "NewAutomation",
+    # Integrations
+    "Integration",
+    "SlackIntegration",
+    "WebhookIntegration",
+]

wandb/automations/_filters/__init__.py

@@ -0,0 +1,40 @@
+from .expressions import FilterExpr, MongoLikeFilter
+from .operators import (
+    And,
+    Contains,
+    Eq,
+    Exists,
+    Gt,
+    Gte,
+    In,
+    Lt,
+    Lte,
+    Ne,
+    Nor,
+    Not,
+    NotIn,
+    Op,
+    Or,
+    Regex,
+)
+
+__all__ = [
+    "And",
+    "Or",
+    "Nor",
+    "Not",
+    "Op",
+    "Gt",
+    "Lt",
+    "Gte",
+    "Lte",
+    "Eq",
+    "Ne",
+    "In",
+    "NotIn",
+    "Contains",
+    "Exists",
+    "Regex",
+    "FilterExpr",
+    "MongoLikeFilter",
+]

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:  # type: ignore[override]
+        # 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:  # type: ignore[override]
+        # 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)