"""ruamel.yaml round-trip helpers preserving comments and HA custom tags.""" from __future__ import annotations import re import threading from collections.abc import Callable from io import StringIO from typing import Any from ruamel.yaml import YAML class _TaggedScalar: """Wrapper that stores a YAML tag + scalar value for lossless round-trip.""" __slots__ = ("tag", "value") def __init__(self, tag: str, value: str) -> None: self.tag = tag self.value = value def __repr__(self) -> str: return f"_TaggedScalar({self.tag!r}, {self.value!r})" def __str__(self) -> str: return self.value def __eq__(self, other: object) -> bool: if not isinstance(other, _TaggedScalar): return NotImplemented return self.tag == other.tag and self.value == other.value def __hash__(self) -> int: return hash((self.tag, self.value)) _HA_TAGS = ( "!include", "!include_dir_list", "!include_dir_named", "!include_dir_merge_list", "!include_dir_merge_named", "!secret", "!env_var", ) def _make_tag_constructor(tag: str) -> Callable[[Any, Any], _TaggedScalar]: """Return a ruamel.yaml constructor function for *tag*.""" def _construct(loader: Any, node: Any) -> _TaggedScalar: return _TaggedScalar(tag, loader.construct_scalar(node)) return _construct def _represent_tagged_scalar(dumper: Any, data: _TaggedScalar) -> Any: """Representer that emits the original tag + scalar value.""" return dumper.represent_scalar(data.tag, data.value) def _register_ha_tags() -> None: """Register HA tag constructors/representers on the shared class registries. ``add_constructor`` / ``add_representer`` mutate class-level registries shared by all ``YAML(typ="rt")`` instances. We call this once at import time; ``make_yaml()`` then only creates a fresh (thread-safe) instance. """ # Use a temporary instance to access the Constructor/Representer classes _tmp = YAML(typ="rt") for tag in _HA_TAGS: _tmp.Constructor.add_constructor(tag, _make_tag_constructor(tag)) _tmp.Representer.add_representer(_TaggedScalar, _represent_tagged_scalar) _register_ha_tags() # Effectively-infinite emitter line width. ruamel's default (~80 columns) # re-wraps long lines on dump; inside a ``>`` folded scalar a new wrap # adjacent to a more-indented line becomes a LITERAL newline on re-parse, # silently corrupting string literals in blocks an edit never touched # (#1720). Never introducing new wraps also keeps untouched long lines # byte-stable across edits. _NEVER_WRAP_WIDTH = 2**31 # A top-level mapping key: starts at column 0, `key:` with nothing (or a # comment) after the colon. Quoted/exotic keys never match — detection # then just falls back to the default style, which is safe. _TOP_LEVEL_KEY_RE = re.compile(r"^[A-Za-z0-9_][^\s:]*:\s*(?:#.*)?$") _DASH_RE = re.compile(r"^( *)- ") # ruamel's compact defaults for block sequences (dash at the parent # column). Used to RESET the shared per-thread instance between dumps. _DEFAULT_SEQ_STYLE = (2, 0) def detect_seq_indent(text: str) -> tuple[int, int] | None: """Detect the file's top-level block-sequence style. Returns ``(sequence, offset)`` for ``YAML.indent()`` — derived from the first list item that directly follows a top-level key — or ``None`` when the file has no such sequence. Only top-level sequences discriminate: nested dashes are indented in BOTH styles. """ lines = text.splitlines() for i, line in enumerate(lines): if not _TOP_LEVEL_KEY_RE.match(line): continue for nxt in lines[i + 1 :]: if not nxt.strip() or nxt.lstrip().startswith("#"): continue m = _DASH_RE.match(nxt) if m: offset = len(m.group(1)) return (offset + 2, offset) break # value is not a sequence — try the next top-level key return None def apply_seq_indent(ry: YAML, style: tuple[int, int] | None) -> None: """Apply a detected sequence style (or the compact default) to *ry*. ``make_yaml()`` instances are cached per-thread, so the style MUST be (re)applied before every dump — passing ``None`` resets to the default instead of leaking the previous file's style. """ sequence, offset = style if style is not None else _DEFAULT_SEQ_STYLE ry.indent(mapping=2, sequence=sequence, offset=offset) def _build_yaml() -> YAML: """Create a fresh round-trip YAML instance with HA tag support.""" ry = YAML(typ="rt") ry.preserve_quotes = True ry.width = _NEVER_WRAP_WIDTH return ry class _YAMLStorage(threading.local): """Thread-local storage for ruamel.yaml instances.""" def __init__(self) -> None: self.yaml = _build_yaml() _STORAGE = _YAMLStorage() def make_yaml() -> YAML: """Return a round-trip YAML instance with HA tag support. The instance is cached per-thread to prevent ruamel.yaml from performing expensive plugin discovery (glob/scandir) on every call, which causes CPU spikes and event loop blocking during bulk edits. Thread-local storage is used because ruamel.yaml instances are not thread-safe. """ try: return _STORAGE.yaml except AttributeError: _STORAGE.yaml = _build_yaml() return _STORAGE.yaml def yaml_dumps(ry: YAML, data: Any) -> str: """Dump *data* to a string using the given YAML instance.""" buf = StringIO() ry.dump(data, buf) return buf.getvalue()