wandb 0.22.2__py3-none-win_amd64.whl → 0.22.3__py3-none-win_amd64.whl

wandb/apis/public/registries/_utils.py

@@ -1,14 +1,11 @@
 from __future__ import annotations
 
 from enum import Enum
-from functools import lru_cache
-from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence
+from functools import lru_cache, partial
+from typing import TYPE_CHECKING, Any, Collection
 
 from wandb._strutils import ensureprefix
-from wandb.sdk.artifacts._validators import (
-    REGISTRY_PREFIX,
-    validate_artifact_types_list,
-)
+from wandb.sdk.artifacts._validators import REGISTRY_PREFIX, validate_artifact_types
 
 if TYPE_CHECKING:
     from wandb_gql import Client
@@ -24,12 +21,35 @@ class Visibility(str, Enum):
 
     @classmethod
     def _missing_(cls, value: object) -> Any:
-        return next((e for e in cls if e.name == value), None)
+        # Allow instantiation from enum names too (e.g. "organization" or "restricted")
+        return cls.__members__.get(value)
 
+    @classmethod
+    def from_gql(cls, value: str) -> Visibility:
+        """Convert a GraphQL `visibility` value to a Visibility enum."""
+        try:
+            return cls(value)
+        except ValueError:
+            expected = ",".join(repr(e.value) for e in cls)
+            raise ValueError(
+                f"Invalid visibility {value!r} from backend. Expected one of: {expected}"
+            ) from None
 
-def format_gql_artifact_types_input(
-    artifact_types: list[str] | None,
-) -> list[dict[str, str]]:
+    @classmethod
+    def from_python(cls, name: str) -> Visibility:
+        """Convert a visibility string (e.g. user-provided or python default value) to a Visibility enum."""
+        try:
+            return cls(name)
+        except ValueError:
+            expected = ",".join(repr(e.name) for e in cls)
+            raise ValueError(
+                f"Invalid visibility {name!r}. Expected one of: {expected}"
+            ) from None
+
+
+def prepare_artifact_types_input(
+    artifact_types: Collection[str] | None,
+) -> list[dict[str, str]] | None:
     """Format the artifact types for the GQL input.
 
     Args:
@@ -38,39 +58,9 @@ def format_gql_artifact_types_input(
     Returns:
         The artifact types for the GQL input.
     """
-    if artifact_types is None:
-        return []
-    return [{"name": typ} for typ in validate_artifact_types_list(artifact_types)]
-
-
-def gql_to_registry_visibility(
-    visibility: str,
-) -> Literal["organization", "restricted"]:
-    """Convert the GQL visibility to the registry visibility.
-
-    Args:
-        visibility: The GQL visibility.
-
-    Returns:
-        The registry visibility.
-    """
-    try:
-        return Visibility(visibility).name
-    except ValueError:
-        raise ValueError(f"Invalid visibility: {visibility!r} from backend")
-
-
-def registry_visibility_to_gql(
-    visibility: Literal["organization", "restricted"],
-) -> str:
-    """Convert the registry visibility to the GQL visibility."""
-    try:
-        return Visibility[visibility].value
-    except LookupError:
-        allowed_str = ", ".join(map(repr, (e.name for e in Visibility)))
-        raise ValueError(
-            f"Invalid visibility: {visibility!r}. Must be one of: {allowed_str}"
-        )
+    if artifact_types:
+        return [{"name": typ} for typ in validate_artifact_types(artifact_types)]
+    return None
 
 
 def ensure_registry_prefix_on_names(query: Any, in_name: bool = False) -> Any:
@@ -81,25 +71,21 @@ def ensure_registry_prefix_on_names(query: Any, in_name: bool = False) -> Any:
     EX: {"name": "model"} -> {"name": "wandb-registry-model"}
     """
     if isinstance((txt := query), str):
-        if in_name:
-            return ensureprefix(txt, REGISTRY_PREFIX)
-        return txt
-    if isinstance((dct := query), Mapping):
+        return ensureprefix(txt, REGISTRY_PREFIX) if in_name else txt
+    if isinstance((dct := query), dict):
         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":
+            if key == "$regex":
                 # For regex operator, we skip transformation of its value.
                 new_dict[key] = obj
+            elif key == "name":
+                new_dict[key] = ensure_registry_prefix_on_names(obj, in_name=True)
             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)
-        )
+    if isinstance((seq := query), (list, tuple)):
+        return list(map(partial(ensure_registry_prefix_on_names, in_name=in_name), seq))
     return query
 
 

wandb/apis/public/registries/registries_search.py

@@ -5,33 +5,38 @@ from __future__ import annotations
 import json
 from typing import TYPE_CHECKING, Any
 
-from pydantic import ValidationError
+from pydantic import PositiveInt, ValidationError
 from typing_extensions import override
 from wandb_gql import gql
 
 from wandb._analytics import tracked
 from wandb.apis.paginator import Paginator
+from wandb.apis.public.artifacts import ArtifactCollection
 from wandb.apis.public.utils import gql_compat
 from wandb.sdk.artifacts._generated import (
     FETCH_REGISTRIES_GQL,
     REGISTRY_COLLECTIONS_GQL,
     REGISTRY_VERSIONS_GQL,
-    ArtifactCollectionType,
     FetchRegistries,
-    RegistriesPage,
     RegistryCollections,
-    RegistryCollectionsPage,
     RegistryVersions,
-    RegistryVersionsPage,
 )
 from wandb.sdk.artifacts._gqlutils import omit_artifact_fields
-from wandb.sdk.artifacts._validators import FullArtifactPath, remove_registry_prefix
+from wandb.sdk.artifacts._models.pagination import (
+    ArtifactMembershipConnection,
+    RegistryCollectionConnection,
+    RegistryConnection,
+)
+from wandb.sdk.artifacts._validators import (
+    SOURCE_COLLECTION_TYPENAME,
+    FullArtifactPath,
+    remove_registry_prefix,
+)
 
 from ._utils import ensure_registry_prefix_on_names
 
 if TYPE_CHECKING:
-    from wandb_gql import Client
-
+    from wandb.apis.public import RetryingClient
     from wandb.sdk.artifacts.artifact import Artifact
 
 
@@ -40,15 +45,14 @@ class Registries(Paginator):
 
     QUERY = gql(FETCH_REGISTRIES_GQL)
 
-    last_response: RegistriesPage | None
-    _last_org_entity: str | None
+    last_response: RegistryConnection | None
 
     def __init__(
         self,
-        client: Client,
+        client: RetryingClient,
         organization: str,
         filter: dict[str, Any] | None = None,
-        per_page: int | None = 100,
+        per_page: PositiveInt = 100,
     ):
         self.client = client
         self.organization = organization
@@ -57,11 +61,8 @@ class Registries(Paginator):
             "organization": organization,
             "filters": json.dumps(self.filter),
         }
-
         super().__init__(client, variables, per_page)
 
-        self._last_org_entity = None
-
     def __next__(self):
         # Implement custom next since its possible to load empty pages because of auth
         self.index += 1
@@ -71,22 +72,28 @@ class Registries(Paginator):
         return self.objects[self.index]
 
     @tracked
-    def collections(self, filter: dict[str, Any] | None = None) -> Collections:
+    def collections(
+        self, filter: dict[str, Any] | None = None, per_page: PositiveInt = 100
+    ) -> Collections:
         return Collections(
             client=self.client,
             organization=self.organization,
             registry_filter=self.filter,
             collection_filter=filter,
+            per_page=per_page,
         )
 
     @tracked
-    def versions(self, filter: dict[str, Any] | None = None) -> Versions:
+    def versions(
+        self, filter: dict[str, Any] | None = None, per_page: PositiveInt = 100
+    ) -> Versions:
         return Versions(
             client=self.client,
             organization=self.organization,
             registry_filter=self.filter,
             collection_filter=None,
             artifact_filter=filter,
+            per_page=per_page,
         )
 
     @property
@@ -97,15 +104,11 @@ class Registries(Paginator):
 
     @property
     def more(self):
-        if self.last_response is None:
-            return True
-        return self.last_response.page_info.has_next_page
+        return (conn := self.last_response) is None or conn.has_next
 
     @property
     def cursor(self):
-        if self.last_response is None:
-            return None
-        return self.last_response.page_info.end_cursor
+        return conn.next_cursor if (conn := self.last_response) else None
 
     @override
     def _update_response(self) -> None:
@@ -117,45 +120,43 @@ class Registries(Paginator):
             )
 
         try:
-            page_data = org_entity.projects
-            self.last_response = RegistriesPage.model_validate(page_data)
-            self._last_org_entity = org_entity.name
+            conn = org_entity.projects
+            self.last_response = RegistryConnection.model_validate(conn)
         except (LookupError, AttributeError, ValidationError) as e:
             raise ValueError("Unexpected response data") from e
 
     def convert_objects(self):
         from wandb.apis.public.registries.registry import Registry
 
-        if (self.last_response is None) or (self._last_org_entity is None):
+        if self.last_response is None:
             return []
 
-        nodes = (e.node for e in self.last_response.edges)
         return [
             Registry(
                 client=self.client,
                 organization=self.organization,
-                entity=self._last_org_entity,
+                entity=node.entity.name,
                 name=remove_registry_prefix(node.name),
-                attrs=node.model_dump(),
+                attrs=node,
             )
-            for node in nodes
+            for node in self.last_response.nodes()
         ]
 
 
-class Collections(Paginator["ArtifactCollection"]):
+class Collections(Paginator[ArtifactCollection]):
     """An lazy iterator of `ArtifactCollection` objects in a Registry."""
 
     QUERY = gql(REGISTRY_COLLECTIONS_GQL)
 
-    last_response: RegistryCollectionsPage | None
+    last_response: RegistryCollectionConnection | None
 
     def __init__(
         self,
-        client: Client,
+        client: RetryingClient,
         organization: str,
         registry_filter: dict[str, Any] | None = None,
         collection_filter: dict[str, Any] | None = None,
-        per_page: int | None = 100,
+        per_page: PositiveInt = 100,
     ):
         self.client = client
         self.organization = organization
@@ -166,7 +167,6 @@ class Collections(Paginator["ArtifactCollection"]):
             "registryFilter": json.dumps(f) if (f := registry_filter) else None,
             "collectionFilter": json.dumps(f) if (f := collection_filter) else None,
             "organization": organization,
-            "collectionTypes": [ArtifactCollectionType.PORTFOLIO],
             "perPage": per_page,
         }
 
@@ -181,48 +181,42 @@ class Collections(Paginator["ArtifactCollection"]):
         return self.objects[self.index]
 
     @tracked
-    def versions(self, filter: dict[str, Any] | None = None) -> Versions:
+    def versions(
+        self, filter: dict[str, Any] | None = None, per_page: PositiveInt = 100
+    ) -> Versions:
         return Versions(
             client=self.client,
             organization=self.organization,
             registry_filter=self.registry_filter,
             collection_filter=self.collection_filter,
             artifact_filter=filter,
+            per_page=per_page,
         )
 
     @property
     def length(self):
-        if self.last_response is None:
-            return None
-        return self.last_response.total_count
+        return conn.total_count if (conn := self.last_response) else None
 
     @property
     def more(self):
-        if self.last_response is None:
-            return True
-        return self.last_response.page_info.has_next_page
+        return (conn := self.last_response) is None or conn.has_next
 
     @property
     def cursor(self):
-        if self.last_response is None:
-            return None
-        return self.last_response.page_info.end_cursor
+        return conn.next_cursor if (conn := self.last_response) else None
 
     @override
     def _update_response(self) -> None:
         data = self.client.execute(self.QUERY, variable_values=self.variables)
         result = RegistryCollections.model_validate(data)
-        if not (
-            (org_data := result.organization)
-            and (org_entity_data := org_data.org_entity)
-        ):
+        if not ((org := result.organization) and (org_entity := org.org_entity)):
             raise ValueError(
                 f"Organization {self.organization!r} not found. Please verify the organization name is correct."
             )
 
         try:
-            page_data = org_entity_data.artifact_collections
-            self.last_response = RegistryCollectionsPage.model_validate(page_data)
+            conn = org_entity.artifact_collections
+            self.last_response = RegistryCollectionConnection.model_validate(conn)
         except (LookupError, AttributeError, ValidationError) as e:
             raise ValueError("Unexpected response data") from e
 
@@ -232,36 +226,36 @@ class Collections(Paginator["ArtifactCollection"]):
         if self.last_response is None:
             return []
 
-        nodes = (e.node for e in self.last_response.edges)
         return [
             ArtifactCollection(
                 client=self.client,
-                entity=project.entity.name,
-                project=project.name,
+                entity=node.project.entity.name,
+                project=node.project.name,
                 name=node.name,
-                type=node.default_artifact_type.name,
+                type=node.type.name,
                 organization=self.organization,
-                attrs=node.model_dump(),
-                is_sequence=False,
+                attrs=node,
             )
-            for node in nodes
-            if (project := node.project)
+            for node in self.last_response.nodes()
+            # We don't _expect_ any registry collections to be
+            # ArtifactSequences, but defensively filter them out anyway.
+            if node.project and (node.typename__ != SOURCE_COLLECTION_TYPENAME)
         ]
 
 
 class Versions(Paginator["Artifact"]):
     """An lazy iterator of `Artifact` objects in a Registry."""
 
-    last_response: RegistryVersionsPage | None
+    last_response: ArtifactMembershipConnection | None
 
     def __init__(
         self,
-        client: Client,
+        client: RetryingClient,
         organization: str,
         registry_filter: dict[str, Any] | None = None,
         collection_filter: dict[str, Any] | None = None,
         artifact_filter: dict[str, Any] | None = None,
-        per_page: int = 100,
+        per_page: PositiveInt = 100,
     ):
         self.client = client
         self.organization = organization
@@ -298,31 +292,24 @@ class Versions(Paginator["Artifact"]):
 
     @property
     def more(self) -> bool:
-        if self.last_response is None:
-            return True
-        return self.last_response.page_info.has_next_page
+        return (conn := self.last_response) is None or conn.has_next
 
     @property
     def cursor(self) -> str | None:
-        if self.last_response is None:
-            return None
-        return self.last_response.page_info.end_cursor
+        return conn.next_cursor if (conn := self.last_response) else None
 
     @override
     def _update_response(self) -> None:
         data = self.client.execute(self.QUERY, variable_values=self.variables)
         result = RegistryVersions.model_validate(data)
-        if not (
-            (org_data := result.organization)
-            and (org_entity_data := org_data.org_entity)
-        ):
+        if not ((org := result.organization) and (org_entity := org.org_entity)):
             raise ValueError(
                 f"Organization {self.organization!r} not found. Please verify the organization name is correct."
             )
 
         try:
-            page_data = org_entity_data.artifact_memberships
-            self.last_response = RegistryVersionsPage.model_validate(page_data)
+            conn = org_entity.artifact_memberships
+            self.last_response = ArtifactMembershipConnection.model_validate(conn)
         except (LookupError, AttributeError, ValidationError) as e:
             raise ValueError("Unexpected response data") from e
 
@@ -331,23 +318,21 @@ class Versions(Paginator["Artifact"]):
 
         if self.last_response is None:
             return []
-
-        nodes = (e.node for e in self.last_response.edges)
         return [
-            Artifact._from_attrs(
-                path=FullArtifactPath(
+            Artifact._from_membership(
+                membership=membership,
+                target=FullArtifactPath(
                     prefix=project.entity.name,
                     project=project.name,
-                    name=f"{collection.name}:v{node.version_index}",
+                    name=f"{collection.name}:v{version_idx}",
                 ),
-                attrs=artifact,
                 client=self.client,
-                aliases=[alias.alias for alias in node.aliases],
             )
-            for node in nodes
+            for membership in self.last_response.nodes()
             if (
-                (collection := node.artifact_collection)
+                (collection := membership.artifact_collection)
                 and (project := collection.project)
-                and (artifact := node.artifact)
+                and membership.artifact
+                and (version_idx := membership.version_index) is not None
             )
         ]