File size: 2,310 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
"""Scopes in which a W&B Automation can be triggered."""
from __future__ import annotations
from typing import Literal, Union
from pydantic import BeforeValidator, Field
from typing_extensions import Annotated, TypeAlias, get_args
from wandb._pydantic import GQLBase
from wandb.automations._generated import (
ArtifactPortfolioScopeFields,
ArtifactSequenceScopeFields,
ProjectScopeFields,
)
from ._validators import LenientStrEnum, to_scope
# NOTE: Re-defined publicly with a more readable name for easier access
class ScopeType(LenientStrEnum):
"""The kind of scope that triggers an automation."""
PROJECT = "PROJECT"
ARTIFACT_COLLECTION = "ARTIFACT_COLLECTION"
class _BaseScope(GQLBase):
scope_type: Annotated[ScopeType, Field(frozen=True)]
class _ArtifactSequenceScope(_BaseScope, ArtifactSequenceScopeFields):
"""An automation scope defined by a specific `ArtifactSequence`."""
scope_type: Literal[ScopeType.ARTIFACT_COLLECTION] = ScopeType.ARTIFACT_COLLECTION
class _ArtifactPortfolioScope(_BaseScope, ArtifactPortfolioScopeFields):
"""An automation scope defined by a specific `ArtifactPortfolio` (e.g. a registry collection)."""
scope_type: Literal[ScopeType.ARTIFACT_COLLECTION] = ScopeType.ARTIFACT_COLLECTION
# for type annotations
ArtifactCollectionScope = Annotated[
Union[_ArtifactSequenceScope, _ArtifactPortfolioScope],
BeforeValidator(to_scope),
Field(discriminator="typename__"),
]
"""An automation scope defined by a specific `ArtifactCollection`."""
# for runtime type checks
ArtifactCollectionScopeTypes: tuple[type, ...] = get_args(
ArtifactCollectionScope.__origin__ # type: ignore[attr-defined]
)
class ProjectScope(_BaseScope, ProjectScopeFields):
"""An automation scope defined by a specific `Project`."""
scope_type: Literal[ScopeType.PROJECT] = ScopeType.PROJECT
# for type annotations
AutomationScope: TypeAlias = Annotated[
Union[_ArtifactSequenceScope, _ArtifactPortfolioScope, ProjectScope],
BeforeValidator(to_scope),
Field(discriminator="typename__"),
]
# for runtime type checks
AutomationScopeTypes: tuple[type, ...] = get_args(AutomationScope.__origin__) # type: ignore[attr-defined]
__all__ = [
"ScopeType",
"ArtifactCollectionScope",
"ProjectScope",
]
|