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

wandb/apis/public/automations.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from itertools import chain
+from typing import TYPE_CHECKING, Any, Iterable, Mapping
+
+from pydantic import ValidationError
+from typing_extensions import override
+from wandb_graphql.language.ast import Document
+
+from wandb.apis.paginator import Paginator, _Client
+
+if TYPE_CHECKING:
+    from wandb.automations import Automation
+    from wandb.automations._generated import ProjectConnectionFields
+
+
+class Automations(Paginator["Automation"]):
+    last_response: ProjectConnectionFields | None
+    _query: Document
+
+    def __init__(
+        self,
+        client: _Client,
+        variables: Mapping[str, Any],
+        per_page: int = 50,
+        _query: Document | None = None,
+    ):
+        super().__init__(client, variables, per_page=per_page)
+        if _query is None:
+            raise RuntimeError(f"Query required for {type(self).__qualname__}")
+        self._query = _query
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more items to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch the raw response data for the current page."""
+        from wandb.automations._generated import ProjectConnectionFields
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["searchScope"]["projects"]
+            self.last_response = ProjectConnectionFields.model_validate(page_data)
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> Iterable[Automation]:
+        """Parse the page data into a list of objects."""
+        from wandb.automations import Automation
+
+        page = self.last_response
+        return [
+            Automation.model_validate(obj)
+            for obj in chain.from_iterable(edge.node.triggers for edge in page.edges)
+        ]

wandb/apis/public/integrations.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Iterable
+
+from pydantic import ValidationError
+from typing_extensions import override
+from wandb_gql import gql
+from wandb_graphql.language.ast import Document
+
+from wandb.apis.paginator import Paginator
+
+if TYPE_CHECKING:
+    from wandb.apis.paginator import _Client
+    from wandb.automations import Integration, SlackIntegration, WebhookIntegration
+    from wandb.automations._generated import (
+        GenericWebhookIntegrationConnectionFields,
+        IntegrationConnectionFields,
+        SlackIntegrationConnectionFields,
+    )
+
+
+class Integrations(Paginator["Integration"]):
+    last_response: IntegrationConnectionFields | None
+    _query: Document
+
+    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
+        from wandb.automations._generated import INTEGRATIONS_BY_ENTITY_GQL
+
+        super().__init__(client, variables, per_page=per_page)
+        # All integrations for entity
+        self._query = gql(INTEGRATIONS_BY_ENTITY_GQL)
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more Integrations to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch and parse the response data for the current page."""
+        from wandb.automations._generated import IntegrationConnectionFields
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["entity"]["integrations"]
+            self.last_response = IntegrationConnectionFields.model_validate(page_data)
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> Iterable[Integration]:
+        """Parse the page data into a list of integrations."""
+        from wandb.automations.integrations import _IntegrationEdge
+
+        page = self.last_response
+        return [_IntegrationEdge.model_validate(edge).node for edge in page.edges]
+
+
+class WebhookIntegrations(Paginator["WebhookIntegration"]):
+    last_response: GenericWebhookIntegrationConnectionFields | None
+    _query: Document
+
+    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
+        from wandb.automations._generated import (
+            GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL,
+        )
+
+        super().__init__(client, variables, per_page=per_page)
+        # Webhook integrations for entity
+        self._query = gql(GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL)
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more webhook integrations to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch and parse the response data for the current page."""
+        from wandb.automations._generated import (
+            GenericWebhookIntegrationConnectionFields,
+        )
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["entity"]["integrations"]
+            self.last_response = (
+                GenericWebhookIntegrationConnectionFields.model_validate(page_data)
+            )
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> Iterable[WebhookIntegration]:
+        """Parse the page data into a list of webhook integrations."""
+        from wandb.automations import WebhookIntegration
+
+        typename = "GenericWebhookIntegration"
+        return [
+            # Filter on typename__ needed because the GQL response still
+            # includes all integration types
+            WebhookIntegration.model_validate(node)
+            for edge in self.last_response.edges
+            if (node := edge.node) and (node.typename__ == typename)
+        ]
+
+
+class SlackIntegrations(Paginator["SlackIntegration"]):
+    last_response: SlackIntegrationConnectionFields | None
+    _query: Document
+
+    def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50):
+        from wandb.automations._generated import SLACK_INTEGRATIONS_BY_ENTITY_GQL
+
+        super().__init__(client, variables, per_page=per_page)
+        # Slack integrations for entity
+        self._query = gql(SLACK_INTEGRATIONS_BY_ENTITY_GQL)
+
+    @property
+    def more(self) -> bool:
+        """Whether there are more Slack integrations to fetch."""
+        if self.last_response is None:
+            return True
+        return self.last_response.page_info.has_next_page
+
+    @property
+    def cursor(self) -> str | None:
+        """The start cursor to use for the next page."""
+        if self.last_response is None:
+            return None
+        return self.last_response.page_info.end_cursor
+
+    @override
+    def _update_response(self) -> None:
+        """Fetch and parse the response data for the current page."""
+        from wandb.automations._generated import SlackIntegrationConnectionFields
+
+        data: dict[str, Any] = self.client.execute(
+            self._query, variable_values=self.variables
+        )
+        try:
+            page_data = data["entity"]["integrations"]
+            self.last_response = SlackIntegrationConnectionFields.model_validate(
+                page_data
+            )
+        except (LookupError, AttributeError, ValidationError) as e:
+            raise ValueError("Unexpected response data") from e
+
+    def convert_objects(self) -> Iterable[SlackIntegration]:
+        """Parse the page data into a list of Slack integrations."""
+        from wandb.automations import SlackIntegration
+
+        typename = "SlackIntegration"
+        return [
+            # Filter on typename__ needed because the GQL response still
+            # includes all integration types
+            SlackIntegration.model_validate(node)
+            for edge in self.last_response.edges
+            if (node := edge.node) and (node.typename__ == typename)
+        ]

wandb/apis/public/projects.py

@@ -1,5 +1,8 @@
 """Public API: projects."""
 
+from contextlib import suppress
+
+from requests import HTTPError
 from wandb_gql import gql
 
 from wandb.apis import public
@@ -153,3 +156,29 @@ class Project(Attrs):
             )
             for e in ret["project"]["sweeps"]["edges"]
         ]
+
+    _PROJECT_ID = gql(
+        """
+        query ProjectID($projectName: String!, $entityName: String!) {
+            project(name: $projectName, entityName: $entityName) {
+                id
+            }
+        }
+        """
+    )
+
+    @property
+    def id(self) -> str:
+        # This is a workaround to ensure that the project ID can be retrieved
+        # on demand, as it generally is not set or fetched on instantiation.
+        # This is necessary if using this project as the scope of a new Automation.
+        with suppress(LookupError):
+            return self._attrs["id"]
+
+        variable_values = {"projectName": self.name, "entityName": self.entity}
+        try:
+            data = self.client.execute(self._PROJECT_ID, variable_values)
+            self._attrs["id"] = data["project"]["id"]
+            return self._attrs["id"]
+        except (HTTPError, LookupError, TypeError) as e:
+            raise ValueError(f"Unable to fetch project ID: {variable_values!r}") from e

wandb/apis/public/registries/_freezable_list.py

@@ -0,0 +1,179 @@
+from itertools import chain
+from typing import (
+    Any,
+    Iterable,
+    Iterator,
+    List,
+    MutableSequence,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+    final,
+    overload,
+)
+
+T = TypeVar("T")
+
+
+@final
+class FreezableList(MutableSequence[T]):
+    """A list-like container type that only allows adding new items.
+
+    It tracks "saved" (immutable) and "draft" (mutable) items.
+    Items can be added, inserted, and removed while in draft state, but once frozen,
+    they become immutable. Unlike a set, duplicate items are allowed in the draft
+    state but duplicates already present in the saved state cannot be added.
+    Any initial items passed to the constructor are saved.
+    """
+
+    def __init__(self, iterable: Union[Iterable[T], None] = None, /) -> None:
+        self._frozen: Tuple[T, ...] = tuple(iterable or ())
+        self._draft: List[T] = []
+
+    def append(self, value: T) -> None:
+        """Append an item to the draft list. No duplicates are allowed."""
+        if (value in self._frozen) or (value in self._draft):
+            return
+        self._draft.append(value)
+
+    def remove(self, value: T) -> None:
+        """Remove the first occurrence of value from the draft list."""
+        if value in self._frozen:
+            raise ValueError(f"Cannot remove item from frozen list: {value!r}")
+        self._draft.remove(value)
+
+    def freeze(self) -> None:
+        """Freeze any draft items by adding them to the saved tuple."""
+        # Filter out duplicates already in saved before extending
+        new_items = tuple(item for item in self._draft if item not in self._frozen)
+        self._frozen = self._frozen + new_items
+        self._draft.clear()
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, Sequence):
+            return NotImplemented
+        return list(self) == list(value)
+
+    def __contains__(self, value: Any) -> bool:
+        return value in self._frozen or value in self._draft
+
+    def __len__(self) -> int:
+        return len(self._frozen) + len(self._draft)
+
+    def __iter__(self) -> Iterator[T]:
+        return iter(chain(self._frozen, self._draft))
+
+    @overload
+    def __getitem__(self, index: int) -> T: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> Sequence[T]: ...
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[T, Sequence[T]]:
+        return [*self._frozen, *self._draft][index]
+
+    @overload
+    def __setitem__(self, index: int, value: T) -> None: ...
+
+    @overload
+    def __setitem__(self, index: slice, value: Iterable[T]) -> None: ...
+
+    def __setitem__(
+        self, index: Union[int, slice], value: Union[T, Iterable[T]]
+    ) -> None:
+        if isinstance(index, slice):
+            # Setting slices might affect saved items, disallow for simplicity
+            raise TypeError(
+                f"{type(self).__name__!r} does not support slice assignment"
+            )
+        else:
+            if value in self._frozen or value in self._draft:
+                return
+
+            # The frozen items are sequentially first and protected from changes
+            len_frozen = len(self._frozen)
+            size = len(self)
+
+            if (index >= size) or (index < -size):
+                raise IndexError("Index out of range")
+
+            draft_index = (index % size) - len_frozen
+            if draft_index < 0:
+                raise ValueError(f"Cannot assign to saved item at index {index!r}")
+            self._draft[draft_index] = value
+
+    @overload
+    def __delitem__(self, index: int) -> None: ...
+
+    @overload
+    def __delitem__(self, index: slice) -> None: ...
+
+    def __delitem__(self, index: Union[int, slice]) -> None:
+        if isinstance(index, slice):
+            raise TypeError(f"{type(self).__name__!r} does not support slice deletion")
+        else:
+            # The frozen items are sequentially first and protected from changes
+            len_frozen = len(self._frozen)
+            size = len(self)
+
+            if (index >= size) or (index < -size):
+                raise IndexError("Index out of range")
+
+            draft_index = (index % size) - len_frozen
+            if draft_index < 0:
+                raise ValueError(f"Cannot delete saved item at index {index!r}")
+            del self._draft[draft_index]
+
+    def insert(self, index: int, value: T) -> None:
+        """Insert item before index.
+
+        Insertion is only allowed at indices corresponding to the draft portion
+        of the list (i.e., index >= len(frozen_items)). Negative indices are
+        interpreted relative to the combined length of frozen and draft items.
+        """
+        if value in self._frozen or value in self._draft:
+            # Silently ignore duplicates, similar to append
+            return
+
+        # The frozen items are sequentially first and protected from changes
+        len_frozen = len(self._frozen)
+        size = len(self)
+
+        # Follow the behavior of `list.insert()` when the index is out of bounds.
+        # - negative out-of-bounds index: prepend.  Will only work if the frozen items are empty.
+        if index < -size and not self._frozen:
+            return self._draft.insert(0, value)
+
+        # - positive out-of-bounds index: append.
+        if index >= size:
+            return self._draft.append(value)
+
+        # - in-bounds index: insert only if into the draft portion.
+        draft_index = (index % size) - len_frozen
+        if draft_index < 0:
+            raise IndexError(
+                f"Cannot insert into the frozen list (index < {len_frozen})"
+            )
+        return self._draft.insert(draft_index, value)
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(frozen={list(self._frozen)!r}, draft={list(self._draft)!r})"
+
+    @property
+    def draft(self) -> Tuple[T, ...]:
+        """A read-only, tuple copy of the current draft items."""
+        return tuple(self._draft)
+
+
+class AddOnlyArtifactTypesList(FreezableList[str]):
+    def remove(self, value: str) -> None:
+        try:
+            super().remove(value)
+        except ValueError:
+            raise ValueError(
+                f"Cannot remove artifact type: {value!r} that has been saved to the registry"
+            )
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(saved={list(self._frozen)!r}, draft={list(self._draft)!r})"

wandb/apis/public/{registries.py → registries/registries_search.py}

@@ -1,7 +1,7 @@
-"""Public API: registries."""
+"""Public API: registries search."""
 
 import json
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 if TYPE_CHECKING:
     from wandb_gql import Client
@@ -11,11 +11,11 @@ from wandb_gql import gql
 import wandb
 from wandb.apis.paginator import Paginator
 from wandb.apis.public.artifacts import ArtifactCollection
+from wandb.apis.public.registries.utils import _ensure_registry_prefix_on_names
 from wandb.sdk.artifacts._graphql_fragments import (
     _gql_artifact_fragment,
     _gql_registry_fragment,
 )
-from wandb.sdk.artifacts._validators import REGISTRY_PREFIX
 
 
 class Registries(Paginator):
@@ -63,7 +63,7 @@ class Registries(Paginator):
         super().__init__(client, variables, per_page)
 
     def __bool__(self):
-        return len(self) > 0 or len(self.objects) > 0
+        return bool(self.objects)
 
     def __next__(self):
         # Implement custom next since its possible to load empty pages because of auth
@@ -128,6 +128,8 @@ class Registries(Paginator):
                 f"Organization '{self.organization}' not found. Please verify the organization name is correct"
             )
 
+        from wandb.apis.public.registries.registry import Registry
+
         return [
             Registry(
                 self.client,
@@ -142,86 +144,6 @@ class Registries(Paginator):
         ]
 
 
-class Registry:
-    """A single registry in the Registry."""
-
-    def __init__(
-        self,
-        client: "Client",
-        organization: str,
-        entity: str,
-        full_name: str,
-        attrs: Dict[str, Any],
-    ):
-        self.client = client
-        self._full_name = full_name
-        self._name = full_name.replace(REGISTRY_PREFIX, "")
-        self._entity = entity
-        self._organization = organization
-        self._description = attrs.get("description", "")
-        self._allow_all_artifact_types = attrs.get(
-            "allowAllArtifactTypesInRegistry", False
-        )
-        self._artifact_types = [
-            t["node"]["name"] for t in attrs.get("artifactTypes", {}).get("edges", [])
-        ]
-        self._id = attrs.get("id", "")
-        self._created_at = attrs.get("createdAt", "")
-        self._updated_at = attrs.get("updatedAt", "")
-
-    @property
-    def full_name(self):
-        return self._full_name
-
-    @property
-    def name(self):
-        return self._name
-
-    @property
-    def entity(self):
-        return self._entity
-
-    @property
-    def organization(self):
-        return self._organization
-
-    @property
-    def description(self):
-        return self._description
-
-    @property
-    def allow_all_artifact_types(self):
-        return self._allow_all_artifact_types
-
-    @property
-    def artifact_types(self):
-        return self._artifact_types
-
-    @property
-    def created_at(self):
-        return self._created_at
-
-    @property
-    def updated_at(self):
-        return self._updated_at
-
-    @property
-    def path(self):
-        return [self.entity, self.name]
-
-    def collections(self, filter: Optional[Dict[str, Any]] = None):
-        registry_filter = {
-            "name": self.full_name,
-        }
-        return Collections(self.client, self.organization, registry_filter, filter)
-
-    def versions(self, filter: Optional[Dict[str, Any]] = None):
-        registry_filter = {
-            "name": self.full_name,
-        }
-        return Versions(self.client, self.organization, registry_filter, None, filter)
-
-
 class Collections(Paginator):
     """Iterator that returns Artifact collections in the Registry."""
 
@@ -303,12 +225,12 @@ class Collections(Paginator):
         self.collection_filter = collection_filter or {}
 
         variables = {
-            "registryFilter": json.dumps(self.registry_filter)
-            if self.registry_filter
-            else None,
-            "collectionFilter": json.dumps(self.collection_filter)
-            if self.collection_filter
-            else None,
+            "registryFilter": (
+                json.dumps(self.registry_filter) if self.registry_filter else None
+            ),
+            "collectionFilter": (
+                json.dumps(self.collection_filter) if self.collection_filter else None
+            ),
             "organization": self.organization,
             "collectionTypes": ["PORTFOLIO"],
             "perPage": per_page,
@@ -383,6 +305,7 @@ class Collections(Paginator):
                 r["node"]["defaultArtifactType"]["name"],
                 self.organization,
                 r["node"],
+                is_sequence=False,
             )
             for r in self.last_response["organization"]["orgEntity"][
                 "artifactCollections"
@@ -460,15 +383,15 @@ class Versions(Paginator):
         )
 
         variables = {
-            "registryFilter": json.dumps(self.registry_filter)
-            if self.registry_filter
-            else None,
-            "collectionFilter": json.dumps(self.collection_filter)
-            if self.collection_filter
-            else None,
-            "artifactFilter": json.dumps(self.artifact_filter)
-            if self.artifact_filter
-            else None,
+            "registryFilter": (
+                json.dumps(self.registry_filter) if self.registry_filter else None
+            ),
+            "collectionFilter": (
+                json.dumps(self.collection_filter) if self.collection_filter else None
+            ),
+            "artifactFilter": (
+                json.dumps(self.artifact_filter) if self.artifact_filter else None
+            ),
             "organization": self.organization,
         }
 
@@ -541,33 +464,3 @@ class Versions(Paginator):
             ]["edges"]
         )
         return artifacts
-
-
-def _ensure_registry_prefix_on_names(query, in_name=False):
-    """Traverse the filter to prepend the `name` key value with the registry prefix unless the value is a regex.
-
-    - in_name: True if we are under a "name" key (or propagating from one).
-
-    EX: {"name": "model"} -> {"name": "wandb-registry-model"}
-    """
-    if isinstance((txt := query), str):
-        if in_name:
-            return txt if txt.startswith(REGISTRY_PREFIX) else f"{REGISTRY_PREFIX}{txt}"
-        return txt
-    if isinstance((dct := query), Mapping):
-        new_dict = {}
-        for key, obj in dct.items():
-            if key == "name":
-                new_dict[key] = _ensure_registry_prefix_on_names(obj, in_name=True)
-            elif key == "$regex":
-                # For regex operator, we skip transformation of its value.
-                new_dict[key] = obj
-            else:
-                # For any other key, propagate the in_name and skip_transform flags as-is.
-                new_dict[key] = _ensure_registry_prefix_on_names(obj, in_name=in_name)
-        return new_dict
-    if isinstance((objs := query), Sequence):
-        return list(
-            map(lambda x: _ensure_registry_prefix_on_names(x, in_name=in_name), objs)
-        )
-    return query