Restore the mode from before the edit, finish the restore after any exception and raise an exit again with the report

esphome/yaml_edit.py

@@ -47,11 +47,13 @@ class LineEdit:
 
 @dataclass
 class Snapshot:
-    """A rewritten file's text before and after; the restore only puts
-    ``original`` back over ``written``, never over the user's own edit."""
+    """A rewritten file's text before and after, and its mode before; the
+    restore only puts ``original`` back over ``written``, never over the
+    user's own edit."""
 
     original: str
     written: str
+    mode: int | None = None
 
 
 def field_line_re(
@@ -170,13 +172,20 @@ def source_line(mapping: ConfigType, name: str) -> tuple[Path, int, str]:
     return doc, line_no, line_at(doc, line_no)
 
 
-def write_keeping_mode(path: Path, text: str, like: Path | None = None) -> None:
-    """Write with the mode of ``like`` (default: the file itself) rather
-    than write_file's 0644; a 0600 secrets file stays 0600."""
+def mode_of(path: Path) -> int:
     try:
-        mode = stat.S_IMODE((like or path).stat().st_mode)
+        return stat.S_IMODE(path.stat().st_mode)
     except OSError as err:
-        raise EsphomeError(f"Could not read the mode of {like or path}: {err}") from err
+        raise EsphomeError(f"Could not read the mode of {path}: {err}") from err
+
+
+def write_keeping_mode(
+    path: Path, text: str, like: Path | None = None, mode: int | None = None
+) -> None:
+    """Write with ``mode``, else the mode of ``like`` (default: the file
+    itself), rather than write_file's 0644; a 0600 secrets file stays 0600."""
+    if mode is None:
+        mode = mode_of(like or path)
     try:
         write_file(path, text, private=True)
     except EsphomeError as err:
@@ -226,7 +235,9 @@ def apply_line_edits(edits: list[LineEdit]) -> dict[Path, Snapshot]:
     originals = {}
     for path, own in by_path.items():
         original = read_text(path)
-        originals[path] = Snapshot(original, rewritten_text(original, own))
+        originals[path] = Snapshot(
+            original, rewritten_text(original, own), mode_of(path)
+        )
     try:
         for path, snapshot in originals.items():
             write_keeping_mode(path, snapshot.written)
@@ -258,19 +269,25 @@ def apply_line_edits(edits: list[LineEdit]) -> dict[Path, Snapshot]:
 
 
 def restore_files(originals: dict[Path, Snapshot]) -> None:
-    """Put back what apply_line_edits wrote; a file that failed or changed
-    meanwhile is left alone and reported. Every file is tried whatever went
-    wrong before; an interrupt is raised again at the end with the report
-    as a note."""
+    """Put back what apply_line_edits wrote, mode included; a file that
+    failed or changed meanwhile is left alone and reported. Every file is
+    tried whatever went wrong before; an interrupt or exit is raised again
+    at the end with the report as a note."""
     failed: list[str] = []
     stop: BaseException | None = None
 
     def report(what: str, err: BaseException) -> None:
         nonlocal stop
-        why = "interrupted" if isinstance(err, KeyboardInterrupt) else str(err)
-        failed.append(f"{what}: {why or type(err).__name__}")
-        if not isinstance(err, Exception):
+        if isinstance(err, Exception):
+            why = str(err) or type(err).__name__
+        else:  # an interrupt or exit says what it is, not its code
+            why = (
+                "interrupted"
+                if isinstance(err, KeyboardInterrupt)
+                else type(err).__name__
+            )
             stop = err
+        failed.append(f"{what}: {why}")
 
     for path, snapshot in originals.items():
         try:
@@ -279,12 +296,12 @@ def restore_files(originals: dict[Path, Snapshot]) -> None:
                 continue  # never written, or already put back
             if current != snapshot.written:
                 raise EsphomeError("changed since it was written, left as is")
-            write_keeping_mode(path, snapshot.original)
-        except (EsphomeError, OSError, KeyboardInterrupt) as err:
+            write_keeping_mode(path, snapshot.original, mode=snapshot.mode)
+        except BaseException as err:  # noqa: BLE001  # pylint: disable=broad-except
             report(str(path), err)
     try:
         compiled_config.invalidate_compiled_config()
-    except (EsphomeError, KeyboardInterrupt) as err:
+    except BaseException as err:  # noqa: BLE001  # pylint: disable=broad-except
         report("the validated caches", err)
     if not failed:
         return

tests/unit_tests/test_yaml_edit.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+import stat
 import sys
 from typing import Any
 from unittest.mock import patch
@@ -289,8 +290,11 @@ def test_apply_writes_reloads_drops_the_cache_and_restores(tmp_path: Path) -> No
     cache = compiled_config_path(CORE.config_filename)
     cache.parent.mkdir(parents=True, exist_ok=True)
     cache.write_text("{}")
+    mode = stat.S_IMODE(path.stat().st_mode)
     originals = apply_line_edits([_name_edit("garage")])
-    assert originals == {path: Snapshot(YAML, YAML.replace("kitchen  #", "garage  #"))}
+    assert originals == {
+        path: Snapshot(YAML, YAML.replace("kitchen  #", "garage  #"), mode)
+    }
     assert path.read_text() == YAML.replace("kitchen  #", "garage  #")
     assert not cache.exists()
     cache.write_text("{}")
@@ -443,11 +447,11 @@ def test_restore_goes_on_after_an_interrupt(tmp_path: Path) -> None:
     calls: list[Path] = []
 
     def interrupt_at(n: int) -> Any:
-        def write(target: Path, text: str) -> None:
+        def write(target: Path, text: str, **kwargs: Any) -> None:
             calls.append(target)
             if len(calls) == n:
                 raise KeyboardInterrupt
-            write_keeping_mode(target, text)
+            write_keeping_mode(target, text, **kwargs)
 
         return write
 
@@ -484,11 +488,11 @@ def test_restore_finishes_after_any_failure(tmp_path: Path) -> None:
     )
     writes: list[Path] = []
 
-    def fail_first(target: Path, text: str) -> None:
+    def fail_first(target: Path, text: str, **kwargs: Any) -> None:
         writes.append(target)
         if len(writes) == 1:
             raise OSError("odd")
-        write_keeping_mode(target, text)
+        write_keeping_mode(target, text, **kwargs)
 
     with (
         patch("esphome.yaml_edit.write_keeping_mode", side_effect=fail_first),
@@ -511,6 +515,45 @@ def test_restore_finishes_after_any_failure(tmp_path: Path) -> None:
     assert "the validated caches: interrupted" in "".join(info.value.__notes__)
 
 
+@pytest.mark.skipif(sys.platform == "win32", reason="posix file modes")
+def test_restore_puts_back_the_mode_from_before_the_edit(tmp_path: Path) -> None:
+    """A write whose chmod failed leaves the file 0600; the restore does not
+    copy that, it has the mode from before."""
+    path = _setup(tmp_path, YAML)
+    path.chmod(0o644)
+    originals = apply_line_edits([_name_edit("garage")])
+    path.chmod(0o600)
+    restore_files(originals)
+    assert path.read_text() == YAML
+    assert path.stat().st_mode & 0o777 == 0o644
+
+
+def test_restore_finishes_after_an_exit(tmp_path: Path) -> None:
+    """A SystemExit from a signal handler still leaves the other files put
+    back and comes out again with the report."""
+    path = _setup(tmp_path, YAML)
+    other = tmp_path / "other.yaml"
+    other.write_bytes(b"x: 1\n")
+    originals = apply_line_edits(
+        [_name_edit("garage"), LineEdit(other, 0, "x: 1", "x: 2")]
+    )
+    writes: list[Path] = []
+
+    def exit_first(target: Path, text: str, **kwargs: Any) -> None:
+        writes.append(target)
+        if len(writes) == 1:
+            raise SystemExit(1)
+        write_keeping_mode(target, text, **kwargs)
+
+    with (
+        patch("esphome.yaml_edit.write_keeping_mode", side_effect=exit_first),
+        pytest.raises(SystemExit) as info,
+    ):
+        restore_files(originals)
+    assert other.read_text() == "x: 1\n"
+    assert f"Could not restore {path}: SystemExit" in "".join(info.value.__notes__)
+
+
 def test_a_value_written_to_the_secrets_file_is_registered(tmp_path: Path) -> None:
     """Only a resolved !secret registers its value, so the main file is
     loaded again after an edit elsewhere; own_documents keeps the registry."""