Repository: ali-ramadhan/matplotloom

Issue:
No video should be produced if exit occurs because of an error or exception
```
min/max: 8.4485 seconds.
1 24 0 2
output_filepath: test_ar.mp4
frames_directory: /tmp/tmpn9xoznqw
1
2
3
ffmpeg -y -framerate 5 -i /tmp/tmpn9xoznqw/frame_%06d.png -vf scale='if(mod(iw,2),-2,iw)':'if(mod(ih,2),-2,ih)':flags=lanczos -c:v libx264 -pix_fmt yuv420p test_ar.mp4

ffmpeg version 4.3.6-0+deb11u1 Copyright (c) 2000-2023 the FFmpeg developers
  built with gcc 10 (Debian 10.2.1-6)
  configuration: --prefix=/usr --extra-version=0+deb11u1 --toolchain=hardened --libdir=/usr/lib/x86_64-linux-gnu --incdir=/usr/include/x86_64-linux-gnu --arch=amd64 --enable-gpl --disable-stripping --enable-avresample --disable-filter=resample --enable-gnutls --enable-ladspa --enable-libaom --enable-libass --enable-libbluray --enable-libbs2b --enable-libcaca --enable-libcdio --enable-libcodec2 --enable-libdav1d --enable-libflite --enable-libfontconfig --enable-libfreetype --enable-libfribidi --enable-libgme --enable-libgsm --enable-libjack --enable-libmp3lame --enable-libmysofa --enable-libopenjpeg --enable-libopenmpt --enable-libopus --enable-libpulse --enable-librabbitmq --enable-librsvg --enable-librubberband --enable-libshine --enable-libsnappy --enable-libsoxr --enable-libspeex --enable-libsrt --enable-libssh --enable-libtheora --enable-libtwolame --enable-libvidstab --enable-libvorbis --enable-libvpx --enable-libwavpack --enable-libwebp --enable-libx265 --enable-libxml2 --enable-libxvid --enable-libzmq --enable-libzvbi --enable-lv2 --enable-omx --enable-openal --enable-opencl --enable-opengl --enable-sdl2 --enable-pocketsphinx --enable-libmfx --enable-libdc1394 --enable-libdrm --enable-libiec61883 --enable-chromaprint --enable-frei0r --enable-libx264 --enable-shared
  libavutil      56. 51.100 / 56. 51.100
  libavcodec     58. 91.100 / 58. 91.100
  libavformat    58. 45.100 / 58. 45.100
  libavdevice    58. 10.100 / 58. 10.100
  libavfilter     7. 85.100 /  7. 85.100
  libavresample   4.  0.  0 /  4.  0.  0
  libswscale      5.  7.100 /  5.  7.100
  libswresample   3.  7.100 /  3.  7.100
  libpostproc    55.  7.100 / 55.  7.100
[image2 @ 0x5556dbfb8240] Could find no file with path '/tmp/tmpn9xoznqw/frame_%06d.png' and index in the range 0-4
/tmp/tmpn9xoznqw/frame_%06d.png: No such file or directory

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[113], line 1
----> 1 plot_autoregressive_hyperparameter_search_keeling(time_series_darts["keeling"], study.trials)

Cell In[112], line 33, in plot_autoregressive_hyperparameter_search_keeling(ts, trials, output_filepath)
     29 gs = GridSpec(2, 2, figure=fig)
     31 print(3)
---> 33 gs.update(width_ratios=[[3, 1], [1, 1]])
     35 ax1 = fig.add_subplot(gs[0, 0])
     36 ax2 = fig.add_subplot(gs[0, 1])

File ~/time-series-forecasting/.venv/lib/python3.12/site-packages/matplotlib/gridspec.py:403, in GridSpec.update(self, **kwargs)
    401         setattr(self, k, v)
    402     else:
--> 403         raise AttributeError(f"{k} is an unknown keyword")
    404 for figmanager in _pylab_helpers.Gcf.figs.values():
    405     for ax in figmanager.canvas.figure.axes:

AttributeError: width_ratios is an unknown keyword
```

Discussion:

Patch:
diff --git a/matplotloom/loom.py b/matplotloom/loom.py
index 0b289dd..84c9f70 100644
--- a/matplotloom/loom.py
+++ b/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,7 +21,6 @@ class Loom:
         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
@@ -29,6 +28,13 @@ class Loom:
         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 @@ class Loom:
         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:


Test patch:
diff --git a/tests/test_loom.py b/tests/test_loom.py
index bbc45e6..a6fa34f 100644
--- a/tests/test_loom.py
+++ b/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"))