feat: initial checkin

Author: bdraco

src/annotatedyaml/__init__.py

@@ -1 +1,30 @@
-__version__ = "0.0.0"
+"""YAML utility functions."""
+
+from .const import SECRET_YAML
+from .dumper import dump, save_yaml
+from .input import UndefinedSubstitution, extract_inputs, substitute
+from .loader import (
+    Secrets,
+    YamlTypeError,
+    load_yaml,
+    load_yaml_dict,
+    parse_yaml,
+    secret_yaml,
+)
+from .objects import Input
+
+__all__ = [
+    "SECRET_YAML",
+    "Input",
+    "Secrets",
+    "UndefinedSubstitution",
+    "YamlTypeError",
+    "dump",
+    "extract_inputs",
+    "load_yaml",
+    "load_yaml_dict",
+    "parse_yaml",
+    "save_yaml",
+    "secret_yaml",
+    "substitute",
+]

src/annotatedyaml/const.py

@@ -0,0 +1,3 @@
+"""Constants."""
+
+SECRET_YAML = "secrets.yaml"

src/annotatedyaml/dumper.py

@@ -0,0 +1,96 @@
+"""Custom dumper and representers."""
+
+from collections import OrderedDict
+from typing import Any
+
+import yaml
+
+from .objects import Input, NodeDictClass, NodeListClass, NodeStrClass
+
+# mypy: allow-untyped-calls, no-warn-return-any
+
+
+try:
+    from yaml import CSafeDumper as FastestAvailableSafeDumper
+except ImportError:
+    from yaml import (  # type: ignore[assignment]
+        SafeDumper as FastestAvailableSafeDumper,
+    )
+
+
+def dump(_dict: dict | list) -> str:
+    """Dump YAML to a string and remove null."""
+    return yaml.dump(
+        _dict,
+        default_flow_style=False,
+        allow_unicode=True,
+        sort_keys=False,
+        Dumper=FastestAvailableSafeDumper,
+    ).replace(": null\n", ":\n")
+
+
+def save_yaml(path: str, data: dict) -> None:
+    """Save YAML to a file."""
+    # Dump before writing to not truncate the file if dumping fails
+    str_data = dump(data)
+    with open(path, "w", encoding="utf-8") as outfile:
+        outfile.write(str_data)
+
+
+# From: https://gist.github.com/miracle2k/3184458
+def represent_odict(  # type: ignore[no-untyped-def]
+    dumper, tag, mapping, flow_style=None
+) -> yaml.MappingNode:
+    """Like BaseRepresenter.represent_mapping but does not issue the sort()."""
+    value: list = []
+    node = yaml.MappingNode(tag, value, flow_style=flow_style)
+    if dumper.alias_key is not None:
+        dumper.represented_objects[dumper.alias_key] = node
+    best_style = True
+    if hasattr(mapping, "items"):
+        mapping = mapping.items()
+    for item_key, item_value in mapping:
+        node_key = dumper.represent_data(item_key)
+        node_value = dumper.represent_data(item_value)
+        if not (isinstance(node_key, yaml.ScalarNode) and not node_key.style):
+            best_style = False
+        if not (isinstance(node_value, yaml.ScalarNode) and not node_value.style):
+            best_style = False
+        value.append((node_key, node_value))
+    if flow_style is None:
+        if dumper.default_flow_style is not None:
+            node.flow_style = dumper.default_flow_style
+        else:
+            node.flow_style = best_style
+    return node
+
+
+def add_representer(klass: Any, representer: Any) -> None:
+    """Add to representer to the dumper."""
+    FastestAvailableSafeDumper.add_representer(klass, representer)
+
+
+add_representer(
+    OrderedDict,
+    lambda dumper, value: represent_odict(dumper, "tag:yaml.org,2002:map", value),
+)
+
+add_representer(
+    NodeDictClass,
+    lambda dumper, value: represent_odict(dumper, "tag:yaml.org,2002:map", value),
+)
+
+add_representer(
+    NodeListClass,
+    lambda dumper, value: dumper.represent_sequence("tag:yaml.org,2002:seq", value),
+)
+
+add_representer(
+    NodeStrClass,
+    lambda dumper, value: dumper.represent_scalar("tag:yaml.org,2002:str", str(value)),
+)
+
+add_representer(
+    Input,
+    lambda dumper, value: dumper.represent_scalar("!input", value.name),
+)

src/annotatedyaml/input.py

@@ -0,0 +1,56 @@
+"""Deal with YAML input."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .objects import Input
+
+
+class UndefinedSubstitution(Exception):
+    """Error raised when we find a substitution that is not defined."""
+
+    def __init__(self, input_name: str) -> None:
+        """Initialize the undefined substitution exception."""
+        super().__init__(f"No substitution found for input {input_name}")
+        self.input = input
+
+
+def extract_inputs(obj: Any) -> set[str]:
+    """Extract input from a structure."""
+    found: set[str] = set()
+    _extract_inputs(obj, found)
+    return found
+
+
+def _extract_inputs(obj: Any, found: set[str]) -> None:
+    """Extract input from a structure."""
+    if isinstance(obj, Input):
+        found.add(obj.name)
+        return
+
+    if isinstance(obj, list):
+        for val in obj:
+            _extract_inputs(val, found)
+        return
+
+    if isinstance(obj, dict):
+        for val in obj.values():
+            _extract_inputs(val, found)
+        return
+
+
+def substitute(obj: Any, substitutions: dict[str, Any]) -> Any:
+    """Substitute values."""
+    if isinstance(obj, Input):
+        if obj.name not in substitutions:
+            raise UndefinedSubstitution(obj.name)
+        return substitutions[obj.name]
+
+    if isinstance(obj, list):
+        return [substitute(val, substitutions) for val in obj]
+
+    if isinstance(obj, dict):
+        return {key: substitute(val, substitutions) for key, val in obj.items()}
+
+    return obj