diff --git a/great_ai/large_file/large_file/large_file_base.py b/great_ai/large_file/large_file/large_file_base.py index ce0631e..22a12fc 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/great_ai/large_file/large_file/large_file_base.py @@ -146,11 +146,13 @@ class LargeFileBase(ABC): ) -> Literal[False]: self._file.close() - if type is None: - if "w" in self._mode: - self.push(Path(self._file.name)) + if "w" in self._mode: + try: + if type is None: + self.push(Path(self._file.name)) + finally: os.unlink(self._file.name) - else: + elif type is not None: logger.exception("Could not finish operation.") return False diff --git a/tests/large_file/test_storage_safety.py b/tests/large_file/test_storage_safety.py index 0c4ba32..7ebceba 100644 --- a/tests/large_file/test_storage_safety.py +++ b/tests/large_file/test_storage_safety.py @@ -56,3 +56,18 @@ def test_cache_quota_counts_directory_contents( monkeypatch.setattr(LargeFileLocal, "max_cache_size", "1 KB") LargeFileLocal("model", "w").push(source) assert not (LargeFileLocal.cache_path / "model-0").exists() + + +def test_failed_write_removes_temporary_file() -> None: + temporary = None + try: + with pytest.raises(ValueError, match="failed"): + with LargeFileLocal("model", "w") as file: + temporary = Path(file.name) + file.write("unfinished") + raise ValueError("failed") + assert temporary is not None + assert not temporary.exists() + finally: + if temporary is not None: + temporary.unlink(missing_ok=True)