mypy+pre-commit

Author: justusschock

src/litlogger/media.py

@@ -370,24 +370,24 @@ def _maybe_as_moviepy_clip(self, data: Any) -> None | Any:
         try:
             # MoviePy 2.x layout
             video_clip_mod = import_module("moviepy.video.VideoClip")
-            VideoClip = getattr(video_clip_mod, "VideoClip")
-            if isinstance(data, VideoClip):
+            video_clip = video_clip_mod.VideoClip
+            if isinstance(data, video_clip):
                 return data
         except Exception:
             pass
 
         try:
             # Older common import path
             editor_mod = import_module("moviepy.editor")
-            VideoClip = getattr(editor_mod, "VideoClip")
-            if isinstance(data, VideoClip):
+            video_clip = editor_mod.VideoClip
+            if isinstance(data, video_clip):
                 return data
         except Exception:
             pass
 
         return None
 
-    def _moviepy_clip_from_array(self, data, fps: float):
+    def _moviepy_clip_from_array(self, data: Any, fps: float) -> Any:
         np = import_module("numpy")
 
         if data.dtype != np.uint8:
@@ -402,8 +402,7 @@ def _moviepy_clip_from_array(self, data, fps: float):
 
         if data.ndim not in (3, 4):
             raise ValueError(
-                f"Unsupported array shape for video: {data.shape}. "
-                "Expected (T,H,W), (T,H,W,C), or (T,C,H,W)."
+                f"Unsupported array shape for video: {data.shape}. Expected (T,H,W), (T,H,W,C), or (T,C,H,W)."
             )
 
         # (T, H, W) -> grayscale => expand to (T, H, W, 1)
@@ -420,34 +419,26 @@ def _moviepy_clip_from_array(self, data, fps: float):
             data = np.repeat(data, 3, axis=-1)
 
         if data.shape[-1] not in (3, 4):
-            raise ValueError(
-                f"Unsupported channel count for video frames: {data.shape[-1]}"
-            )
+            raise ValueError(f"Unsupported channel count for video frames: {data.shape[-1]}")
 
         # Drop alpha for now unless you explicitly want to preserve/use masks.
         if data.shape[-1] == 4:
             data = data[..., :3]
 
         try:
             # MoviePy 2.x
-            ImageSequenceClip = getattr(
-                import_module("moviepy.video.io.ImageSequenceClip"),
-                "ImageSequenceClip",
-            )
+            image_sequence_clip = import_module("moviepy.video.io.ImageSequenceClip").ImageSequenceClip
         except Exception:
             # Older common path
-            ImageSequenceClip = getattr(
-                import_module("moviepy.editor"),
-                "ImageSequenceClip",
-            )
+            image_sequence_clip = import_module("moviepy.editor").ImageSequenceClip
 
         # list(...) avoids some ndarray edge cases in callers and matches
         # common usage for frame sequences.
-        return ImageSequenceClip(list(data), fps=fps)
+        return image_sequence_clip(list(data), fps=fps)
 
-    def _write_moviepy_clip(self, clip, path: str, fps: float) -> None:
+    def _write_moviepy_clip(self, clip: Any, path: str, fps: float) -> None:
         # For MP4, libx264 is the usual sensible default.
-        kwargs = {
+        kwargs: dict[str, Any] = {
             "fps": fps,
             "logger": None,
         }

tests/integrations/test_standalone_media.py

@@ -377,3 +377,63 @@ def test_new_dict_api_file_download():
             project_id=project_id,
             body=LitLoggerServiceDeleteMetricsStreamBody(ids=[stream_id]),
         )
+
+
+@pytest.mark.cloud()
+def test_new_dict_api_video_download():
+    """Test the new dict-like API for uploading and downloading videos."""
+    from litlogger import Video
+
+    experiment_name = f"standalone_dict_video_test-{uuid.uuid4().hex}"
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        exp = litlogger.init(
+            name=experiment_name,
+            teamspace="oss-litlogger",
+            root_dir=tmpdir,
+        )
+
+        video_path = os.path.join(tmpdir, "preview.mp4")
+        with open(video_path, "wb") as f:
+            f.write(b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isom")
+
+        exp["preview"] = Video(video_path)
+
+        assert isinstance(exp["preview"], Video)
+        assert exp["preview"].name == "preview"
+
+        litlogger.finalize()
+
+        exp2 = litlogger.init(name=experiment_name, teamspace="oss-litlogger")
+
+        download_path = os.path.join(tmpdir, "downloaded_preview.mp4")
+        video_downloaded = False
+        last_exception = None
+        for attempt in range(30):
+            try:
+                preview = exp2["preview"]
+                if isinstance(preview, Video):
+                    preview.save(download_path)
+                    video_downloaded = True
+                    break
+            except Exception as e:
+                last_exception = e
+            if attempt < 29:
+                sleep(1)
+
+        litlogger.finalize()
+
+        if video_downloaded:
+            assert os.path.exists(download_path)
+            with open(video_path, "rb") as original, open(download_path, "rb") as downloaded:
+                assert downloaded.read() == original.read()
+        else:
+            print(f"\nWarning: Could not download video. Last exception: {last_exception}")
+
+        project_id = exp._teamspace.id
+        stream_id = exp._metrics_store.id
+        client = LitRestClient()
+        client.lit_logger_service_delete_metrics_stream(
+            project_id=project_id,
+            body=LitLoggerServiceDeleteMetricsStreamBody(ids=[stream_id]),
+        )

tests/unittests/test_experiment_media.py

@@ -11,7 +11,7 @@
 import pytest
 from lightning_sdk.lightning_cloud.openapi import V1MediaType
 from litlogger.experiment import Experiment
-from litlogger.media import File, Image, Model, Text
+from litlogger.media import File, Image, Model, Text, Video
 from litlogger.series import Series
 
 
@@ -92,6 +92,15 @@ def test_setitem_text(self):
         assert exp._static_files["notes"] is t
         exp._set_static_file.assert_called_once_with("notes", t)
 
+    def test_setitem_video(self):
+        exp = _make_exp()
+        video = Video("preview.mp4")
+        exp["preview"] = video
+
+        assert exp._key_types["preview"] == "static_file"
+        assert exp._static_files["preview"] is video
+        exp._set_static_file.assert_called_once_with("preview", video)
+
     def test_overwrite_same_type(self):
         """Overwriting a static_file key with another File is allowed."""
         exp = _make_exp()
@@ -188,6 +197,48 @@ def test_non_file_media_uses_media_api(self):
         assert kwargs["media_type"] == V1MediaType.IMAGE
         assert exp._stats.media_logged == 1
 
+    def test_video_media_uses_media_api(self):
+        exp = MagicMock(spec=Experiment)
+        exp._media_api = MagicMock()
+        exp._metrics_store = MagicMock()
+        exp._metrics_store.id = "store-1"
+        exp._teamspace = MagicMock()
+        exp._stats = MagicMock()
+        exp._stats.media_logged = 0
+        exp._media_type_to_v1 = lambda media_type: Experiment._media_type_to_v1(exp, media_type)
+        exp._upload_media = (
+            lambda name, file_path, media_type, step=None, epoch=None, caption=None: Experiment._upload_media(
+                exp,
+                name,
+                file_path,
+                media_type,
+                step=step,
+                epoch=epoch,
+                caption=caption,
+            )
+        )
+        exp._upload_media_value = (
+            lambda key, value, name=None, step=None, epoch=None, caption=None: Experiment._upload_media_value(
+                exp,
+                key,
+                value,
+                name=name,
+                step=step,
+                epoch=epoch,
+                caption=caption,
+            )
+        )
+
+        video = Video("preview.mp4")
+        Experiment._set_static_file(exp, "preview", video)
+
+        exp._media_api.upload_media.assert_called_once()
+        _, kwargs = exp._media_api.upload_media.call_args
+        assert kwargs["name"] == "preview"
+        assert kwargs["file_path"] == "preview.mp4"
+        assert kwargs["media_type"] == V1MediaType.VIDEO
+        assert exp._stats.media_logged == 1
+
     @patch.object(Model, "_log_model", return_value="owner/team/exp-model:latest")
     def test_model_artifact_uses_litmodels(self, mock_log_model):
         exp = Experiment.__new__(Experiment)
@@ -208,6 +259,7 @@ def test_model_artifact_uses_litmodels(self, mock_log_model):
         mock_log_model.assert_called_once_with(
             experiment_name="exp",
             teamspace=exp._teamspace,
+            key="checkpoint",
             experiment=exp,
             cloud_account="acc-1",
         )
@@ -235,6 +287,7 @@ def test_model_object_uses_litmodels(self, mock_log_model):
         mock_log_model.assert_called_once_with(
             experiment_name="exp",
             teamspace=exp._teamspace,
+            key="model-object",
             experiment=exp,
             cloud_account="acc-1",
         )
@@ -300,6 +353,14 @@ def test_append_text_to_series(self):
         assert len(exp["logs"]) == 1
         assert exp._key_types["logs"] == "file_series"
 
+    def test_append_video_to_series(self):
+        exp = _make_exp()
+        video = Video("preview.mp4")
+        exp["clips"].append(video)
+
+        assert len(exp["clips"]) == 1
+        assert exp._key_types["clips"] == "file_series"
+
 
 class TestFileSeriesBindings:
     """Test that _log_file_series_value binds name and _download_fn."""
@@ -400,6 +461,48 @@ def test_non_file_series_uses_media_api(self):
         assert kwargs["media_type"] == V1MediaType.TEXT
         assert exp._stats.media_logged == 1
 
+    def test_video_series_uses_media_api(self):
+        exp = MagicMock(spec=Experiment)
+        exp._media_api = MagicMock()
+        exp._metrics_store = MagicMock()
+        exp._metrics_store.id = "store-1"
+        exp._teamspace = MagicMock()
+        exp._stats = MagicMock()
+        exp._stats.media_logged = 0
+        exp._media_type_to_v1 = lambda media_type: Experiment._media_type_to_v1(exp, media_type)
+        exp._upload_media = (
+            lambda name, file_path, media_type, step=None, epoch=None, caption=None: Experiment._upload_media(
+                exp,
+                name,
+                file_path,
+                media_type,
+                step=step,
+                epoch=epoch,
+                caption=caption,
+            )
+        )
+        exp._upload_media_value = (
+            lambda key, value, name=None, step=None, epoch=None, caption=None: Experiment._upload_media_value(
+                exp,
+                key,
+                value,
+                name=name,
+                step=step,
+                epoch=epoch,
+                caption=caption,
+            )
+        )
+
+        video = Video("preview.mp4")
+        Experiment._log_file_series_value(exp, "clips", video, 2, step=7)
+
+        exp._media_api.upload_media.assert_called_once()
+        _, kwargs = exp._media_api.upload_media.call_args
+        assert kwargs["name"] == "clips"
+        assert kwargs["step"] == 7
+        assert kwargs["media_type"] == V1MediaType.VIDEO
+        assert exp._stats.media_logged == 1
+
     @patch.object(Model, "_log_model", return_value="owner/team/exp-model-series:latest")
     def test_model_series_uses_series_key_for_remote_binding(self, mock_log_model):
         exp = Experiment.__new__(Experiment)
@@ -420,6 +523,7 @@ def test_model_series_uses_series_key_for_remote_binding(self, mock_log_model):
         mock_log_model.assert_called_once_with(
             experiment_name="exp",
             teamspace=exp._teamspace,
+            key="models",
             experiment=exp,
             cloud_account="acc-1",
         )
@@ -460,6 +564,13 @@ def test_getitem_returns_text(self):
 
         assert exp["notes"] is t
 
+    def test_getitem_returns_video(self):
+        exp = _make_exp()
+        video = Video("preview.mp4")
+        exp["preview"] = video
+
+        assert exp["preview"] is video
+
     def test_save_without_upload_raises(self):
         """File.save() fails if not yet uploaded (no _download_fn)."""
         f = File("local.txt")
@@ -749,6 +860,42 @@ def test_rebuilds_static_media_with_wrapper(self):
         assert wrapped.name == "preview"
         assert wrapped._download_fn is not None
 
+    def test_rebuilds_static_video_with_wrapper(self):
+        media = MagicMock()
+        media.name = "preview"
+        media.storage_path = "media/preview.mp4"
+        media.cluster_id = "cloud-1"
+        media.media_type = V1MediaType.VIDEO
+        media.id = "media-1"
+
+        exp = MagicMock(spec=Experiment)
+        exp._key_types = {}
+        exp._metadata_values = {}
+        exp._static_files = {}
+        exp._series = {}
+        exp._metrics_store = MagicMock()
+        exp._metrics_store.id = "store-1"
+        exp._update_metrics_store = MagicMock()
+        exp._metrics_store.tags = []
+        exp._metrics_store.artifacts = []
+        exp._metrics_api = MagicMock()
+        exp._teamspace = MagicMock()
+        exp._teamspace.id = "ts-1"
+        exp._media_api = MagicMock()
+        exp._media_api.client.lit_logger_service_list_lit_logger_media.return_value.media = [media]
+        exp._wrap_media_file = lambda media_name, media_type: Experiment._wrap_media_file(exp, media_name, media_type)
+        exp._create_media_download_fn = lambda storage_path, cloud_account=None: Experiment._create_media_download_fn(
+            exp, storage_path, cloud_account
+        )
+        exp._resumed_steps = {}
+
+        Experiment._rebuild_state(exp)
+
+        wrapped = exp._static_files["preview"]
+        assert isinstance(wrapped, Video)
+        assert wrapped.name == "preview"
+        assert wrapped._download_fn is not None
+
     def test_rebuilds_media_series_with_wrapper(self):
         media0 = MagicMock()
         media0.name = "logs/0"