If a shebang already exists when adding a preamble, preserve it. (Che…

…rry-pick of #19133) (#19137) Currently if a file contains a shebang but no preamble, the premable linter will break the shebang by prepending the preamble. This change adjusts the regex match to preserve the existing shebang. Co-authored-by: Stu Hood <stuhood@gmail.com>
pantsbuild · May 24, 2023 · e0e7eab · e0e7eab
1 parent 5eda12f
commit e0e7eab
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 5 deletions.
diff --git a/src/python/pants/backend/tools/preamble/rules.py b/src/python/pants/backend/tools/preamble/rules.py
@@ -23,10 +23,20 @@ class PreambleRequest(FmtFilesRequest):
 
 @memoized
 def _template_checker_regex(template: str) -> re.Pattern:
-    maybe_shebang = r"(#!.*\n)?"
+    shebang = r"(?P<shebang>#!.*\n)"
+    if re.match(shebang, template):
+        # If the template already contains a shebang, don't attempt to match one.
+        shebang = r"(?P<shebang>)"
+    maybe_shebang = f"{shebang}?"
+
     subbed = string.Template(template).safe_substitute(year=r"====YEAR====")
     raw_regex = re.escape(subbed).replace("====YEAR====", r"\d{4}")
-    return re.compile(maybe_shebang + raw_regex)
+    maybe_template = f"(?P<template>{raw_regex})?"
+
+    # Wrap in `?s` to enable matching newlines with `.`.
+    body = "(?s:(?P<body>.*))"
+
+    return re.compile(maybe_shebang + maybe_template + body)
 
 
 @memoized
@@ -57,10 +67,15 @@ async def preamble_fmt(
     for path, template in template_by_path.items():
         regex = _template_checker_regex(template)
         file_content = contents_by_path[path].content
-        if not regex.match(file_content.decode("utf-8")):
-            new_content = _substituted_template(template).encode("utf-8") + file_content
+        matched = regex.fullmatch(file_content.decode("utf-8"))
+        # The regex produced by `_template_checker_regex` is infallible.
+        assert matched
+        if not matched.group("template"):
+            shebang = matched.group("shebang")
+            body = matched.group("body")
+            new_content = (shebang if shebang else "") + _substituted_template(template) + body
             contents_by_path[path] = dataclasses.replace(
-                contents_by_path[path], content=new_content
+                contents_by_path[path], content=new_content.encode("utf-8")
             )
 
     output_snapshot = await Get(Snapshot, CreateDigest(contents_by_path.values()))

diff --git a/src/python/pants/backend/tools/preamble/rules_integration_test.py b/src/python/pants/backend/tools/preamble/rules_integration_test.py
@@ -141,6 +141,27 @@ def test_ignores_shebang(rule_runner: RuleRunner) -> None:
     assert fmt_result.did_change is False
 
 
+def test_preserves_shebang(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "foo.py": "#!/usr/bin/env python3\n...",
+        }
+    )
+    fmt_result = run_preamble(
+        rule_runner,
+        {
+            "**/*.py": "# Copyright\n",
+        },
+    )
+
+    assert fmt_result.did_change is True
+    assert fmt_result.output == rule_runner.make_snapshot(
+        {
+            "foo.py": "#!/usr/bin/env python3\n# Copyright\n...",
+        }
+    )
+
+
 def test_preamble_includes_shebang(rule_runner: RuleRunner) -> None:
     files_before = {"foo.py": "#!/usr/bin/env python3\n# Copyright"}