feat: built-in to_yaml filter for jinja yaml config loader

Author: Lariel Fernandes

src/mlopus/utils/jinja_yaml.py

@@ -1,4 +1,5 @@
 import os
+import textwrap
 from collections import defaultdict
 from pathlib import Path
 from typing import Any, Literal, Iterator, Tuple, Callable
@@ -47,6 +48,11 @@ def load_jinja_yaml_configs(
     :param file_extensions: File extensions to use. Defaults to `{'.yml', '.yaml'}`. The leading dot is ignored.
     """
 
+    custom_filters = custom_filters or {}
+    for func in _FILTERS:
+        if func.__name__ not in custom_filters:
+            custom_filters[func.__name__] = func
+
     namespace_files: dict[str, list[Path]] = defaultdict(list)
     for namespace, file_path in _iter_files_with_namespaces(Path(base_path), file_extensions or {".yml", ".yaml"}):
         namespace_files[namespace].append(file_path)
@@ -94,7 +100,7 @@ def load_jinja_yaml_configs(
     return result
 
 
-def _iter_files_with_namespaces(base_path: Path, extensions: list[str]) -> Iterator[Tuple[str, Path]]:
+def _iter_files_with_namespaces(base_path: Path, extensions: set[str]) -> Iterator[Tuple[str, Path]]:
     """Iterate (namespace, path) tuples for every YAML file in the specified path.
 
     Namespaces are determined by the file's top dir or by its name before any double underscores.
@@ -108,3 +114,35 @@ def _iter_files_with_namespaces(base_path: Path, extensions: list[str]) -> Itera
         rel_path = file_path.relative_to(base_path)
         namespace = rel_path.parts[0] if len(rel_path.parts) > 1 else file_path.stem.split("__")[0]
         yield namespace, file_path
+
+
+class _NotSet:
+    pass
+
+
+class _Filters:
+    """Built-in jinja filters."""
+
+    @staticmethod
+    def to_yaml(
+        arg: Any,
+        *,
+        indent: int = 0,
+        if_none: Any | None = _NotSet,
+        if_falsy: Any | None = _NotSet,
+    ) -> str:
+        if arg is None and if_none is not _NotSet:
+            arg = if_none
+
+        if not arg and if_falsy is not _NotSet:
+            arg = if_falsy
+
+        encoded = yaml.safe_dump(arg).removesuffix("\n").removesuffix("\n...")
+
+        if prefix := indent * " ":
+            encoded = textwrap.indent(encoded, prefix).removeprefix(prefix)
+
+        return encoded
+
+
+_FILTERS = (_Filters.to_yaml,)

src/tests/test_utils/test_jinja_yaml.py

@@ -21,6 +21,13 @@ def test_load_jinja_yaml_configs_basic():
               bar: foo
             with_suffix: {{ "foo" | add_suffix }}
             empty_key: ""
+            this_is_none: null
+            this_is_falsy_str: ""
+            this_is_falsy_int: 0
+            this_is_falsy_float: 0.0
+            dict_with_nones:
+              key1: null
+              key2: [null]
         """)
         )
 
@@ -40,6 +47,12 @@ def test_load_jinja_yaml_configs_basic():
               {% for key, value in common.conn_defaults.items() %}
               {{ key }}: {{ value }}
               {% endfor %}
+            this_is_none: {{ common.this_is_none | to_yaml(if_none="") }}
+            this_is_falsy_str: {{ common.this_is_falsy_str | to_yaml(if_falsy=None) }}
+            this_is_falsy_int: {{ common.this_is_falsy_int | to_yaml(if_falsy=None) }}
+            this_is_falsy_float: {{ common.this_is_falsy_float | to_yaml(if_falsy=None) }}
+            dict_with_nones:
+              {{ common.dict_with_nones | to_yaml(indent=2) }}
         """)
         )
 
@@ -71,3 +84,10 @@ def test_load_jinja_yaml_configs_basic():
         assert result["database"]["missing_gets_fallback"] == "fallback"
         assert result["database"]["empty_gets_fallback"] == "fallback"
         assert result["database"]["conn_defaults"] == result["common"]["conn_defaults"]
+
+        assert result["database"]["this_is_none"] == ""
+        assert result["database"]["this_is_falsy_str"] is None
+        assert result["database"]["this_is_falsy_int"] is None
+        assert result["database"]["this_is_falsy_float"] is None
+
+        assert result["database"]["dict_with_nones"] == {"key1": None, "key2": [None]}