sqlmesh.utils.yaml

  1from __future__ import annotations
  2
  3import getpass
  4import io
  5import typing as t
  6from decimal import Decimal
  7from os import getenv
  8from pathlib import Path
  9
 10from ruamel import yaml
 11from ruamel.yaml.constructor import SafeConstructor
 12
 13from sqlmesh.core.constants import VAR
 14from sqlmesh.utils.errors import SQLMeshError
 15from sqlmesh.utils.jinja import ENVIRONMENT, create_var
 16
 17JINJA_METHODS = {
 18    "env_var": lambda key, default=None: getenv(key, default),
 19    "user": lambda: getpass.getuser(),
 20}
 21
 22
 23def YAML(typ: t.Optional[str] = "safe") -> yaml.YAML:
 24    yaml_obj = yaml.YAML(typ=typ)
 25
 26    # Ruamel doesn't know how to serialize Decimal values. This is problematic when,
 27    # e.g., we're trying to auto-generate a unit test whose body contains Decimal data.
 28    # This is a best-effort approach to solve this by serializing them as strings.
 29    yaml_obj.representer.add_representer(
 30        Decimal, lambda dumper, data: dumper.represent_str(str(data))
 31    )
 32
 33    return yaml_obj
 34
 35
 36class SafeConstructorOverride(SafeConstructor):
 37    def check_mapping_key(
 38        self,
 39        node: t.Any,
 40        key_node: t.Any,
 41        mapping: t.Any,
 42        key: t.Any,
 43        value: t.Any,
 44    ) -> bool:
 45        """This function normally returns True if key is unique.
 46
 47        It is only used by the construct_mapping function. By always returning True,
 48        keys will always be updated and so the last value will be kept for mappings.
 49        """
 50        return True
 51
 52
 53def load(
 54    source: str | Path,
 55    raise_if_empty: bool = True,
 56    render_jinja: bool = True,
 57    allow_duplicate_keys: bool = False,
 58    variables: t.Optional[t.Dict[str, t.Any]] = None,
 59    keep_last_duplicate_key: bool = False,
 60) -> t.Dict:
 61    """Loads a YAML object from either a raw string or a file."""
 62    path: t.Optional[Path] = None
 63
 64    if isinstance(source, Path):
 65        path = source
 66        with open(source, "r", encoding="utf-8") as file:
 67            source = file.read()
 68
 69    if render_jinja:
 70        source = ENVIRONMENT.from_string(source).render(
 71            {
 72                **JINJA_METHODS,
 73                VAR: create_var(variables or {}),
 74            }
 75        )
 76
 77    yaml = YAML()
 78    if allow_duplicate_keys and keep_last_duplicate_key:
 79        yaml.Constructor = SafeConstructorOverride
 80    yaml.allow_duplicate_keys = allow_duplicate_keys
 81    contents = yaml.load(source)
 82    if contents is None:
 83        if raise_if_empty:
 84            error_path = f" '{path}'" if path else ""
 85            raise SQLMeshError(f"YAML source{error_path} can't be empty.")
 86        return {}
 87
 88    return contents
 89
 90
 91@t.overload
 92def dump(value: t.Any, stream: io.IOBase) -> None: ...
 93
 94
 95@t.overload
 96def dump(value: t.Any) -> str: ...
 97
 98
 99def dump(value: t.Any, stream: t.Optional[io.IOBase] = None) -> t.Optional[str]:
100    """Dumps a ruamel.yaml loaded object and converts it into a string or writes it to a stream."""
101    result = io.StringIO()
102    YAML(typ=None).dump(value, stream or result)
103    return None if stream else result.getvalue()