Merge pull request #20 from tkdchen/STONEBLD-3422-retain-indentation-style

Retain block sequence indentation

Author: tkdchen

src/pipeline_migration/migrate.py

@@ -14,8 +14,8 @@
 from pipeline_migration.utils import dump_yaml, load_yaml
 from pipeline_migration.registry import Container, Registry, ImageIndex
 from pipeline_migration.quay import QuayTagInfo, list_active_repo_tags
-from pipeline_migration.utils import is_true
 from pipeline_migration.types import FilePath
+from pipeline_migration.utils import is_true, YAMLStyle
 
 ANNOTATION_HAS_MIGRATION: Final[str] = "dev.konflux-ci.task.has-migration"
 ANNOTATION_IS_MIGRATION: Final[str] = "dev.konflux-ci.task.is-migration"
@@ -105,15 +105,17 @@ def resolve_pipeline(pipeline_file: FilePath) -> Generator[FilePath, Any, None]:
     :type pipeline_file: str
     :return: a generator yielding a file containing the pipeline definition.
     """
+    yaml_style = YAMLStyle.detect(pipeline_file)
     origin_pipeline = load_yaml(pipeline_file)
+
     if not isinstance(origin_pipeline, dict):
         raise ValueError(f"Given file {pipeline_file} is not a YAML mapping.")
 
     kind = origin_pipeline.get("kind")
     if kind == TEKTON_KIND_PIPELINE:
         yield pipeline_file
         pl_yaml = load_yaml(pipeline_file)
-        dump_yaml(pipeline_file, pl_yaml)
+        dump_yaml(pipeline_file, pl_yaml, style=yaml_style)
     elif kind == TEKTON_KIND_PIPELINE_RUN:
         spec = origin_pipeline.get("spec") or {}
         if "pipelineSpec" in spec:
@@ -125,7 +127,7 @@ def resolve_pipeline(pipeline_file: FilePath) -> Generator[FilePath, Any, None]:
             yield temp_pipeline_file
             modified_pipeline = load_yaml(temp_pipeline_file)
             spec["pipelineSpec"] = modified_pipeline["spec"]
-            dump_yaml(pipeline_file, origin_pipeline)
+            dump_yaml(pipeline_file, origin_pipeline, style=yaml_style)
         elif "pipelineRef" in spec:
             # Pipeline definition can be referenced here, via either git-resolver or a name field
             # pointing to YAML file under the .tekton/.

src/pipeline_migration/utils.py

@@ -1,17 +1,100 @@
+from dataclasses import dataclass, field
 from typing import Any
-from ruamel.yaml import YAML
 
 from pipeline_migration.types import FilePath
 
+from ruamel.yaml import YAML, CommentedMap, CommentedSeq
+from ruamel.yaml.comments import CommentedBase
+
+
+__all__ = [
+    "BlockSequenceIndentation",
+    "create_yaml_obj",
+    "dump_yaml",
+    "is_true",
+    "load_yaml",
+    "YAMLStyle",
+]
+
 
 def is_true(value: str) -> bool:
     return value.strip().lower() == "true"
 
 
-def create_yaml_obj():
+@dataclass
+class BlockSequenceIndentation:
+    indentations: dict[int, int] = field(default_factory=dict)
+
+    @property
+    def is_consistent(self) -> bool:
+        return len(self.levels) == 1
+
+    @property
+    def levels(self) -> list[int]:
+        return list(self.indentations.keys())
+
+
+@dataclass
+class _NodePath:
+    node: CommentedBase
+    field: str = ""
+
+
+@dataclass
+class YAMLStyle:
+    indentation: BlockSequenceIndentation
+
+    preserve_quotes: bool = True
+    width: int = 8192
+
+    @classmethod
+    def _detect_block_sequence_indentation(cls, node: CommentedBase) -> BlockSequenceIndentation:
+
+        parent_nodes: list[_NodePath] = []
+        block_seq_indentations = BlockSequenceIndentation()
+        indentations = block_seq_indentations.indentations
+
+        def _walk(node: CommentedBase) -> None:
+            if isinstance(node, CommentedMap):
+                parent_nodes.append(_NodePath(node=node))
+                for key, value in node.items():
+                    parent_nodes[-1].field = key
+                    _walk(value)
+                parent_nodes.pop()
+            elif isinstance(node, CommentedSeq):
+                levels = node.lc.col - parent_nodes[-1].node.lc.col
+                if levels in indentations:
+                    indentations[levels] += 1
+                else:
+                    indentations[levels] = 1
+                for item in node:
+                    _walk(item)
+
+        _walk(node)
+
+        return block_seq_indentations
+
+    @classmethod
+    def detect(cls, file_path: FilePath) -> "YAMLStyle":
+        doc = load_yaml(file_path)
+        indentation = cls._detect_block_sequence_indentation(doc)
+        return cls(indentation=indentation)
+
+
+def create_yaml_obj(style: YAMLStyle | None = None):
     yaml = YAML()
-    yaml.preserve_quotes = True
-    yaml.width = 8192
+    if style is None:
+        return yaml
+
+    yaml.preserve_quotes = style.preserve_quotes
+    yaml.width = style.width
+
+    offset = 0
+    if style.indentation.is_consistent:
+        offset = style.indentation.levels[0]
+    sequence = offset + 2
+    yaml.indent(sequence=sequence, offset=offset)
+
     return yaml
 
 
@@ -20,6 +103,6 @@ def load_yaml(yaml_file: FilePath) -> Any:
         return create_yaml_obj().load(f)
 
 
-def dump_yaml(yaml_file: FilePath, data: Any):
+def dump_yaml(yaml_file: FilePath, data: Any, style: YAMLStyle | None = None) -> None:
     with open(yaml_file, "w", encoding="utf-8") as f:
-        create_yaml_obj().dump(data, f)
+        create_yaml_obj(style).dump(data, f)

tests/conftest.py

@@ -1,5 +1,6 @@
 import pytest
 from copy import deepcopy
+from textwrap import dedent
 from typing import Final
 
 import responses
@@ -107,3 +108,43 @@ def _mock(c: Container, has_migration=False, previous_migration_bundle: str | No
         responses.add(responses.GET, f"https://{c.manifest_url()}", json=manifest_json)
 
     return _mock
+
+
+@pytest.fixture
+def pipeline_yaml():
+    return dedent(
+        """\
+        apiVersion: tekton.dev/v1
+        kind: Pipeline
+        metadata:
+          name: pl
+        spec:
+          tasks:
+          - name: clone
+            image: debian:latest
+            script: |
+              git clone https://git.host/project
+        """
+    )
+
+
+@pytest.fixture(params=["no_indents", "2_spaces_indents"])
+def pipeline_yaml_with_various_indent_styles(request, pipeline_yaml):
+    match request.param:
+        case "no_indents":
+            return pipeline_yaml
+        case "2_spaces_indents":
+            return dedent(
+                """\
+                apiVersion: tekton.dev/v1
+                kind: Pipeline
+                metadata:
+                  name: pl
+                spec:
+                  tasks:
+                    - name: clone
+                      image: debian:latest
+                      script: |
+                        git clone https://git.host/project
+                """
+            )

tests/test_cli.py

@@ -23,7 +23,8 @@
     ensure_container,
 )
 
-from tests.test_migrate import APP_IMAGE_REPO, PIPELINE_DEFINITION, TASK_BUNDLE_CLONE
+from pipeline_migration.utils import YAMLStyle
+from tests.test_migrate import APP_IMAGE_REPO, TASK_BUNDLE_CLONE
 from tests.utils import generate_digest, generate_git_sha
 
 
@@ -218,20 +219,29 @@ def _mock_quay_list_tags(self):
                 json={"tags": tags, "page": 1, "has_additional": False},
             )
 
-    def _mock_pipeline_file(self, tmp_path: Path) -> Path:
+    def _mock_pipeline_file(self, tmp_path: Path, content: str) -> Path:
         tekton_dir = tmp_path / ".tekton"
         tekton_dir.mkdir()
         pipeline_file = tekton_dir / "component-pipeline.yaml"
-        pipeline_file.write_text(PIPELINE_DEFINITION)
+        pipeline_file.write_text(content)
         return pipeline_file
 
     @responses.activate
     @pytest.mark.parametrize("use_linked_migrations", [True, False])
-    def test_apply_migrations(self, use_linked_migrations, monkeypatch, tmp_path):
+    def test_apply_migrations(
+        self,
+        use_linked_migrations,
+        pipeline_yaml_with_various_indent_styles,
+        monkeypatch,
+        tmp_path,
+    ):
         monkeypatch.setattr("pipeline_migration.migrate.Registry", MockRegistry)
         self._mock_quay_list_tags()
 
-        pipeline_file = self._mock_pipeline_file(tmp_path)
+        pipeline_file = self._mock_pipeline_file(tmp_path, pipeline_yaml_with_various_indent_styles)
+
+        # Verified later
+        origin_style = YAMLStyle.detect(pipeline_file)
 
         tb_upgrades = [
             {
@@ -292,6 +302,14 @@ def _subprocess_run(cmd, *args, **kwargs):
 
         assert entry_point() is None
 
+        # Verify result formatting
+        cur_style = YAMLStyle.detect(pipeline_file)
+        assert cur_style.indentation.is_consistent
+        if origin_style.indentation.is_consistent:
+            assert origin_style.indentation.levels == cur_style.indentation.levels
+        else:
+            assert cur_style.indentation.levels == [0]
+
 
 def test_entry_point_should_catch_error(monkeypatch, caplog):
     cli_cmd = ["pmt", "--use-legacy-resolver", "-u", json.dumps(UPGRADES)]

tests/test_migrate.py

@@ -646,26 +646,10 @@ def _fetch_migration_file(image: str, digest: str) -> str | None:
         assert script_content == migrations[0].migration_script
 
 
-PIPELINE_DEFINITION: Final = dedent(
-    """\
-apiVersion: tekton.dev/v1
-kind: Pipeline
-metadata:
-  name: pl
-spec:
-  tasks:
-  - name: clone
-    image: debian:latest
-    script: |
-      git clone https://git.host/project
-"""
-)
-
-
 class TestApplyMigrations:
 
     @pytest.mark.parametrize("chdir", [True, False])
-    def test_apply_migrations(self, chdir, tmp_path, monkeypatch):
+    def test_apply_migrations(self, pipeline_yaml, chdir, tmp_path, monkeypatch):
         """Ensure applying all resolved migrations"""
 
         renovate_upgrades = [
@@ -703,7 +687,7 @@ def test_apply_migrations(self, chdir, tmp_path, monkeypatch):
         tekton_dir = tmp_path / ".tekton"
         tekton_dir.mkdir()
         package_file = tekton_dir / "pipeline.yaml"
-        package_file.write_text(PIPELINE_DEFINITION)
+        package_file.write_text(pipeline_yaml)
 
         test_context = {"executed_scripts": []}
         counter = itertools.count()
@@ -733,7 +717,9 @@ def subprocess_run(*args, **kwargs):
             with pytest.raises(ValueError, match="Pipeline file does not exist: .+"):
                 manager.apply_migrations()
 
-    def test_raise_error_if_migration_process_fails(self, caplog, monkeypatch, tmp_path):
+    def test_raise_error_if_migration_process_fails(
+        self, pipeline_yaml, caplog, monkeypatch, tmp_path
+    ):
         caplog.set_level(logging.DEBUG, logger="migrate")
 
         renovate_upgrades = [
@@ -771,7 +757,7 @@ def test_raise_error_if_migration_process_fails(self, caplog, monkeypatch, tmp_p
         tekton_dir = tmp_path / ".tekton"
         tekton_dir.mkdir()
         package_file = tekton_dir / "pipeline.yaml"
-        package_file.write_text(PIPELINE_DEFINITION)
+        package_file.write_text(pipeline_yaml)
 
         def _mkstemp(*args, **kwargs):
             tmp_file_path = tmp_path / "migration_file"