From f629b16cd2a91070508e776e33168c83c2714d67 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 27 Oct 2026 16:53:49 +0000 Subject: [PATCH] Remove temporary model files after failed writes --- great_ai/large_file/large_file/large_file_base.py | 10 ++++++---- tests/large_file/test_storage_safety.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) 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)