diff --git a/simvue/client.py b/simvue/client.py index 915e9db0..8b184fed 100644 --- a/simvue/client.py +++ b/simvue/client.py @@ -32,7 +32,7 @@ from .models import FOLDER_REGEX, NAME_REGEX from .config.user import SimvueConfiguration from .api.request import get_json_from_response -from .api.objects import Run, Folder, Tag, Artifact, Alert +from .api.objects import Run, Folder, Tag, Artifact, Alert, FileArtifact, ObjectArtifact CONCURRENT_DOWNLOADS = 10 @@ -42,8 +42,10 @@ def _download_artifact_to_file( - artifact: Artifact, output_dir: pathlib.Path | None + artifact: FileArtifact | ObjectArtifact, output_dir: pathlib.Path | None ) -> None: + if not artifact.name: + raise RuntimeError(f"Expected artifact '{artifact.id}' to have a name") _output_file = (output_dir or pathlib.Path.cwd()).joinpath(artifact.name) # If this is a hierarchical structure being downloaded, need to create directories _output_file.parent.mkdir(parents=True, exist_ok=True) @@ -475,13 +477,14 @@ def list_artifacts( ) # type: ignore def _retrieve_artifacts_from_server( - self, run_id: str, name: str, count: int | None = None - ) -> typing.Generator[tuple[str, Artifact], None, None]: - return Artifact.get( - runs=json.dumps([run_id]), - filters=json.dumps([f"name == {name}"]), - count=count, - ) # type: ignore + self, run_id: str, name: str + ) -> FileArtifact | ObjectArtifact | None: + return Artifact.from_name( + run_id=run_id, + name=name, + server_url=self._user_config.server.url, + server_token=self._user_config.server.token, + ) @prettify_pydantic @pydantic.validate_call @@ -529,12 +532,14 @@ def get_artifact( RuntimeError if retrieval of artifact from the server failed """ - _artifact = Artifact.from_name( - run_id=run_id, - name=name, - server_url=self._user_config.server.url, - server_token=self._user_config.server.token, - ) + _artifact = self._retrieve_artifacts_from_server(run_id, name) + + if not _artifact: + raise ObjectNotFoundError( + obj_type="artifact", + name=name, + extra=f"for run '{run_id}'", + ) _content = b"".join(_artifact.download_content()) @@ -574,12 +579,14 @@ def get_artifact_as_file( if there was a failure during retrieval of information from the server """ - _artifacts = self._retrieve_artifacts_from_server(run_id, name, count=1) + _artifact = self._retrieve_artifacts_from_server(run_id, name) - try: - _id, _artifact = next(_artifacts) - except StopIteration as e: - raise ValueError(f"No artifact '{name}' found for run '{run_id}'") from e + if not _artifact: + raise ObjectNotFoundError( + obj_type="artifact", + name=name, + extra=f"for run '{run_id}'", + ) _download_artifact_to_file(_artifact, output_dir)