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 @@ 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 @@ -29,6 +28,13 @@ def __init__( 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: 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"))