wandb 0.19.10__py3-none-macosx_11_0_arm64.whl → 0.19.11__py3-none-macosx_11_0_arm64.whl

wandb/automations/_generated/update_automation.py

@@ -0,0 +1,17 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/automations/
+
+from __future__ import annotations
+
+from typing import Optional
+
+from wandb._pydantic import GQLBase
+
+from .fragments import UpdateAutomationResult
+
+
+class UpdateAutomation(GQLBase):
+    result: Optional[UpdateAutomationResult]
+
+
+UpdateAutomation.model_rebuild()

wandb/automations/_utils.py

@@ -2,30 +2,63 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional, Protocol, TypedDict
+from typing import Any, Collection, Final, Optional, Protocol, TypedDict
 
-from wandb._pydantic import to_json
-from wandb.automations.events import InputEvent
-from wandb.automations.scopes import InputScope
+from pydantic import Field
+from typing_extensions import Annotated, Self, Unpack
 
+from wandb._pydantic import GQLBase, GQLId, computed_field, model_validator, to_json
+
+from ._filters import MongoLikeFilter
 from ._generated import (
     CreateFilterTriggerInput,
+    QueueJobActionInput,
     TriggeredActionConfig,
     UpdateFilterTriggerInput,
 )
+from ._validators import to_input_action
 from .actions import (
     ActionType,
     DoNothing,
-    DoNotification,
-    DoWebhook,
     InputAction,
     SavedAction,
+    SendNotification,
+    SendWebhook,
+)
+from .automations import Automation, NewAutomation
+from .events import EventType, InputEvent, RunMetricFilter, _WrappedSavedEventFilter
+from .scopes import AutomationScope, ScopeType
+
+EXCLUDED_INPUT_EVENTS: Final[Collection[EventType]] = frozenset(
+    {
+        EventType.UPDATE_ARTIFACT_ALIAS,
+    }
 )
+"""Event types that should not be assigned when creating/updating automations."""
 
-if TYPE_CHECKING:
-    from typing_extensions import Unpack
+EXCLUDED_INPUT_ACTIONS: Final[Collection[ActionType]] = frozenset(
+    {
+        ActionType.QUEUE_JOB,
+    }
+)
+"""Action types that should not be assigned when creating/updating automations."""
+
+ALWAYS_SUPPORTED_EVENTS: Final[Collection[EventType]] = frozenset(
+    {
+        EventType.CREATE_ARTIFACT,
+        EventType.LINK_ARTIFACT,
+        EventType.ADD_ARTIFACT_ALIAS,
+    }
+)
+"""Event types that we can safely assume all contemporary server versions support."""
 
-    from .automations import Automation, NewAutomation
+ALWAYS_SUPPORTED_ACTIONS: Final[Collection[ActionType]] = frozenset(
+    {
+        ActionType.NOTIFICATION,
+        ActionType.GENERIC_WEBHOOK,
+    }
+)
+"""Action types that we can safely assume all contemporary server versions support."""
 
 
 class HasId(Protocol):
@@ -36,88 +69,169 @@ def extract_id(obj: HasId | str) -> str:
     return obj.id if hasattr(obj, "id") else obj
 
 
-# NOTE: `QueueJobActionInput` for defining a Launch job is deprecated,
-# so we deliberately don't currently expose it in the API for *creating* automations.
-_ACTION_CONFIG_KEYS: dict[ActionType, str] = {
+# ---------------------------------------------------------------------------
+ACTION_CONFIG_KEYS: dict[ActionType, str] = {
     ActionType.NOTIFICATION: "notification_action_input",
     ActionType.GENERIC_WEBHOOK: "generic_webhook_action_input",
     ActionType.NO_OP: "no_op_action_input",
+    ActionType.QUEUE_JOB: "queue_job_action_input",
 }
 
 
 class InputActionConfig(TriggeredActionConfig):
     """A `TriggeredActionConfig` that prepares the action config for saving an automation."""
 
-    notification_action_input: Optional[DoNotification] = None
-    generic_webhook_action_input: Optional[DoWebhook] = None
+    # NOTE: `QueueJobActionInput` for defining a Launch job is deprecated,
+    # so while it's allowed here to update EXISTING mutations, we don't
+    # currently expose it through the public API to create NEW automations.
+    queue_job_action_input: Optional[QueueJobActionInput] = None
+
+    notification_action_input: Optional[SendNotification] = None
+    generic_webhook_action_input: Optional[SendWebhook] = None
     no_op_action_input: Optional[DoNothing] = None
 
 
-def prepare_action_input(obj: InputAction | SavedAction) -> InputActionConfig:
-    """Return a `TriggeredActionConfig` as required in the input schema of CreateFilterTriggerInput."""
-    key = _ACTION_CONFIG_KEYS[obj.action_type]
-    # Delegate the validators for each ActionInput type to handle custom logic
-    # for converting from saved action types to an input action type.
-    return InputActionConfig.model_validate({key: obj})
+def prepare_action_config_input(obj: SavedAction | InputAction) -> dict[str, Any]:
+    """Prepare the `TriggeredActionConfig` input, nesting the action input inside the appropriate key.
+
+    This is necessary to conform to the schemas for:
+    - CreateFilterTriggerInput
+    - UpdateFilterTriggerInput
+    """
+    # Delegate to inner validators to convert SavedAction -> InputAction types, if needed.
+    obj = to_input_action(obj)
+    return InputActionConfig(**{ACTION_CONFIG_KEYS[obj.action_type]: obj}).model_dump()
+
+
+def prepare_event_filter_input(
+    obj: _WrappedSavedEventFilter | MongoLikeFilter | RunMetricFilter,
+) -> str:
+    """Prepare the `EventFilter` input, unnesting the filter if needed and serializing to JSON.
+
+    This is necessary to conform to the schemas for:
+    - CreateFilterTriggerInput
+    - UpdateFilterTriggerInput
+    """
+    # Input event filters are nested one level deeper than saved event filters.
+    # Note that this is NOT the case for run/run metric filters.
+    #
+    # Yes, this is confusing.  It's also necessary to conform to under-the-hood
+    # schemas and logic in the backend.
+    filter_to_serialize = (
+        obj.filter if isinstance(obj, _WrappedSavedEventFilter) else obj
+    )
+    return to_json(filter_to_serialize)
 
 
-class AutomationParams(TypedDict, total=False):
+class WriteAutomationsKwargs(TypedDict, total=False):
     """Keyword arguments that can be passed to create or update an automation."""
 
     name: str
     description: str
     enabled: bool
-
-    scope: InputScope
+    scope: AutomationScope
     event: InputEvent
     action: InputAction
 
 
-def prepare_create_input(
-    obj: NewAutomation, **updates: Unpack[AutomationParams]
-) -> CreateFilterTriggerInput:
-    """Prepares the payload to create an automation in a GraphQL request."""
-    from .automations import PreparedAutomation
-
-    # Apply any updates to the properties of the automation
-    updated = obj.model_copy(update=updates)
-    prepared = PreparedAutomation.model_validate(updated)
-
-    # Prepare the input as required for the GraphQL request
-    return CreateFilterTriggerInput(
-        name=prepared.name,
-        description=prepared.description,
-        enabled=prepared.enabled,
-        scope_type=prepared.scope.scope_type,
-        scope_id=prepared.scope.id,
-        triggering_event_type=prepared.event.event_type,
-        event_filter=to_json(prepared.event.filter),
-        triggered_action_type=prepared.action.action_type,
-        triggered_action_config=prepare_action_input(prepared.action).model_dump(),
-    )
+class ValidatedCreateInput(GQLBase, extra="forbid", frozen=True):
+    """Validated automation parameters, prepared for creating a new automation.
 
+    Note: Users should never need to instantiate this class directly.
+    """
 
-def prepare_update_input(
-    obj: Automation, **updates: Unpack[AutomationParams]
+    name: str
+    description: Optional[str] = None
+    enabled: bool = True
+
+    # ------------------------------------------------------------------------------
+    # Set on instantiation, but used to derive other fields and deliberately
+    # EXCLUDED from the final GraphQL request vars
+    event: Annotated[InputEvent, Field(exclude=True)]
+    action: Annotated[InputAction, Field(exclude=True)]
+
+    # ------------------------------------------------------------------------------
+    # Derived fields to match the input schemas
+    @computed_field
+    def scope_type(self) -> ScopeType:
+        return self.event.scope.scope_type
+
+    @computed_field
+    def scope_id(self) -> GQLId:
+        return self.event.scope.id
+
+    @computed_field
+    def triggering_event_type(self) -> EventType:
+        return self.event.event_type
+
+    @computed_field
+    def event_filter(self) -> str:
+        return prepare_event_filter_input(self.event.filter)
+
+    @computed_field
+    def triggered_action_type(self) -> ActionType:
+        return self.action.action_type
+
+    @computed_field
+    def triggered_action_config(self) -> dict[str, Any]:
+        return prepare_action_config_input(self.action)
+
+    # ------------------------------------------------------------------------------
+    # Custom validation
+    @model_validator(mode="after")
+    def _forbid_legacy_event_types(self) -> Self:
+        if (type_ := self.event.event_type) in EXCLUDED_INPUT_EVENTS:
+            raise ValueError(f"{type_!r} events cannot be assigned to automations.")
+        return self
+
+    @model_validator(mode="after")
+    def _forbid_legacy_action_types(self) -> Self:
+        if (type_ := self.action.action_type) in EXCLUDED_INPUT_ACTIONS:
+            raise ValueError(f"{type_!r} actions cannot be assigned to automations.")
+        return self
+
+
+def prepare_to_create(
+    obj: NewAutomation | None = None,
+    /,
+    **kwargs: Unpack[WriteAutomationsKwargs],
+) -> CreateFilterTriggerInput:
+    """Prepares the payload to create an automation in a GraphQL request."""
+    # Validate all input variables, and prepare as expected by the GraphQL request.
+    # - if an object is provided, override its fields with any keyword args
+    # - otherwise, instantiate from the keyword args
+
+    # NOTE: `exclude_none=True` drops fields that are still `None`.
+    #
+    # This assumes that `None` is good enough for now as a sentinel
+    # "unset" value.  If this proves insufficient, revisit in the future,
+    # as it should be reasonably easy to implement a custom sentinel
+    # type later on.
+    obj_dict = {**obj.model_dump(exclude_none=True), **kwargs} if obj else kwargs
+    validated = ValidatedCreateInput(**obj_dict)
+    return CreateFilterTriggerInput.model_validate(validated)
+
+
+def prepare_to_update(
+    obj: Automation | None = None,
+    /,
+    **kwargs: Unpack[WriteAutomationsKwargs],
 ) -> UpdateFilterTriggerInput:
     """Prepares the payload to update an automation in a GraphQL request."""
-    from .events import SavedEventFilter
-
-    updated = obj.model_copy(update=updates)
+    # Validate all values:
+    # - if an object is provided, override its fields with any keyword args
+    # - otherwise, instantiate from the keyword args
+    v_obj = Automation(**{**dict(obj or {}), **kwargs})
 
     return UpdateFilterTriggerInput(
-        id=updated.id,
-        name=updated.name,
-        description=updated.description,
-        enabled=updated.enabled,
-        scope_type=updated.scope.scope_type,
-        scope_id=updated.scope.id,
-        triggering_event_type=updated.event.event_type,
-        event_filter=to_json(
-            updated.event.filter.filter
-            if isinstance(updated.event.filter, SavedEventFilter)
-            else updated.event.filter
-        ),
-        triggered_action_type=updated.action.action_type,
-        triggered_action_config=prepare_action_input(updated.action).model_dump(),
+        id=v_obj.id,
+        name=v_obj.name,
+        description=v_obj.description,
+        enabled=v_obj.enabled,
+        scope_type=v_obj.scope.scope_type,
+        scope_id=v_obj.scope.id,
+        triggering_event_type=v_obj.event.event_type,
+        event_filter=prepare_event_filter_input(v_obj.event.filter),
+        triggered_action_type=v_obj.action.action_type,
+        triggered_action_config=prepare_action_config_input(v_obj.action),
     )

wandb/automations/_validators.py

@@ -1,16 +1,53 @@
 from __future__ import annotations
 
+from enum import Enum
 from functools import singledispatch
 from itertools import chain
 from typing import Any, TypeVar
 
+from pydantic_core import PydanticUseDefault
+
 from ._filters import And, FilterExpr, In, Nor, Not, NotIn, Op, Or
 
 T = TypeVar("T")
 
 
-def validate_scope(v: Any) -> Any:
-    """Convert a familiar wandb `Project` or `ArtifactCollection` object to an automation scope."""
+class LenientStrEnum(str, Enum):
+    """A string enum allowing for case-insensitive lookups by value.
+
+    May include other internal customizations if needed.
+
+    Note: This is a bespoke, internal implementation and NOT intended as a
+    backport of `enum.StrEnum` from Python 3.11+.
+    """
+
+    def __repr__(self) -> str:
+        return self.name
+
+    @classmethod
+    def _missing_(cls, value: object) -> Any:
+        # Accept case-insensitive enum values
+        if isinstance(value, str):
+            v = value.lower()
+            return next((e for e in cls if e.value.lower() == v), None)
+        return None
+
+
+def default_if_none(v: Any) -> Any:
+    """A before-validator validator that coerces `None` to the default field value instead."""
+    # https://docs.pydantic.dev/2.11/api/pydantic_core/#pydantic_core.PydanticUseDefault
+    if v is None:
+        raise PydanticUseDefault
+    return v
+
+
+def upper_if_str(v: Any) -> Any:
+    return v.strip().upper() if isinstance(v, str) else v
+
+
+# ----------------------------------------------------------------------------
+def to_scope(v: Any) -> Any:
+    """Convert eligible objects, including pre-existing `wandb` types, to an automation scope."""
     from wandb.apis.public import ArtifactCollection, Project
 
     from .scopes import ProjectScope, _ArtifactPortfolioScope, _ArtifactSequenceScope
@@ -23,6 +60,61 @@ def validate_scope(v: Any) -> Any:
     return v
 
 
+def to_saved_action(v: Any) -> Any:
+    """If necessary (and possible), convert the object to a saved action."""
+    from .actions import (
+        DoNothing,
+        SavedNoOpAction,
+        SavedNotificationAction,
+        SavedWebhookAction,
+        SendNotification,
+        SendWebhook,
+    )
+
+    if isinstance(v, SendNotification):
+        return SavedNotificationAction(
+            integration={"id": v.integration_id},
+            **v.model_dump(exclude={"integration_id"}),
+        )
+    if isinstance(v, SendWebhook):
+        return SavedWebhookAction(
+            integration={"id": v.integration_id},
+            **v.model_dump(exclude={"integration_id"}),
+        )
+    if isinstance(v, DoNothing):
+        return SavedNoOpAction.model_validate(v)
+
+    return v
+
+
+def to_input_action(v: Any) -> Any:
+    """If necessary (and possible), convert the object to an input action."""
+    from .actions import (
+        DoNothing,
+        SavedNoOpAction,
+        SavedNotificationAction,
+        SavedWebhookAction,
+        SendNotification,
+        SendWebhook,
+    )
+
+    if isinstance(v, SavedNotificationAction):
+        return SendNotification(
+            integration_id=v.integration.id,
+            **v.model_dump(exclude={"integration"}),
+        )
+    if isinstance(v, SavedWebhookAction):
+        return SendWebhook(
+            integration_id=v.integration.id,
+            **v.model_dump(exclude={"integration"}),
+        )
+    if isinstance(v, SavedNoOpAction):
+        return DoNothing.model_validate(v)
+
+    return v
+
+
+# ----------------------------------------------------------------------------
 @singledispatch
 def simplify_op(op: Op | FilterExpr) -> Op | FilterExpr:
     """Simplify a MongoDB filter by removing and unnesting redundant operators."""