From 59c6df5e460185dbe1deeb6790076e30e97150bb Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Thu, 10 Aug 2017 17:57:53 -0700
Subject: [PATCH] When possible, preserve config format on autoupdate

---
 pre_commit/commands/autoupdate.py | 45 ++++++++++++++++++++++---
 tests/commands/autoupdate_test.py | 56 +++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 5 deletions(-)

diff --git a/pre_commit/commands/autoupdate.py b/pre_commit/commands/autoupdate.py
index 620a8a6e..69ff2782 100644
--- a/pre_commit/commands/autoupdate.py
+++ b/pre_commit/commands/autoupdate.py
@@ -1,6 +1,7 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import re
 from collections import OrderedDict
 
 from aspy.yaml import ordered_dump
@@ -65,6 +66,44 @@ def _update_repo(repo_config, runner, tags_only):
     return new_config
 
 
+SHA_LINE_RE = re.compile(r'^(\s+)sha:(\s*)([^\s#]+)(.*)$', re.DOTALL)
+SHA_LINE_FMT = '{}sha:{}{}{}'
+
+
+def _write_new_config_file(path, output):
+    original_contents = open(path).read()
+    output = remove_defaults(output, CONFIG_SCHEMA)
+    new_contents = ordered_dump(output, **C.YAML_DUMP_KWARGS)
+
+    lines = original_contents.splitlines(True)
+    sha_line_indices_rev = list(reversed([
+        i for i, line in enumerate(lines) if SHA_LINE_RE.match(line)
+    ]))
+
+    for line in new_contents.splitlines(True):
+        if SHA_LINE_RE.match(line):
+            # It's possible we didn't identify the sha lines in the original
+            if not sha_line_indices_rev:
+                break
+            line_index = sha_line_indices_rev.pop()
+            original_line = lines[line_index]
+            orig_match = SHA_LINE_RE.match(original_line)
+            new_match = SHA_LINE_RE.match(line)
+            lines[line_index] = SHA_LINE_FMT.format(
+                orig_match.group(1), orig_match.group(2),
+                new_match.group(3), orig_match.group(4),
+            )
+
+    # If we failed to intelligently rewrite the sha lines, fall back to the
+    # pretty-formatted yaml output
+    to_write = ''.join(lines)
+    if ordered_load(to_write) != output:
+        to_write = new_contents
+
+    with open(path, 'w') as f:
+        f.write(to_write)
+
+
 def autoupdate(runner, tags_only):
     """Auto-update the pre-commit config to the latest versions of repos."""
     retv = 0
@@ -100,10 +139,6 @@ def autoupdate(runner, tags_only):
             output_configs.append(repo_config)
 
     if changed:
-        with open(runner.config_file_path, 'w') as config_file:
-            config_file.write(ordered_dump(
-                remove_defaults(output_configs, CONFIG_SCHEMA),
-                **C.YAML_DUMP_KWARGS
-            ))
+        _write_new_config_file(runner.config_file_path, output_configs)
 
     return retv
diff --git a/tests/commands/autoupdate_test.py b/tests/commands/autoupdate_test.py
index 8dac48c4..1920610a 100644
--- a/tests/commands/autoupdate_test.py
+++ b/tests/commands/autoupdate_test.py
@@ -1,5 +1,6 @@
 from __future__ import unicode_literals
 
+import pipes
 import shutil
 from collections import OrderedDict
 
@@ -123,6 +124,61 @@ def test_autoupdate_out_of_date_repo(
     assert out_of_date_repo.head_sha in after
 
 
+def test_does_not_reformat(
+        out_of_date_repo, mock_out_store_directory, in_tmpdir,
+):
+    fmt = (
+        '-   repo: {}\n'
+        '    sha: {}  # definitely the version I want!\n'
+        '    hooks:\n'
+        '    -   id: foo\n'
+        '        # These args are because reasons!\n'
+        '        args: [foo, bar, baz]\n'
+    )
+    config = fmt.format(out_of_date_repo.path, out_of_date_repo.original_sha)
+    with open(C.CONFIG_FILE, 'w') as f:
+        f.write(config)
+
+    autoupdate(Runner('.', C.CONFIG_FILE), tags_only=False)
+    after = open(C.CONFIG_FILE).read()
+    expected = fmt.format(out_of_date_repo.path, out_of_date_repo.head_sha)
+    assert after == expected
+
+
+def test_loses_formatting_when_not_detectable(
+        out_of_date_repo, mock_out_store_directory, in_tmpdir,
+):
+    """A best-effort attempt is made at updating sha without rewriting
+    formatting.  When the original formatting cannot be detected, this
+    is abandoned.
+    """
+    config = (
+        '[\n'
+        '    {{\n'
+        '        repo: {}, sha: {},\n'
+        '        hooks: [\n'
+        '            # A comment!\n'
+        '            {{id: foo}},\n'
+        '        ],\n'
+        '    }}\n'
+        ']\n'.format(
+            pipes.quote(out_of_date_repo.path), out_of_date_repo.original_sha,
+        )
+    )
+    with open(C.CONFIG_FILE, 'w') as f:
+        f.write(config)
+
+    autoupdate(Runner('.', C.CONFIG_FILE), tags_only=False)
+    after = open(C.CONFIG_FILE).read()
+    expected = (
+        '-   repo: {}\n'
+        '    sha: {}\n'
+        '    hooks:\n'
+        '    -   id: foo\n'
+    ).format(out_of_date_repo.path, out_of_date_repo.head_sha)
+    assert after == expected
+
+
 @pytest.yield_fixture
 def tagged_repo(out_of_date_repo):
     with cwd(out_of_date_repo.path):