wandb 0.22.0__py3-none-macosx_12_0_arm64.whl → 0.22.2__py3-none-macosx_12_0_arm64.whl

wandb/sdk/artifacts/_generated/operations.py

@@ -33,6 +33,7 @@ __all__ = [
     "REGISTRY_VERSIONS_GQL",
     "RUN_INPUT_ARTIFACTS_GQL",
     "RUN_OUTPUT_ARTIFACTS_GQL",
+    "TYPE_INFO_GQL",
     "UNLINK_ARTIFACT_GQL",
     "UPDATE_ARTIFACT_GQL",
     "UPDATE_ARTIFACT_PORTFOLIO_GQL",
@@ -456,17 +457,7 @@ query RunOutputArtifacts($entity: String!, $project: String!, $runName: String!,
   project(name: $project, entityName: $entity) {
     run(name: $runName) {
       outputArtifacts(after: $cursor, first: $perPage) {
-        totalCount
-        edges {
-          node {
-            ...ArtifactFragment
-          }
-          cursor
-        }
-        pageInfo {
-          endCursor
-          hasNextPage
-        }
+        ...RunOutputArtifactConnectionFragment
       }
     }
   }
@@ -519,6 +510,20 @@ fragment ArtifactFragmentWithoutAliases on Artifact {
   createdAt
   updatedAt
 }
+
+fragment RunOutputArtifactConnectionFragment on ArtifactConnection {
+  totalCount
+  edges {
+    node {
+      ...ArtifactFragment
+    }
+    cursor
+  }
+  pageInfo {
+    endCursor
+    hasNextPage
+  }
+}
 """
 
 RUN_INPUT_ARTIFACTS_GQL = """
@@ -526,17 +531,7 @@ query RunInputArtifacts($entity: String!, $project: String!, $runName: String!,
   project(name: $project, entityName: $entity) {
     run(name: $runName) {
       inputArtifacts(after: $cursor, first: $perPage) {
-        totalCount
-        edges {
-          node {
-            ...ArtifactFragment
-          }
-          cursor
-        }
-        pageInfo {
-          endCursor
-          hasNextPage
-        }
+        ...RunInputArtifactConnectionFragment
       }
     }
   }
@@ -589,6 +584,20 @@ fragment ArtifactFragmentWithoutAliases on Artifact {
   createdAt
   updatedAt
 }
+
+fragment RunInputArtifactConnectionFragment on InputArtifactConnection {
+  totalCount
+  edges {
+    node {
+      ...ArtifactFragment
+    }
+    cursor
+  }
+  pageInfo {
+    endCursor
+    hasNextPage
+  }
+}
 """
 
 FETCH_LINKED_ARTIFACTS_GQL = """
@@ -1040,6 +1049,27 @@ mutation UnlinkArtifact($artifactID: ID!, $artifactPortfolioID: ID!) {
 }
 """
 
+TYPE_INFO_GQL = """
+query TypeInfo($name: String!) {
+  __type(name: $name) {
+    ...TypeInfoFragment
+  }
+}
+
+fragment TypeInfoFragment on __Type {
+  name
+  fields {
+    name
+    args {
+      name
+    }
+  }
+  inputFields {
+    name
+  }
+}
+"""
+
 REGISTRY_VERSIONS_GQL = """
 query RegistryVersions($organization: String!, $registryFilter: JSONString, $collectionFilter: JSONString, $artifactFilter: JSONString, $cursor: String, $perPage: Int) {
   organization(name: $organization) {

wandb/sdk/artifacts/_generated/run_input_artifacts.py

@@ -3,13 +3,13 @@
 
 from __future__ import annotations
 
-from typing import List, Optional
+from typing import Optional
 
 from pydantic import Field
 
 from wandb._pydantic import GQLBase
 
-from .fragments import ArtifactFragment
+from .fragments import RunInputArtifactConnectionFragment
 
 
 class RunInputArtifacts(GQLBase):
@@ -21,31 +21,11 @@ class RunInputArtifactsProject(GQLBase):
 
 
 class RunInputArtifactsProjectRun(GQLBase):
-    input_artifacts: Optional[RunInputArtifactsProjectRunInputArtifacts] = Field(
+    input_artifacts: Optional[RunInputArtifactConnectionFragment] = Field(
         alias="inputArtifacts"
     )
 
 
-class RunInputArtifactsProjectRunInputArtifacts(GQLBase):
-    total_count: int = Field(alias="totalCount")
-    edges: List[RunInputArtifactsProjectRunInputArtifactsEdges]
-    page_info: RunInputArtifactsProjectRunInputArtifactsPageInfo = Field(
-        alias="pageInfo"
-    )
-
-
-class RunInputArtifactsProjectRunInputArtifactsEdges(GQLBase):
-    node: Optional[ArtifactFragment]
-    cursor: str
-
-
-class RunInputArtifactsProjectRunInputArtifactsPageInfo(GQLBase):
-    end_cursor: Optional[str] = Field(alias="endCursor")
-    has_next_page: bool = Field(alias="hasNextPage")
-
-
 RunInputArtifacts.model_rebuild()
 RunInputArtifactsProject.model_rebuild()
 RunInputArtifactsProjectRun.model_rebuild()
-RunInputArtifactsProjectRunInputArtifacts.model_rebuild()
-RunInputArtifactsProjectRunInputArtifactsEdges.model_rebuild()

wandb/sdk/artifacts/_generated/run_output_artifacts.py

@@ -3,13 +3,13 @@
 
 from __future__ import annotations
 
-from typing import List, Optional
+from typing import Optional
 
 from pydantic import Field
 
 from wandb._pydantic import GQLBase
 
-from .fragments import ArtifactFragment
+from .fragments import RunOutputArtifactConnectionFragment
 
 
 class RunOutputArtifacts(GQLBase):
@@ -21,31 +21,11 @@ class RunOutputArtifactsProject(GQLBase):
 
 
 class RunOutputArtifactsProjectRun(GQLBase):
-    output_artifacts: Optional[RunOutputArtifactsProjectRunOutputArtifacts] = Field(
+    output_artifacts: Optional[RunOutputArtifactConnectionFragment] = Field(
         alias="outputArtifacts"
     )
 
 
-class RunOutputArtifactsProjectRunOutputArtifacts(GQLBase):
-    total_count: int = Field(alias="totalCount")
-    edges: List[RunOutputArtifactsProjectRunOutputArtifactsEdges]
-    page_info: RunOutputArtifactsProjectRunOutputArtifactsPageInfo = Field(
-        alias="pageInfo"
-    )
-
-
-class RunOutputArtifactsProjectRunOutputArtifactsEdges(GQLBase):
-    node: Optional[ArtifactFragment]
-    cursor: str
-
-
-class RunOutputArtifactsProjectRunOutputArtifactsPageInfo(GQLBase):
-    end_cursor: Optional[str] = Field(alias="endCursor")
-    has_next_page: bool = Field(alias="hasNextPage")
-
-
 RunOutputArtifacts.model_rebuild()
 RunOutputArtifactsProject.model_rebuild()
 RunOutputArtifactsProjectRun.model_rebuild()
-RunOutputArtifactsProjectRunOutputArtifacts.model_rebuild()
-RunOutputArtifactsProjectRunOutputArtifactsEdges.model_rebuild()

wandb/sdk/artifacts/_generated/type_info.py

@@ -0,0 +1,19 @@
+# Generated by ariadne-codegen
+# Source: tools/graphql_codegen/artifacts/
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import Field
+
+from wandb._pydantic import GQLBase
+
+from .fragments import TypeInfoFragment
+
+
+class TypeInfo(GQLBase):
+    type: Optional[TypeInfoFragment] = Field(alias="__type")
+
+
+TypeInfo.model_rebuild()

wandb/sdk/artifacts/_gqlutils.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from functools import lru_cache
+
+from wandb_gql import Client, gql
+
+from ._generated import TYPE_INFO_GQL, TypeInfo, TypeInfoFragment
+
+OMITTABLE_ARTIFACT_FIELDS = frozenset(
+    {
+        "ttlDurationSeconds",
+        "ttlIsInherited",
+        "aliases",
+        "tags",
+        "historyStep",
+    }
+)
+
+
+@lru_cache(maxsize=16)
+def type_info(client: Client, typename: str) -> TypeInfoFragment | None:
+    """Returns the type info for a given GraphQL type."""
+    data = client.execute(gql(TYPE_INFO_GQL), variable_values={"name": typename})
+    return TypeInfo.model_validate(data).type
+
+
+def supports_enable_tracking_var(client: Client) -> bool:
+    """Returns True if the server supports the `enableTracking` variable for the `Project.artifact(...)` field."""
+    typ = type_info(client, "Project")
+    if (
+        typ
+        and typ.fields
+        and (art_field := next((f for f in typ.fields if f.name == "artifact"), None))
+    ):
+        return any("enableTracking" == arg.name for arg in art_field.args)
+    return False
+
+
+def allowed_fields(client: Client, typename: str) -> set[str]:
+    """Returns the allowed field names for a given GraphQL type."""
+    typ = type_info(client, typename)
+    return {f.name for f in typ.fields} if (typ and typ.fields) else set()
+
+
+def omit_artifact_fields(client: Client) -> set[str]:
+    """Return names of Artifact fields to remove from GraphQL requests (for server compatibility)."""
+    return set(OMITTABLE_ARTIFACT_FIELDS) - allowed_fields(client, "Artifact")

wandb/sdk/artifacts/_models/__init__.py

@@ -0,0 +1,4 @@
+"""Pydantic model classes and related helpers for artifacts code.
+
+Excludes GraphQL-generated classes.
+"""

wandb/sdk/artifacts/_models/base_model.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from abc import ABC
+
+from pydantic import ConfigDict
+
+from wandb._pydantic import JsonableModel
+
+
+# Abstract base class with a common default configuration that's shared by all
+# pydantic classes in artifacts code (excluding GraphQL-generated types).
+class ArtifactsBase(JsonableModel, ABC):
+    # See: https://docs.pydantic.dev/latest/api/config/#pydantic.config.ConfigDict
+    model_config = ConfigDict(
+        # Most likely, some fields won't be pydantic types
+        arbitrary_types_allowed=True,
+        # Assume instances of the same class have already been validated to save time,
+        # but validate subclasses in case they override the default behavior.
+        revalidate_instances="subclass-instances",
+    )

wandb/sdk/artifacts/_validators.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import json
 import re
-from dataclasses import astuple, dataclass, field, replace
+from dataclasses import dataclass, field, replace
 from functools import wraps
 from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, TypeVar, cast
 
@@ -277,34 +277,62 @@ def remove_registry_prefix(project: str) -> str:
 
 @pydantic_dataclass
 class ArtifactPath:
-    #: The collection name.
     name: str
-    #: The project name, which can also be a registry name.
+    """The collection or artifact version name."""
     project: Optional[str] = None  # noqa: UP045
-    #: Prefix is often an org or entity name.
+    """The project name."""
     prefix: Optional[str] = None  # noqa: UP045
+    """Typically the entity or org name."""
 
     @classmethod
     def from_str(cls, path: str) -> Self:
-        """Instantiate by parsing an artifact path."""
-        if len(parts := path.split("/")) <= 3:
-            return cls(*reversed(parts))
-        raise ValueError(
-            f"Expected a valid path like `name`, `project/name`, or `prefix/project/name`.  Got: {path!r}"
-        )
+        """Instantiate by parsing a string artifact path.
+
+        Raises:
+            ValueError: If the string is not a valid artifact path.
+        """
+        # Separate the alias first, which may itself contain slashes.
+        # If there's no alias, note that both sep and alias will be empty.
+        collection_path, sep, alias = path.partition(":")
+
+        prefix, project = None, None  # defaults, if missing
+        if len(parts := collection_path.split("/")) == 1:
+            name = parts[0]
+        elif len(parts) == 2:
+            project, name = parts
+        elif len(parts) == 3:
+            prefix, project, name = parts
+        else:
+            raise ValueError(f"Invalid artifact path: {path!r}")
+        return cls(prefix=prefix, project=project, name=f"{name}{sep}{alias}")
 
     def to_str(self) -> str:
         """Returns the slash-separated string representation of the path."""
-        return "/".join(filter(bool, reversed(astuple(self))))
+        ordered_parts = (self.prefix, self.project, self.name)
+        return "/".join(part for part in ordered_parts if part)
 
     def with_defaults(
         self,
+        *,
         prefix: str | None = None,
         project: str | None = None,
     ) -> Self:
-        """Returns this path with missing values set to the given defaults."""
+        """Returns a copy of this path with missing values set to the given defaults."""
         return replace(
             self,
             prefix=self.prefix or prefix,
             project=self.project or project,
         )
+
+    def is_registry_path(self) -> bool:
+        """Returns True if this path appears to be a registry path."""
+        return bool((p := self.project) and is_artifact_registry_project(p))
+
+
+@pydantic_dataclass
+class FullArtifactPath(ArtifactPath):
+    """Same as ArtifactPath, but with all parts required."""
+
+    name: str
+    project: str
+    prefix: str