Don't invoke ffmpeg if an exception occured during looming

matplotloom/loom.py

@@ -12,7 +12,7 @@ class Loom:
     def __init__(
         self,
         output_filepath: Union[Path, str],
-        frames_directory: Union[Path, str] = Path(TemporaryDirectory().name),
+        frames_directory: Union[Path, str, None] = None,
         fps: int = 30,
         keep_frames: bool = False,
         overwrite: bool = False,
@@ -21,14 +21,20 @@ def __init__(
         savefig_kwargs: dict = {}
     ) -> None:
         self.output_filepath = Path(output_filepath)
-        self.frames_directory = Path(frames_directory)
         self.fps = fps
         self.keep_frames = keep_frames
         self.overwrite = overwrite
         self.verbose = verbose
         self.parallel = parallel
         self.savefig_kwargs = savefig_kwargs
 
+        self._temp_dir = None
+        if frames_directory is None:
+            self._temp_dir = TemporaryDirectory()
+            self.frames_directory = Path(self._temp_dir.name)
+        else:
+            self.frames_directory = Path(frames_directory)
+
         if not self.parallel:
             self.frame_counter = 0
         else:
@@ -53,8 +59,23 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.save_video()
-        return
+        try:
+            if exc_type is None:
+                self.save_video()
+            else:
+                if self.verbose:
+                    print(f"An error occurred: {exc_type.__name__}: {exc_value}")
+                    print("Animation was not saved.")
+        finally:
+            if not self.keep_frames:
+                for frame_filepath in self.frame_filepaths:
+                    if frame_filepath.exists():
+                        frame_filepath.unlink()
+
+            if self._temp_dir:
+                self._temp_dir.cleanup()
+
+        return False  # Propagate the exception if there was one
 
     def save_frame(self, fig, frame_number=None):
         if not self.parallel:

tests/test_loom.py

@@ -1,64 +1,88 @@
 import pytest
 
+import matplotlib.pyplot as plt
+
 from pathlib import Path
 from matplotlib.figure import Figure
 from matplotloom import Loom
 
-@pytest.fixture
-def basic_loom(tmp_path):
+def test_init_loom(tmp_path):
     output_filepath = tmp_path / "output.mp4"
-    return Loom(output_filepath=output_filepath, verbose=True)
+    loom = Loom(output_filepath=output_filepath, verbose=True)
 
-@pytest.fixture
-def loom_with_custom_dir(tmp_path):
-    output_filepath = tmp_path / "output.mp4"
-    frames_directory = tmp_path / "frames"
-    return Loom(output_filepath=output_filepath, frames_directory=frames_directory, verbose=True)
+    assert loom.output_directory.exists()
+    assert loom.frames_directory.exists()
+
+def test_init_loom_with_custom_frames_directory(tmp_path):
+    output_filepath = tmp_path / "output.gif"
+    custom_frames_dir = tmp_path / "frames"
+
+    loom_with_custom_dir = Loom(
+        output_filepath=output_filepath,
+        frames_directory=custom_frames_dir,
+        verbose=True
+    )
 
-def test_init(basic_loom, loom_with_custom_dir):
-    assert basic_loom.output_directory.exists()
-    assert basic_loom.frames_directory.exists()
+    assert loom_with_custom_dir.output_directory.exists()
     assert loom_with_custom_dir.frames_directory.exists()
 
-# Test the saving of a single frame
-def test_save_frame(basic_loom):
+def test_save_frame(tmp_path):
+    output_filepath = tmp_path / "output.mp4"
+    loom = Loom(output_filepath=output_filepath, verbose=True)
+
     fig = Figure()
     ax = fig.subplots()
     ax.plot([0, 1], [0, 1])
 
-    basic_loom.save_frame(fig)
-    assert len(basic_loom.frame_filepaths) == 1
-    assert basic_loom.frame_filepaths[0].exists()
+    loom.save_frame(fig)
+    assert len(loom.frame_filepaths) == 1
+    assert loom.frame_filepaths[0].exists()
 
-# This test might require a mock of subprocess.Popen or an actual ffmpeg installation
-def test_video_creation(basic_loom):
+def test_video_creation(tmp_path):
+    output_filepath = tmp_path / "output.mp4"
+    loom = Loom(output_filepath=output_filepath, verbose=True)
+
     fig = Figure()
     ax = fig.subplots()
     ax.plot([0, 1], [0, 1])
 
-    basic_loom.save_frame(fig)
-    basic_loom.save_video()
-    assert basic_loom.output_filepath.exists()
+    loom.save_frame(fig)
+    loom.save_video()
+    assert loom.output_filepath.exists()
 
-# Ensure frames are kept or deleted according to the keep_frames flag
-def test_keep_frames(basic_loom):
+def test_dont_keep_frames(tmp_path):
+    output_filepath = tmp_path / "output.mp4"
+    loom = Loom(
+        output_filepath=output_filepath,
+        keep_frames=False,
+        verbose=True
+    )
+
     fig = Figure()
     ax = fig.subplots()
     ax.plot([0, 1], [0, 1])
 
-    basic_loom.save_frame(fig)
-    basic_loom.save_video()
+    loom.save_frame(fig)
+    loom.save_video()
+
+    assert not loom.frame_filepaths[0].exists()
+
+def test_keep_frames(tmp_path):
+    output_filepath = tmp_path / "output.mp4"
+    loom = Loom(
+        output_filepath=output_filepath,
+        keep_frames=True,
+        verbose=True
+    )
 
-    # Check if frames are deleted (default behavior)
-    assert not basic_loom.frame_filepaths[0].exists()
+    fig = Figure()
+    ax = fig.subplots()
+    ax.plot([0, 1], [0, 1])
 
-    # Test with keep_frames=True
-    basic_loom.keep_frames = True
-    basic_loom.save_frame(fig)
-    basic_loom.save_video()
+    loom.save_frame(fig)
+    loom.save_video()
 
-    # Check if frames are kept
-    assert basic_loom.frame_filepaths[1].exists()
+    assert loom.frame_filepaths[0].exists()
 
 def test_context_manager(tmp_path):
     output_filepath = tmp_path / "output.mp4"
@@ -68,28 +92,39 @@ def test_context_manager(tmp_path):
         ax.plot([0, 1], [0, 1])
         loom.save_frame(fig)
 
-    # After exiting the context, the video file should exist
     assert output_filepath.exists()
 
 def test_overwrite(tmp_path):
     output_filepath = tmp_path / "output.mp4"
 
-    # Create a dummy file to simulate an existing output file
     with open(output_filepath, "w") as f:
         f.write("Dummy content")
 
-    # Test with overwrite=False (default)
     with pytest.raises(FileExistsError):
         Loom(output_filepath)
 
-    # Test with overwrite=True
     loom = Loom(output_filepath, overwrite=True)
+
     fig = Figure()
     ax = fig.subplots()
     ax.plot([0, 1], [0, 1])
     loom.save_frame(fig)
     loom.save_video()
 
-    # Check if the output file was overwritten
     assert output_filepath.exists()
     assert output_filepath.stat().st_size > len("Dummy content")
+
+def test_loom_error_handling(tmp_path):
+    output_file = tmp_path / "test_error.mp4"
+
+    with pytest.raises(ValueError, match="Test error"):
+        with Loom(output_file, verbose=True) as loom:
+            fig, ax = plt.subplots()
+            ax.plot([1, 2, 3], [1, 2, 3])
+            loom.save_frame(fig)
+            raise ValueError("Test error")
+
+    assert not output_file.exists()
+
+    frames_dir = Path(loom.frames_directory)
+    assert not any(frames_dir.glob("frame_*.png"))