Remove temporary model files after failed writes

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:
-            if "w" in self._mode:
+            try:
-                self.push(Path(self._file.name))
+                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

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)