fix migrate-config for purelib yaml

asottile · asottile · commit 222c62bc5d29 · 2024-10-08T11:46:48.000-04:00
diff --git a/pre_commit/commands/migrate_config.py b/pre_commit/commands/migrate_config.py
@@ -47,7 +47,8 @@ def _migrate_map(contents: str) -> str:
 
 
 def _preserve_style(n: ScalarNode, *, s: str) -> str:
-    return f'{n.style}{s}{n.style}'
+    style = n.style or ''
+    return f'{style}{s}{style}'
 
 
 def _fix_stage(n: ScalarNode) -> str:
diff --git a/tests/commands/migrate_config_test.py b/tests/commands/migrate_config_test.py
@@ -1,10 +1,26 @@
 from __future__ import annotations
 
+from unittest import mock
+
 import pytest
+import yaml
 
 import pre_commit.constants as C
 from pre_commit.clientlib import InvalidConfigError
 from pre_commit.commands.migrate_config import migrate_config
+from pre_commit.yaml import yaml_compose
+
+
+@pytest.fixture(autouse=True, params=['c', 'pure'])
+def switch_pyyaml_impl(request):
+    if request.param == 'c':
+        yield
+    else:
+        with mock.patch.dict(
+                yaml_compose.keywords,
+                {'Loader': yaml.SafeLoader},
+        ):
+            yield
 
 
 def test_migrate_config_normal_format(tmpdir, capsys):
