diff --git a/client/starwhale/core/dataset/type.py b/client/starwhale/core/dataset/type.py index 9b663d1ee1..7d16253e22 100644 --- a/client/starwhale/core/dataset/type.py +++ b/client/starwhale/core/dataset/type.py @@ -165,6 +165,7 @@ class ArtifactType(Enum): Video = "video" Audio = "audio" Text = "text" + Link = "link" _TBAType = t.TypeVar("_TBAType", bound="BaseArtifact") @@ -213,6 +214,8 @@ def reflect(cls, raw_data: bytes, data_type: t.Dict[str, t.Any]) -> BaseArtifact return Audio(raw_data, mime_type=mime_type, shape=shape) elif not dtype or dtype == ArtifactType.Binary.value: return Binary(raw_data) + elif dtype == ArtifactType.Link.value: + return cls.reflect(raw_data, data_type["data_type"]) else: raise NoSupportError(f"Artifact reflect error: {data_type}") @@ -436,7 +439,7 @@ def __init__( data_type: t.Optional[BaseArtifact] = None, with_local_fs_data: bool = False, ) -> None: - self.type = "link" + self.type = ArtifactType.Link self.uri = uri.strip() self.offset = offset self.size = size diff --git a/client/tests/sdk/test_dataset.py b/client/tests/sdk/test_dataset.py index 633254eb33..f24d222e9f 100644 --- a/client/tests/sdk/test_dataset.py +++ b/client/tests/sdk/test_dataset.py @@ -317,3 +317,8 @@ def test_reflect(self) -> None: with self.assertRaises(NoSupportError): BaseArtifact.reflect(b"", data_type={"type": 1}) + + link_audio = BaseArtifact.reflect( + b"link", data_type={"type": "link", "data_type": {"type": "audio"}} + ) + assert isinstance(link_audio, Audio)