TheRock/rocm-systems/projects/rocprofiler-compute/tools/config_management/utils_ruamel.py

# Copyright (c) Advanced Micro Devices, Inc.
# SPDX-License-Identifier:  MIT

from pathlib import Path
from typing import Any, Union

from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedMap

# --- Round-trip YAML (for writing) ---
RT_YAML = YAML(typ="rt")
RT_YAML.preserve_quotes = True
RT_YAML.width = 4096  # prevent unwanted line wrapping
RT_YAML.indent(mapping=2, sequence=2, offset=0)
RT_YAML.explicit_start = False
RT_YAML.explicit_end = False

# --- Read-only YAML (safe loader) ---
RO_YAML = YAML(typ="safe")
RO_YAML.width = 4096


def load_yaml(
    filepath: Union[str, Path],
    *,
    round_trip: bool = False,
) -> Any:
    path = Path(filepath)
    if not path.exists():
        raise FileNotFoundError(f"YAML file not found: {path}")

    yaml = RT_YAML if round_trip else RO_YAML

    with open(path, "r", encoding="utf-8") as f:
        return yaml.load(f) or CommentedMap()


def save_yaml(data: Any, filepath: Union[str, Path]) -> None:
    path = Path(filepath)
    path.parent.mkdir(parents=True, exist_ok=True)

    with open(path, "w", encoding="utf-8") as f:
        RT_YAML.dump(data, f)


def strip_existing_header(yaml_data) -> None:
    ca = getattr(yaml_data, "ca", None)
    if not ca or not hasattr(ca, "comment") or ca.comment is None:
        return

    original = ca.comment

    cleaned = []

    for block in original:
        if block is None:
            cleaned.append(None)
            continue

        new_block = [token for token in block if "AUTOGENERATED" not in token.value]

        if not new_block:
            cleaned.append(None)
        else:
            cleaned.append(new_block)

    if len(cleaned) < 2:
        cleaned.extend([None] * (2 - len(cleaned)))

    ca.comment = cleaned