migrate-config rewrites deprecated stages

asottile · asottile · commit 5679399d905a · 2024-09-16T20:36:33.000-04:00
diff --git a/pre_commit/commands/migrate_config.py b/pre_commit/commands/migrate_config.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import functools
+import itertools
 import textwrap
 from typing import Callable
 
@@ -49,6 +50,10 @@ def _preserve_style(n: ScalarNode, *, s: str) -> str:
     return f'{n.style}{s}{n.style}'
 
 
+def _fix_stage(n: ScalarNode) -> str:
+    return _preserve_style(n, s=f'pre-{n.value}')
+
+
 def _migrate_composed(contents: str) -> str:
     tree = yaml_compose(contents)
     rewrites: list[tuple[ScalarNode, Callable[[ScalarNode], str]]] = []
@@ -76,6 +81,22 @@ def _migrate_composed(contents: str) -> str:
         if node.value == 'python_venv':
             rewrites.append((node, python_venv_replace))
 
+    # stages rewrites
+    default_stages_matcher = (MappingValue('default_stages'), SequenceItem())
+    default_stages_match = match(tree, default_stages_matcher)
+    hook_stages_matcher = (
+        MappingValue('repos'),
+        SequenceItem(),
+        MappingValue('hooks'),
+        SequenceItem(),
+        MappingValue('stages'),
+        SequenceItem(),
+    )
+    hook_stages_match = match(tree, hook_stages_matcher)
+    for node in itertools.chain(default_stages_match, hook_stages_match):
+        if node.value in {'commit', 'push', 'merge-commit'}:
+            rewrites.append((node, _fix_stage))
+
     rewrites.sort(reverse=True, key=lambda nf: nf[0].start_mark.index)
 
     src_parts = []
diff --git a/tests/commands/migrate_config_test.py b/tests/commands/migrate_config_test.py
@@ -213,6 +213,48 @@ def test_migrate_config_quoted_python_venv(tmp_path):
     assert cfg.read_text() == expected
 
 
+def test_migrate_config_default_stages(tmp_path):
+    src = '''\
+default_stages: [commit, push, merge-commit, commit-msg]
+repos: []
+'''
+    expected = '''\
+default_stages: [pre-commit, pre-push, pre-merge-commit, commit-msg]
+repos: []
+'''
+    cfg = tmp_path.joinpath('cfg.yaml')
+    cfg.write_text(src)
+    assert migrate_config(str(cfg)) == 0
+    assert cfg.read_text() == expected
+
+
+def test_migrate_config_hook_stages(tmp_path):
+    src = '''\
+repos:
+-   repo: local
+    hooks:
+    -   id: example
+        name: example
+        entry: example
+        language: system
+        stages: ["commit", "push", "merge-commit", "commit-msg"]
+'''
+    expected = '''\
+repos:
+-   repo: local
+    hooks:
+    -   id: example
+        name: example
+        entry: example
+        language: system
+        stages: ["pre-commit", "pre-push", "pre-merge-commit", "commit-msg"]
+'''
+    cfg = tmp_path.joinpath('cfg.yaml')
+    cfg.write_text(src)
+    assert migrate_config(str(cfg)) == 0
+    assert cfg.read_text() == expected
+
+
 def test_migrate_config_invalid_yaml(tmpdir):
     contents = '['
     cfg = tmpdir.join(C.CONFIG_FILE)
