|
| 1 | +# License: MIT |
| 2 | +# Copyright © 2025 Frequenz Energy-as-a-Service GmbH |
| 3 | + |
| 4 | +"""Patch existing dotted-key TOML files with new microgrid config values. |
| 5 | +
|
| 6 | +Unlike `dump_map`, which rebuilds a TOML document from scratch, this module |
| 7 | +only adds missing leaves and missing microgrid entries. Values already on |
| 8 | +disk always win, so comments, field order and number formatting survive. |
| 9 | +""" |
| 10 | + |
| 11 | +from collections.abc import Mapping |
| 12 | +from pathlib import Path |
| 13 | +from typing import Any |
| 14 | + |
| 15 | +import tomlkit |
| 16 | +from tomlkit import TOMLDocument |
| 17 | + |
| 18 | +from frequenz.gridpool import MicrogridConfig |
| 19 | + |
| 20 | +from ._dump_config import _format_value, _iter_leaves |
| 21 | + |
| 22 | + |
| 23 | +def patch_file(path: Path, configs: dict[str, MicrogridConfig]) -> str: |
| 24 | + """Patch the TOML file at `path` with any leaves missing from `configs`. |
| 25 | +
|
| 26 | + Args: |
| 27 | + path: Path to the existing TOML file to patch. |
| 28 | + configs: Mapping from microgrid ID (as string) to `MicrogridConfig`. |
| 29 | +
|
| 30 | + Returns: |
| 31 | + The patched TOML text; the caller is responsible for writing it back. |
| 32 | + """ |
| 33 | + return patch_text(path.read_text(), configs) |
| 34 | + |
| 35 | + |
| 36 | +def patch_text(original: str, configs: dict[str, MicrogridConfig]) -> str: |
| 37 | + """Patch dotted-key TOML text with any leaves missing from `configs`. |
| 38 | +
|
| 39 | + Args: |
| 40 | + original: The existing TOML text to patch. |
| 41 | + configs: Mapping from microgrid ID (as string) to `MicrogridConfig`. |
| 42 | +
|
| 43 | + Returns: |
| 44 | + The patched TOML text. |
| 45 | + """ |
| 46 | + doc = tomlkit.parse(original) |
| 47 | + schema = MicrogridConfig.Schema() |
| 48 | + |
| 49 | + # Leaves needing a whole new sub-table can't be inserted via item assignment |
| 50 | + # without tomlkit falling back to bracket-header syntax, so they are spliced |
| 51 | + # into the rendered text instead. |
| 52 | + orphans: dict[str, list[tuple[list[str], Any]]] = {} |
| 53 | + |
| 54 | + for mid, cfg in configs.items(): |
| 55 | + dumped = schema.dump(cfg) |
| 56 | + assert isinstance(dumped, dict) |
| 57 | + leaves = _iter_leaves([], dumped) |
| 58 | + if not leaves: |
| 59 | + continue |
| 60 | + |
| 61 | + if mid not in doc: |
| 62 | + _append_new_entry(doc, mid, leaves) |
| 63 | + continue |
| 64 | + |
| 65 | + for path, value in leaves: |
| 66 | + if _leaf_exists(doc, mid, path): |
| 67 | + continue |
| 68 | + if not _insert_leaf(doc, mid, path, value): |
| 69 | + orphans.setdefault(mid, []).append((path, value)) |
| 70 | + |
| 71 | + text = tomlkit.dumps(doc) |
| 72 | + if orphans: |
| 73 | + text = _splice_orphans(text, orphans) |
| 74 | + return text |
| 75 | + |
| 76 | + |
| 77 | +def _leaf_exists(doc: TOMLDocument, mid: str, path: list[str]) -> bool: |
| 78 | + """Whether the dotted key `mid.path...` already has a value in `doc`.""" |
| 79 | + node: Any = doc |
| 80 | + for key in (mid, *path): |
| 81 | + if not isinstance(node, Mapping) or key not in node: |
| 82 | + return False |
| 83 | + node = node[key] |
| 84 | + return True |
| 85 | + |
| 86 | + |
| 87 | +def _insert_leaf(doc: TOMLDocument, mid: str, path: list[str], value: Any) -> bool: |
| 88 | + """Insert a leaf directly onto its existing parent table, if there is one. |
| 89 | +
|
| 90 | + Args: |
| 91 | + doc: The document being patched, mutated in place. |
| 92 | + mid: Microgrid ID the leaf belongs to. |
| 93 | + path: Dotted key path under `mid` for the leaf. |
| 94 | + value: The leaf's value. |
| 95 | +
|
| 96 | + Returns: |
| 97 | + `True` if the leaf was inserted, `False` if a whole new sub-table is |
| 98 | + needed instead (left for the caller to handle). |
| 99 | + """ |
| 100 | + node: Any = doc |
| 101 | + matched = 0 |
| 102 | + for key in (mid, *path[:-1]): |
| 103 | + nxt = node[key] if isinstance(node, Mapping) and key in node else None |
| 104 | + if not isinstance(nxt, Mapping): |
| 105 | + break |
| 106 | + node = nxt |
| 107 | + matched += 1 |
| 108 | + |
| 109 | + full_path = [mid, *path] |
| 110 | + if matched != len(full_path) - 1: |
| 111 | + return False |
| 112 | + node[full_path[-1]] = _format_value(value) |
| 113 | + return True |
| 114 | + |
| 115 | + |
| 116 | +def _append_new_entry( |
| 117 | + doc: TOMLDocument, mid: str, leaves: list[tuple[list[str], Any]] |
| 118 | +) -> None: |
| 119 | + """Append a brand-new microgrid entry at the end of the document. |
| 120 | +
|
| 121 | + Args: |
| 122 | + doc: The document being patched, mutated in place. |
| 123 | + mid: Microgrid ID of the new entry. |
| 124 | + leaves: Flattened `(path, value)` pairs for the new entry, in |
| 125 | + dataclass field order. |
| 126 | + """ |
| 127 | + if doc.body: |
| 128 | + doc.add(tomlkit.nl()) |
| 129 | + for path, value in leaves: |
| 130 | + doc.append(tomlkit.key([mid, *path]), _format_value(value)) |
| 131 | + |
| 132 | + |
| 133 | +def _render_lines(mid: str, leaves: list[tuple[list[str], Any]]) -> list[str]: |
| 134 | + """Render `(path, value)` leaves as standalone `mid.path = value` lines.""" |
| 135 | + tmp = tomlkit.document() |
| 136 | + for path, value in leaves: |
| 137 | + tmp.append(tomlkit.key([mid, *path]), _format_value(value)) |
| 138 | + return tomlkit.dumps(tmp).splitlines(keepends=True) |
| 139 | + |
| 140 | + |
| 141 | +def _splice_orphans(text: str, orphans: dict[str, list[tuple[list[str], Any]]]) -> str: |
| 142 | + """Insert each microgrid's orphaned leaves right after its own last line. |
| 143 | +
|
| 144 | + Args: |
| 145 | + text: The already-rendered document text. |
| 146 | + orphans: Mapping from microgrid ID to its `(path, value)` leaves |
| 147 | + that need a brand-new sub-table. |
| 148 | +
|
| 149 | + Returns: |
| 150 | + `text` with the orphaned leaves inserted. |
| 151 | + """ |
| 152 | + lines = text.splitlines(keepends=True) |
| 153 | + for mid, leaves in orphans.items(): |
| 154 | + insert_at = _last_line_index_for_mid(lines, mid) |
| 155 | + new_lines = _render_lines(mid, leaves) |
| 156 | + lines[insert_at + 1 : insert_at + 1] = new_lines |
| 157 | + return "".join(lines) |
| 158 | + |
| 159 | + |
| 160 | +def _last_line_index_for_mid(lines: list[str], mid: str) -> int: |
| 161 | + """Index of the last line belonging to `mid` (its key starts with `mid.`).""" |
| 162 | + prefix = f"{mid}." |
| 163 | + for idx in range(len(lines) - 1, -1, -1): |
| 164 | + if lines[idx].lstrip().startswith(prefix): |
| 165 | + return idx |
| 166 | + raise ValueError(f"No existing line found for microgrid {mid!r}.") |
0 commit comments