wandb 0.15.4__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/data_types/_dtypes.py

@@ -13,8 +13,7 @@ from wandb.util import (
 np = get_module("numpy")  # intentionally not required
 
 if t.TYPE_CHECKING:
-    from wandb.sdk.artifacts.local_artifact import Artifact as ArtifactInCreation
-    from wandb.sdk.artifacts.public_artifact import Artifact as DownloadedArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
 _TYPES_STRIPPED = not (sys.version_info.major == 3 and sys.version_info.minor >= 6)
 if not _TYPES_STRIPPED:
@@ -77,7 +76,7 @@ class TypeRegistry:
 
     @staticmethod
     def type_from_dict(
-        json_dict: t.Dict[str, t.Any], artifact: t.Optional["DownloadedArtifact"] = None
+        json_dict: t.Dict[str, t.Any], artifact: t.Optional["Artifact"] = None
     ) -> "Type":
         wb_type = json_dict.get("wb_type")
         if wb_type is None:
@@ -135,7 +134,7 @@ class TypeRegistry:
 
 def _params_obj_to_json_obj(
     params_obj: t.Any,
-    artifact: t.Optional["ArtifactInCreation"] = None,
+    artifact: t.Optional["Artifact"] = None,
 ) -> t.Any:
     """Helper method."""
     if params_obj.__class__ == dict:
@@ -152,7 +151,7 @@ def _params_obj_to_json_obj(
 
 
 def _json_obj_to_params_obj(
-    json_obj: t.Any, artifact: t.Optional["DownloadedArtifact"] = None
+    json_obj: t.Any, artifact: t.Optional["Artifact"] = None
 ) -> t.Any:
     """Helper method."""
     if json_obj.__class__ == dict:
@@ -222,9 +221,7 @@ class Type:
         else:
             return InvalidType()
 
-    def to_json(
-        self, artifact: t.Optional["ArtifactInCreation"] = None
-    ) -> t.Dict[str, t.Any]:
+    def to_json(self, artifact: t.Optional["Artifact"] = None) -> t.Dict[str, t.Any]:
         """Generate a jsonable dictionary serialization the type.
 
         If overridden by subclass, ensure that `from_json` is equivalently overridden.
@@ -249,7 +246,7 @@ class Type:
     def from_json(
         cls,
         json_dict: t.Dict[str, t.Any],
-        artifact: t.Optional["DownloadedArtifact"] = None,
+        artifact: t.Optional["Artifact"] = None,
     ) -> "Type":
         """Construct a new instance of the type using a JSON dictionary.
 
@@ -756,9 +753,7 @@ class NDArrayType(Type):
 
         return InvalidType()
 
-    def to_json(
-        self, artifact: t.Optional["ArtifactInCreation"] = None
-    ) -> t.Dict[str, t.Any]:
+    def to_json(self, artifact: t.Optional["Artifact"] = None) -> t.Dict[str, t.Any]:
         # custom override to support serialization path outside of params internal dict
         res = {
             "wb_type": self.name,

wandb/sdk/data_types/base_types/json_metadata.py

@@ -9,7 +9,7 @@ from .._private import MEDIA_TMP
 from .media import Media
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ...wandb_run import Run as LocalRun
 
@@ -40,7 +40,7 @@ class JSONMetadata(Media):
     def get_media_subdir(cls: Type["JSONMetadata"]) -> str:
         return os.path.join("media", "metadata", cls.type_name())
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = self.type_name()
 

wandb/sdk/data_types/base_types/media.py

@@ -14,10 +14,9 @@ from wandb.sdk.lib.paths import LogicalPath
 from .wb_value import WBValue
 
 if TYPE_CHECKING:  # pragma: no cover
-    import numpy as np  # type: ignore
+    import numpy as np
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ...wandb_run import Run as LocalRun
 
@@ -144,7 +143,7 @@ class Media(WBValue):
             self._path = new_path
             _datatypes_callback(media_path)
 
-    def to_json(self, run: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run: Union["LocalRun", "Artifact"]) -> dict:
         """Serialize the object into a JSON blob.
 
         Uses run or artifact to store additional data. If `run_or_artifact` is a
@@ -254,7 +253,7 @@ class Media(WBValue):
 
     @classmethod
     def from_json(
-        cls: Type["Media"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["Media"], json_obj: dict, source_artifact: "Artifact"
     ) -> "Media":
         """Likely will need to override for any more complicated media objects."""
         return cls(source_artifact.get_path(json_obj["path"]).download())
@@ -316,4 +315,4 @@ def _numpy_arrays_to_lists(
     # Protects against logging non serializable objects
     elif isinstance(payload, Media):
         return str(payload.__class__.__name__)
-    return payload
+    return payload  # type: ignore

wandb/sdk/data_types/base_types/wb_value.py

@@ -3,8 +3,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Type, Uni
 from wandb import util
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ...wandb_run import Run as LocalRun
 
@@ -36,19 +35,19 @@ def _server_accepts_client_ids() -> bool:
 
 
 class _WBValueArtifactSource:
-    artifact: "PublicArtifact"
+    artifact: "Artifact"
     name: Optional[str]
 
-    def __init__(self, artifact: "PublicArtifact", name: Optional[str] = None) -> None:
+    def __init__(self, artifact: "Artifact", name: Optional[str] = None) -> None:
         self.artifact = artifact
         self.name = name
 
 
 class _WBValueArtifactTarget:
-    artifact: "LocalArtifact"
+    artifact: "Artifact"
     name: Optional[str]
 
-    def __init__(self, artifact: "LocalArtifact", name: Optional[str] = None) -> None:
+    def __init__(self, artifact: "Artifact", name: Optional[str] = None) -> None:
         self.artifact = artifact
         self.name = name
 
@@ -73,7 +72,7 @@ class WBValue:
         self._artifact_source = None
         self._artifact_target = None
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         """Serialize the object into a JSON blob.
 
         Uses current run or artifact to store additional data.
@@ -90,7 +89,7 @@ class WBValue:
 
     @classmethod
     def from_json(
-        cls: Type["WBValue"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["WBValue"], json_obj: dict, source_artifact: "Artifact"
     ) -> "WBValue":
         """Deserialize a `json_obj` into it's class representation.
 
@@ -126,7 +125,7 @@ class WBValue:
 
     @staticmethod
     def init_from_json(
-        json_obj: dict, source_artifact: "PublicArtifact"
+        json_obj: dict, source_artifact: "Artifact"
     ) -> Optional["WBValue"]:
         """Initialize a `WBValue` from a JSON blob based on the class that creatd it.
 
@@ -186,7 +185,7 @@ class WBValue:
         raise NotImplementedError
 
     def _set_artifact_source(
-        self, artifact: "PublicArtifact", name: Optional[str] = None
+        self, artifact: "Artifact", name: Optional[str] = None
     ) -> None:
         assert (
             self._artifact_source is None
@@ -196,7 +195,7 @@ class WBValue:
         self._artifact_source = _WBValueArtifactSource(artifact, name)
 
     def _set_artifact_target(
-        self, artifact: "LocalArtifact", name: Optional[str] = None
+        self, artifact: "Artifact", name: Optional[str] = None
     ) -> None:
         assert (
             self._artifact_target is None
@@ -232,7 +231,7 @@ class WBValue:
         elif (
             self._artifact_target
             and self._artifact_target.name
-            and self._artifact_target.artifact._logged_artifact is not None
+            and self._artifact_target.artifact._is_draft_save_started()
             and not util._is_offline()
             and not _server_accepts_client_ids()
         ):
@@ -263,7 +262,7 @@ class WBValue:
         elif (
             self._artifact_target
             and self._artifact_target.name
-            and self._artifact_target.artifact._logged_artifact is not None
+            and self._artifact_target.artifact._is_draft_save_started()
             and not util._is_offline()
             and not _server_accepts_client_ids()
         ):

wandb/sdk/data_types/helper_types/bounding_boxes_2d.py

@@ -8,8 +8,7 @@ from wandb.util import has_num
 from ..base_types.json_metadata import JSONMetadata
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ...wandb_run import Run as LocalRun
 
@@ -274,7 +273,7 @@ class BoundingBoxes2D(JSONMetadata):
                 raise TypeError("A box's caption must be a string")
         return True
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         if isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run):
             return super().to_json(run_or_artifact)
         elif isinstance(run_or_artifact, wandb.Artifact):
@@ -283,10 +282,10 @@ class BoundingBoxes2D(JSONMetadata):
             # an object with a _type key. Will need to push this change to the UI first to ensure backwards compat
             return self._val
         else:
-            raise ValueError("to_json accepts wandb_run.Run or wandb_artifact.Artifact")
+            raise ValueError("to_json accepts wandb_run.Run or wandb.Artifact")
 
     @classmethod
     def from_json(
-        cls: Type["BoundingBoxes2D"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["BoundingBoxes2D"], json_obj: dict, source_artifact: "Artifact"
     ) -> "BoundingBoxes2D":
         return cls({"box_data": json_obj}, "")

wandb/sdk/data_types/helper_types/classes.py

@@ -5,8 +5,7 @@ from .. import _dtypes
 from ..base_types.media import Media
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ...wandb_run import Run as LocalRun
 
@@ -31,13 +30,11 @@ class Classes(Media):
     def from_json(
         cls: Type["Classes"],
         json_obj: dict,
-        source_artifact: Optional["PublicArtifact"],
+        source_artifact: Optional["Artifact"],
     ) -> "Classes":
         return cls(json_obj.get("class_set"))  # type: ignore
 
-    def to_json(
-        self, run_or_artifact: Optional[Union["LocalRun", "LocalArtifact"]]
-    ) -> dict:
+    def to_json(self, run_or_artifact: Optional[Union["LocalRun", "Artifact"]]) -> dict:
         json_obj = {}
         # This is a bit of a hack to allow _ClassesIdType to
         # be able to operate fully without an artifact in play.
@@ -116,7 +113,7 @@ class _ClassesIdType(_dtypes.Type):
     def from_obj(cls, py_obj: Optional[Any] = None) -> "_dtypes.Type":
         return cls(py_obj)
 
-    def to_json(self, artifact: Optional["LocalArtifact"] = None) -> Dict[str, Any]:
+    def to_json(self, artifact: Optional["Artifact"] = None) -> Dict[str, Any]:
         cl_dict = super().to_json(artifact)
         # TODO (tss): Refactor this block with the similar one in wandb.Image.
         # This is a bit of a smell that the classes object does not follow
@@ -137,7 +134,7 @@ class _ClassesIdType(_dtypes.Type):
     def from_json(
         cls,
         json_dict: Dict[str, Any],
-        artifact: Optional["PublicArtifact"] = None,
+        artifact: Optional["Artifact"] = None,
     ) -> "_dtypes.Type":
         classes_obj = None
         if (

wandb/sdk/data_types/helper_types/image_mask.py

@@ -10,8 +10,7 @@ from .._private import MEDIA_TMP
 from ..base_types.media import Media
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ...wandb_run import Run as LocalRun
 
@@ -184,14 +183,14 @@ class ImageMask(Media):
 
     @classmethod
     def from_json(
-        cls: Type["ImageMask"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["ImageMask"], json_obj: dict, source_artifact: "Artifact"
     ) -> "ImageMask":
         return cls(
             {"path": source_artifact.get_path(json_obj["path"]).download()},
             key="",
         )
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
 
         if isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run):
@@ -201,7 +200,7 @@ class ImageMask(Media):
             # Nothing special to add (used to add "digest", but no longer used.)
             return json_dict
         else:
-            raise ValueError("to_json accepts wandb_run.Run or wandb_artifact.Artifact")
+            raise ValueError("to_json accepts wandb_run.Run or wandb.Artifact")
 
     @classmethod
     def type_name(cls: Type["ImageMask"]) -> str:

wandb/sdk/data_types/histogram.py

@@ -6,9 +6,9 @@ from wandb import util
 from .base_types.wb_value import WBValue
 
 if TYPE_CHECKING:  # pragma: no cover
-    import numpy as np  # type: ignore
+    import numpy as np
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -84,7 +84,7 @@ class Histogram(WBValue):
         if len(self.histogram) + 1 != len(self.bins):
             raise ValueError("len(bins) must be len(histogram) + 1")
 
-    def to_json(self, run: Optional[Union["LocalRun", "LocalArtifact"]] = None) -> dict:
+    def to_json(self, run: Optional[Union["LocalRun", "Artifact"]] = None) -> dict:
         return {"_type": self._log_type, "values": self.histogram, "bins": self.bins}
 
     def __sizeof__(self) -> int:

wandb/sdk/data_types/html.py

@@ -10,8 +10,7 @@ from .base_types.media import BatchableMedia
 if TYPE_CHECKING:  # pragma: no cover
     from typing import TextIO
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -78,14 +77,14 @@ class Html(BatchableMedia):
     def get_media_subdir(cls: Type["Html"]) -> str:
         return os.path.join("media", "html")
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = self._log_type
         return json_dict
 
     @classmethod
     def from_json(
-        cls: Type["Html"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["Html"], json_obj: dict, source_artifact: "Artifact"
     ) -> "Html":
         return cls(source_artifact.get_path(json_obj["path"]).download(), inject=False)
 

wandb/sdk/data_types/image.py

@@ -18,12 +18,11 @@ from .helper_types.image_mask import ImageMask
 
 if TYPE_CHECKING:  # pragma: no cover
     import matplotlib  # type: ignore
-    import numpy as np  # type: ignore
+    import numpy as np
     import torch  # type: ignore
     from PIL.Image import Image as PILImage
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -309,7 +308,7 @@ class Image(BatchableMedia):
 
     @classmethod
     def from_json(
-        cls: Type["Image"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["Image"], json_obj: dict, source_artifact: "Artifact"
     ) -> "Image":
         classes: Optional[Classes] = None
         if json_obj.get("classes") is not None:
@@ -388,7 +387,7 @@ class Image(BatchableMedia):
                     run, key, step, id_, ignore_copy_err=ignore_copy_err
                 )
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = Image._log_type
         json_dict["format"] = self.format

wandb/sdk/data_types/molecule.py

@@ -15,7 +15,7 @@ if TYPE_CHECKING:  # pragma: no cover
 
     import rdkit.Chem  # type: ignore
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -205,7 +205,7 @@ class Molecule(BatchableMedia):
     def get_media_subdir(cls: Type["Molecule"]) -> str:
         return os.path.join("media", "molecule")
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = self._log_type
         if self._caption:

wandb/sdk/data_types/object_3d.py

@@ -30,9 +30,9 @@ from ._private import MEDIA_TMP
 from .base_types.media import BatchableMedia
 
 if TYPE_CHECKING:  # pragma: no cover
-    import numpy as np  # type: ignore
+    import numpy as np
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -312,7 +312,7 @@ class Object3D(BatchableMedia):
     def get_media_subdir(cls: Type["Object3D"]) -> str:
         return os.path.join("media", "object3D")
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = Object3D._log_type
 

wandb/sdk/data_types/plotly.py

@@ -15,7 +15,7 @@ if TYPE_CHECKING:  # pragma: no cover
     import pandas as pd
     import plotly  # type: ignore
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -76,7 +76,7 @@ class Plotly(Media):
     def get_media_subdir(cls: Type["Plotly"]) -> str:
         return os.path.join("media", "plotly")
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = self._log_type
         return json_dict

wandb/sdk/data_types/saved_model.py

@@ -29,8 +29,7 @@ if TYPE_CHECKING:  # pragma: no cover
     import tensorflow  # type: ignore
     import torch  # type: ignore
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -39,7 +38,7 @@ DEBUG_MODE = False
 
 
 def _add_deterministic_dir_to_artifact(
-    artifact: "LocalArtifact", dir_name: str, target_dir_root: str
+    artifact: "Artifact", dir_name: str, target_dir_root: str
 ) -> str:
     file_paths = []
     for dirpath, _, filenames in os.walk(dir_name, topdown=True):
@@ -51,7 +50,7 @@ def _add_deterministic_dir_to_artifact(
     return target_path
 
 
-def _load_dir_from_artifact(source_artifact: "PublicArtifact", path: str) -> str:
+def _load_dir_from_artifact(source_artifact: "Artifact", path: str) -> str:
     dl_path = None
 
     # Look through the entire manifest to find all of the files in the directory.
@@ -126,7 +125,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
 
     @classmethod
     def from_json(
-        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "Artifact"
     ) -> "_SavedModel":
         path = json_obj["path"]
 
@@ -144,7 +143,7 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         # and specified adapter.
         return cls(dl_path)
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         # Unlike other data types, we do not allow adding to a Run directly. There is a
         # bit of tech debt in the other data types which requires the input to `to_json`
         # to accept a Run or Artifact. However, Run additions should be deprecated in the future.
@@ -320,7 +319,7 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
 
     @classmethod
     def from_json(
-        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "PublicArtifact"
+        cls: Type["_SavedModel"], json_obj: dict, source_artifact: "Artifact"
     ) -> "_PicklingSavedModel":
         backup_path = [p for p in sys.path]
         if (
@@ -337,7 +336,7 @@ class _PicklingSavedModel(_SavedModel[SavedModelObjType]):
 
         return inst  # type: ignore
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_obj = super().to_json(run_or_artifact)
         assert isinstance(run_or_artifact, wandb.Artifact)
         if self._dep_py_files_path is not None:

wandb/sdk/data_types/trace_tree.py

@@ -13,12 +13,12 @@ from dataclasses import dataclass, field
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-from wandb.data_types import _json_helper
+import wandb.data_types
 from wandb.sdk.data_types import _dtypes
 from wandb.sdk.data_types.base_types.media import Media
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -103,7 +103,7 @@ class WBTraceTree(Media):
     def get_media_subdir(cls) -> str:
         return "media/wb_trace_tree"
 
-    def to_json(self, run: Optional[Union["LocalRun", "LocalArtifact"]]) -> dict:
+    def to_json(self, run: Optional[Union["LocalRun", "Artifact"]]) -> dict:
         res = {}
         res["_type"] = self._log_type
         # Here we use `dumps` to put things into string format. This is because
@@ -142,7 +142,7 @@ def _fallback_serialize(obj: Any) -> str:
 def _safe_serialize(obj: dict) -> str:
     try:
         return json.dumps(
-            _json_helper(obj, None),
+            wandb.data_types._json_helper(obj, None),
             skipkeys=True,
             default=_fallback_serialize,
         )

wandb/sdk/data_types/video.py

@@ -13,9 +13,9 @@ from .base_types.media import BatchableMedia
 if TYPE_CHECKING:  # pragma: no cover
     from typing import TextIO
 
-    import numpy as np  # type: ignore
+    import numpy as np
 
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     from ..wandb_run import Run as LocalRun
 
@@ -133,7 +133,7 @@ class Video(BatchableMedia):
             required='wandb.Video requires moviepy and imageio when passing raw data.  Install with "pip install moviepy imageio"',
         )
         tensor = self._prepare_video(self.data)
-        _, self._height, self._width, self._channels = tensor.shape
+        _, self._height, self._width, self._channels = tensor.shape  # type: ignore
 
         # encode sequence of images into gif string
         clip = mpy.ImageSequenceClip(list(tensor), fps=self._fps)
@@ -170,7 +170,7 @@ class Video(BatchableMedia):
     def get_media_subdir(cls: Type["Video"]) -> str:
         return os.path.join("media", "videos")
 
-    def to_json(self, run_or_artifact: Union["LocalRun", "LocalArtifact"]) -> dict:
+    def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = self._log_type