"""ruamel.yaml round-trip helpers preserving comments and HA custom tags.""" from __future__ import annotations import math import re import threading from collections.abc import Callable from datetime import date, datetime from io import StringIO from typing import Any from ruamel.yaml import YAML from ruamel.yaml.scalarbool import ScalarBoolean 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() def _jsonify_float(node: float) -> float | str: """Narrow a float to something json can encode. ``.inf``/``.nan`` are valid YAML with no JSON encoding, so they render back to their YAML source form — the same treatment a tag gets. """ if math.isnan(node): return ".nan" if math.isinf(node): return ".inf" if node > 0 else "-.inf" return float(node) def yaml_jsonify(node: Any) -> Any: """Convert a round-trip node into JSON-serializable plain Python. An HA tag is rendered back to its SOURCE form (``!secret api_key``), never resolved: the value behind a ``!secret`` lives in secrets.yaml and is not looked up here, so a parsed view carries no plaintext-secret surface — the same property the round-trip text view has. Lives here because this module owns ``_TaggedScalar`` and the tag registry. ruamel's scalar types subclass the builtins (``ScalarInt``/``ScalarFloat``/ ``ScalarString``), so they are narrowed to the plain type; timestamps (``!!timestamp``) come back as ``date``/``datetime``, which json cannot encode, and become ISO strings. Non-finite floats (``.inf``/``.nan``) have no JSON encoding either, so they render back to their YAML source form — the same treatment a tag gets. """ if isinstance(node, _TaggedScalar): return f"{node.tag} {node.value}".strip() if isinstance(node, dict): return {str(key): yaml_jsonify(value) for key, value in node.items()} if isinstance(node, (list, tuple)): return [yaml_jsonify(item) for item in node] # Both branches must precede int: plain bool subclasses int, and a bool # carrying an anchor loads as ruamel's ScalarBoolean, which subclasses int # WITHOUT subclassing bool — so an `enabled: &flag true` would otherwise # serialize as 1. if node is None or isinstance(node, bool): return node if isinstance(node, ScalarBoolean): return bool(node) if isinstance(node, int): return int(node) if isinstance(node, float): return _jsonify_float(node) if isinstance(node, str): return str(node) if isinstance(node, (datetime, date)): return node.isoformat() return str(node)