diff --git a/src/icepack/helper.py b/src/icepack/helper.py index 7251c4a..354933a 100644 --- a/src/icepack/helper.py +++ b/src/icepack/helper.py @@ -285,6 +285,13 @@ def __init__(self, path, mode='r'): else: self._entries = {} + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, exc_traceback): + failed = exc_type is not None + self.close(silent=failed) + def add_entry(self, key, path): """Add an entry with the content of path (may be None).""" if self._mode != 'w': diff --git a/tests/test_helper_zip.py b/tests/test_helper_zip.py index 4c86e4e..4958507 100644 --- a/tests/test_helper_zip.py +++ b/tests/test_helper_zip.py @@ -48,6 +48,16 @@ def test_add_entry_after_metdata(self, src_path, zip_file): with pytest.raises(InvalidArchiveError): zip_file.add_entry('bar', foo) + def test_context_manager(self, src_path, zip_path): + with Zip(zip_path, mode='w') as zip_file: + zip_file.add_entry('foo', src_path / 'foo') + meta_path = src_path / 'qux' / 'quux' + zip_file.add_metadata(meta_path, meta_path) + temp_dir = zip_file._temp_dir + assert temp_dir.exists() + assert not temp_dir.exists() + assert zip_path.exists() + class TestReadMode: """Test read operations.""" @@ -66,3 +76,12 @@ def test_regular_zip(self, shared_datadir): path = shared_datadir / 'zips' / 'infozip.zip' with pytest.raises(InvalidArchiveError): zip_file = Zip(path) + + def test_context_manager(self, shared_datadir): + path = shared_datadir / 'zips' / 'zip-helper.zip' + with Zip(path) as zip_file: + path = zip_file.extract_entry('foo') + assert File.sha256(path) == 'b5bb9d8014a0f9b1d61e21e796d78dccdf1352f23cd32812f4850b878ae4944c' # noqa + temp_dir = zip_file._temp_dir + assert temp_dir.exists() + assert not temp_dir.exists()