Write formatted files atomically to avoid corruption on write failure

CHANGES.md

@@ -8,6 +8,10 @@
 - Add support for NO_COLOR environment variable to disable ANSI output (#5129)
 - No spurious target version warning when runtime version is included in a
   --target-version flag (#5167)
+- Fix in-place formatting truncating or corrupting the source file if writing the
+  reformatted contents fails partway through, e.g. because the disk is full. Black now
+  writes to a temporary file in the same directory and atomically replaces the original
+  once the write has fully succeeded (#5207)
 
 ### Highlights
 

src/black/__init__.py

@@ -1,8 +1,11 @@
 import io
 import json
+import os
 import platform
 import re
+import shutil
 import sys
+import tempfile
 import tokenize
 import traceback
 from collections.abc import (
@@ -998,8 +1001,27 @@ def format_file_in_place(
     dst_contents = header.decode(encoding) + dst_contents
 
     if write_back == WriteBack.YES:
-        with open(src, "w", encoding=encoding, newline=newline) as f:
-            f.write(dst_contents)
+        # Write to a temporary file in the same directory (so it's on the same
+        # filesystem) and atomically replace the target with it. This avoids
+        # leaving the file truncated or corrupted if writing fails partway
+        # through, e.g. because the disk is full (see issue #2479).
+        # If `src` is itself a symlink, write through to its resolved target
+        # instead of replacing the symlink with a regular file.
+        real_src = src.resolve() if src.is_symlink() else src
+        fd, tmp_file = tempfile.mkstemp(
+            dir=real_src.parent, prefix=f".{real_src.name}.", suffix=".tmp"
+        )
+        tmp_path = Path(tmp_file)
+        try:
+            with os.fdopen(fd, "w", encoding=encoding, newline=newline) as f:
+                f.write(dst_contents)
+                f.flush()
+                os.fsync(f.fileno())
+            shutil.copymode(real_src, tmp_path)
+            os.replace(tmp_path, real_src)
+        except BaseException:
+            tmp_path.unlink(missing_ok=True)
+            raise
     elif write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
         now = datetime.now(timezone.utc)
         src_name = f"{src}\t{then}"

tests/test_black.py

@@ -170,6 +170,56 @@ def test_one_empty_line_ff(self) -> None:
                 os.unlink(tmp_file)
             self.assertFormatEqual(expected, actual)
 
+    def test_ff_preserves_file_on_write_failure(self) -> None:
+        # Regression test for #2479: if writing the reformatted contents fails
+        # partway through (e.g. because the disk is full), the original file
+        # must be left untouched instead of being truncated or corrupted.
+        source = "print('hello')\n"
+        with TemporaryDirectory() as workspace:
+            tmp_file = Path(workspace) / "source.py"
+            tmp_file.write_text(source, encoding="utf-8")
+            with patch(
+                "black.os.fsync", side_effect=OSError("No space left on device")
+            ):
+                with self.assertRaises(OSError):
+                    ff(tmp_file, write_back=black.WriteBack.YES)
+            self.assertEqual(tmp_file.read_text(encoding="utf-8"), source)
+            # No leftover temporary file should remain in the directory.
+            leftover = [p for p in Path(workspace).iterdir() if p != tmp_file]
+            self.assertEqual(leftover, [])
+
+    def test_ff_preserves_permissions(self) -> None:
+        if system() == "Windows":
+            return
+
+        with TemporaryDirectory() as workspace:
+            tmp_file = Path(workspace) / "source.py"
+            tmp_file.write_text("print('hello' )\n", encoding="utf-8")
+            tmp_file.chmod(0o640)
+            self.assertTrue(ff(tmp_file, write_back=black.WriteBack.YES))
+            self.assertEqual(tmp_file.stat().st_mode & 0o777, 0o640)
+
+    def test_ff_writes_through_symlink(self) -> None:
+        # Regression test: formatting a file via a symlinked path must write
+        # through to the real target, not replace the symlink with a regular
+        # file (which atomic replace-by-rename would otherwise do).
+        with TemporaryDirectory() as workspace:
+            workspace_path = Path(workspace)
+            real_file = workspace_path / "real.py"
+            real_file.write_text("print('hello' )\n", encoding="utf-8")
+            symlink = workspace_path / "link.py"
+            try:
+                symlink.symlink_to(real_file)
+            except (OSError, NotImplementedError) as e:
+                self.skipTest(f"Can't create symlinks: {e}")
+
+            self.assertTrue(ff(symlink, write_back=black.WriteBack.YES))
+            self.assertTrue(symlink.is_symlink())
+            # Compare resolved paths rather than the raw readlink() target, since
+            # Windows may normalize it to an extended-length \\?\ path.
+            self.assertEqual(symlink.resolve(), real_file.resolve())
+            self.assertEqual(real_file.read_text(encoding="utf-8"), 'print("hello")\n')
+
     def test_piping(self) -> None:
         _, source, expected = read_data_from_file(
             PROJECT_ROOT / "src/black/__init__.py"