llvm_tools: Add atomic_write to patch_utils.py

This allows a utility function which can write to an
arbitrary file without risking an incomplete write error,
causing issues and creating an invalid edit.

This function ensures that the file is only swapped if the
file write was successful, otherwise the file is deleted.

BUG=None
TEST=./patch_utils_unittest.py

Change-Id: Iedc3297b0e59d216f027e6ff125f92bc4d088c7d
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/3685569
Reviewed-by: George Burgess <gbiv@chromium.org>
Commit-Queue: Jordan Abrahams-Whitehead <ajordanr@google.com>
Tested-by: Jordan Abrahams-Whitehead <ajordanr@google.com>

2 files changed, 60 insertions, 1 deletions

diff --git a/llvm_tools/patch_utils.py b/llvm_tools/patch_utils.py
index 2f282990..9117ba72 100644
--- a/llvm_tools/patch_utils.py
+++ b/llvm_tools/patch_utils.py
@@ -5,13 +5,14 @@
 """Provides patch utilities for PATCHES.json file handling."""
 
 import collections
+import contextlib
 import dataclasses
 import io
 from pathlib import Path
 import re
 import subprocess
 import sys
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 
 CHECKED_FILE_RE = re.compile(r'^checking file\s+(.*)$')
@@ -21,6 +22,38 @@ HUNK_END_RE = re.compile(r'^--\s*$')
 PATCH_SUBFILE_HEADER_RE = re.compile(r'^\+\+\+ [ab]/(.*)$')
 
 
+@contextlib.contextmanager
+def atomic_write(fp: Union[Path, str], mode='w', *args, **kwargs):
+  """Write to a filepath atomically.
+
+  This works by a temp file swap, created with a .tmp suffix in
+  the same directory briefly until being renamed to the desired
+  filepath.
+
+  Args:
+    fp: Filepath to open.
+    mode: File mode; can be 'w', 'wb'. Default 'w'.
+    *args: Passed to Path.open as nargs.
+    **kwargs: Passed to Path.open as kwargs.
+
+  Raises:
+    ValueError when the mode is invalid.
+  """
+  if isinstance(fp, str):
+    fp = Path(fp)
+  if mode not in ('w', 'wb'):
+    raise ValueError(f'mode {mode} not accepted')
+  temp_fp = fp.with_suffix(fp.suffix + '.tmp')
+  try:
+    with temp_fp.open(mode, *args, **kwargs) as f:
+      yield f
+  except:
+    if temp_fp.is_file():
+      temp_fp.unlink()
+    raise
+  temp_fp.rename(fp)
+
+
 @dataclasses.dataclass
 class Hunk:
   """Represents a patch Hunk."""
diff --git a/llvm_tools/patch_utils_unittest.py b/llvm_tools/patch_utils_unittest.py
index 3dfe52b2..bef5ae5f 100755
--- a/llvm_tools/patch_utils_unittest.py
+++ b/llvm_tools/patch_utils_unittest.py
@@ -6,6 +6,7 @@
 """Unit tests for the patch_utils.py file."""
 
 from pathlib import Path
+import tempfile
 import unittest
 import unittest.mock as mock
 
@@ -15,6 +16,31 @@ import patch_utils as pu
 class TestPatchUtils(unittest.TestCase):
   """Test the patch_utils."""
 
+  def test_atomic_write(self):
+    """Test that atomic write safely writes."""
+    prior_contents = 'This is a test written by patch_utils_unittest.py\n'
+    new_contents = 'I am a test written by patch_utils_unittest.py\n'
+    with tempfile.TemporaryDirectory(prefix='patch_utils_unittest') as dirname:
+      dirpath = Path(dirname)
+      filepath = dirpath / 'test_atomic_write.txt'
+      with filepath.open('w', encoding='utf-8') as f:
+        f.write(prior_contents)
+
+      def _t():
+        with pu.atomic_write(filepath, encoding='utf-8') as f:
+          f.write(new_contents)
+          raise Exception('Expected failure')
+
+      self.assertRaises(Exception, _t)
+      with filepath.open(encoding='utf-8') as f:
+        lines = f.readlines()
+      self.assertEqual(lines[0], prior_contents)
+      with pu.atomic_write(filepath, encoding='utf-8') as f:
+        f.write(new_contents)
+      with filepath.open(encoding='utf-8') as f:
+        lines = f.readlines()
+      self.assertEqual(lines[0], new_contents)
+
   def test_from_to_dict(self):
     """Test to and from dict conversion."""
     d = TestPatchUtils._default_json_dict()