Support logical sorting extra functionality

tests/integration/test_extra.py

@@ -4,7 +4,7 @@
 from ruamel.yaml import YAML
 
 from yamx import YAMX
-from yamx.extra import ResolvingContext, extract_toggles, resolve_toggles
+from yamx.extra import ResolvingContext, extract_toggles, resolve_toggles, sort_logical
 
 
 @pytest.mark.parametrize(
@@ -399,3 +399,214 @@ def test_resolve_failed(raw_config):
     context = ResolvingContext({"defines": {"toggle_a": True}})
     with pytest.raises(Exception, match="Unsupported toggle condition: "):
         resolve_toggles(data, context)
+
+
+@pytest.mark.parametrize(
+    "config, input_dict, res_dict",
+    [
+        # root key sort
+        (
+            {(): {"key_order": ["b", "a"]}},
+            """
+a: 2
+b: 1
+""",
+            """
+b: 1
+a: 2
+""",
+        ),
+        # nested dict key sort
+        (
+            {("nested1", "nested2", "nested3"): {"key_order": ["b", "a"]}},
+            """
+nested1:
+  nested2:
+    nested3:
+      a: 2
+      b: 1
+""",
+            """
+nested1:
+  nested2:
+    nested3:
+      b: 1
+      a: 2
+""",
+        ),
+        # nested list key sort
+        (
+            {("[]", "nested", "[]"): {"key_order": ["b", "a"]}},
+            """
+- nested:
+  - a: 2
+    b: 1
+""",
+            """
+- nested:
+  - b: 1
+    a: 2
+""",
+        ),
+        # root item sort
+        (
+            {(): {"item_order": ["a"]}},
+            """
+- a: 2
+- a: 1
+""",
+            """
+- a: 1
+- a: 2
+""",
+        ),
+        # root item sort by 2 keys
+        (
+            {(): {"item_order": ["a", "b"]}},
+            """
+- a: 2
+  b: 2
+- b: 1
+  a: 2
+- a: 1
+  b: 2
+- a: 1
+  b: 1
+""",
+            """
+- a: 1
+  b: 1
+- a: 1
+  b: 2
+- b: 1
+  a: 2
+- a: 2
+  b: 2
+""",
+        ),
+        (
+            # nested item sort
+            {("nested", "[]", "nested1"): {"item_order": ["a"]}},
+            """
+nested:
+- nested1:
+  - a: 2
+  - a: 1
+""",
+            """
+nested:
+- nested1:
+  - a: 1
+  - a: 2
+""",
+        ),
+        (
+            # item sort with nested key
+            {(): {"item_order": [("a", "b")]}},
+            """
+- a:
+    b: 2
+- a:
+    b: 1
+""",
+            """
+- a:
+    b: 1
+- a:
+    b: 2
+""",
+        ),
+        # key sort with conditional blocks
+        (
+            {("[]",): {"key_order": ["a", "b"]}},
+            """
+- b: 1
+  a: 1
+# {% if defines.get("toggle") %}
+- b: 2
+  a: 2
+# {% else %}
+- b: 3
+  a: 3
+# {% endif %}
+""",
+            """
+- a: 1
+  b: 1
+# {% if defines.get("toggle") %}
+- a: 2
+  b: 2
+# {% else %}
+- a: 3
+  b: 3
+# {% endif %}
+""",
+        ),
+        # key sort with nested conditional blocks
+        (
+            {("[]", "nested"): {"key_order": ["a", "b"]}},
+            """
+- nested:
+    b: 1
+    a: 1
+# {% if defines.get("toggle") %}
+- nested:
+    # {% if defines.get("toggle1") %}
+    b: 2
+    a: 2
+    # {% else %}
+    b: 3
+    a: 3
+    # {% endif %}
+# {% endif %}
+""",
+            """
+- nested:
+    a: 1
+    b: 1
+# {% if defines.get("toggle") %}
+- nested:
+    # {% if defines.get("toggle1") %}
+    a: 2
+    b: 2
+    # {% else %}
+    a: 3
+    b: 3
+    # {% endif %}
+# {% endif %}
+""",
+        ),
+        # partial key sort with conditional blocks
+        (
+            {("nested",): {"key_order": ["a", "b", "c", "d"]}},
+            """
+nested:
+  c: 1
+  # {% if defines.get("toggle") %}
+  d: 3
+  a: 2
+  # {% else %}
+  a: 3
+  # {% endif %}
+  b: 1
+""",
+            """
+nested:
+  b: 1
+  c: 1
+  # {% if defines.get("toggle") %}
+  a: 2
+  d: 3
+  # {% else %}
+  a: 3
+  # {% endif %}
+""",
+        ),
+    ],
+)
+def test_sort_logical(config, input_dict, res_dict):
+    yamx = YAMX(sort_keys=False)
+    data = yamx.load_from_string(input_dict)
+    sorted_data = sort_logical(data, config)
+    resolved_str = yamx.dump_to_string(sorted_data)
+    assert resolved_str == res_dict.strip()

yamx/extra.py

@@ -1,10 +1,11 @@
 # not very supported pieces of functionality
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Set
+from typing import Any, Dict, Final, Optional, Set, Tuple, Union
 
 from jinja2 import nodes
 
+from yamx.constants import CONDITIONAL_KEY_PREFIX
 from yamx.containers.data import (
     Condition,
     ConditionalData,
@@ -15,6 +16,8 @@
 from yamx.loader.utils import get_jinja_env
 from yamx.utils import strtobool
 
+LIST_SYMBOL: Final[str] = "[]"
+
 
 @dataclass(frozen=True)
 class ResolvingContext:
@@ -173,3 +176,100 @@ def resolve_condition(
     resolved = env.from_string(jinja_ast).render(context)
 
     return bool(strtobool(resolved))
+
+
+def _config_by_key(config: dict, key: str):
+    """Get config by key"""
+    return {k[1:]: v for k, v in config.items() if len(k) and k[0] == key}
+
+
+def _extract_item_by_key(obj, key: Union[Tuple[str, ...], str]):
+    """Extract item by key"""
+    if isinstance(key, str):
+        return obj[key]
+    if len(key) == 1:
+        return obj[key[0]]
+
+    if isinstance(obj, ConditionalMap):
+        try:
+            return _extract_item_by_key(obj[key[0]], key[1:])
+        except KeyError:
+            return None
+    elif isinstance(obj, ConditionalSeq):
+        if key[0] == LIST_SYMBOL:
+            return [item[key[1]] for item in obj]
+        return obj[key]
+    else:
+        return None
+
+
+def sort_logical(obj: Any, config: dict, toggled: bool = False) -> Any:
+    """Perform logical sorting based on the provided structure"""
+    if isinstance(obj, ConditionalData):
+        return ConditionalData(sort_logical(obj.data, config))
+    if isinstance(obj, ConditionalMap):
+        print(obj)
+        res_dict = {}
+        # sorting config for processed dict is under empty tuple key
+        order_config = config.get(())
+        if order_config:
+            key_order = order_config.get("key_order", [])
+            ordered_obj = sorted(
+                obj.items(),
+                key=lambda kv: key_order.index(kv[0])
+                if kv[0] in key_order
+                else len(key_order),
+            )
+        else:
+            ordered_obj = obj.items()
+
+        for key, value in ordered_obj:
+            if toggled or key.startswith(CONDITIONAL_KEY_PREFIX):
+                key_config = config
+            else:
+                key_config = _config_by_key(config, key)
+            # if sorting config is empty, stop processing nested structures
+            if len(key_config):
+                res_dict[key] = sort_logical(value, key_config)
+            else:
+                res_dict[key] = value
+
+        return ConditionalMap(res_dict)
+    elif isinstance(obj, ConditionalSeq):
+        res_seq = []
+        order_config = config.get(())
+        if order_config:
+            item_order = order_config.get("item_order", [])
+            ordered_obj = sorted(
+                obj,
+                key=lambda item: [
+                    _extract_item_by_key(item, key) for key in item_order
+                ],
+            )
+        else:
+            ordered_obj = obj
+        if toggled:
+            item_config = config
+        else:
+            item_config = _config_by_key(config, LIST_SYMBOL)
+        if item_config:
+            res_seq = [sort_logical(item, item_config) for item in ordered_obj]
+        else:
+            res_seq = ordered_obj
+        return ConditionalSeq(res_seq)
+    elif isinstance(obj, ConditionalGroup):
+
+        return ConditionalGroup(
+            condition=obj.condition,
+            body=sort_logical(obj.body, config, toggled=True),
+            else_body=sort_logical(obj.else_body, config, toggled=True),
+            elif_bodies=tuple(
+                ConditionalGroup(
+                    condition=elif_body.condition,
+                    body=sort_logical(elif_body.body, config, toggled=True),
+                )
+                for elif_body in obj.elif_bodies
+            ),
+        )
+    else:
+        return obj

yamx/loader/preprocessor.py

@@ -143,10 +143,10 @@ def _process_conditions(
     raw_data: str, typ: ConditionalBlockType, condition: Optional[Condition] = None
 ) -> str:
     """
-    # here we rely on the fact that yaml structure is valid and we can:
-    # * load it with YAML loader
-    # * set dynamically generated tags based on conditions
-    # * dump it back to string form
+    here we rely on the fact that yaml structure is valid and we can:
+    * load it with YAML loader
+    * set dynamically generated tags based on conditions
+    * dump it back to string form
     """
     global UNIQUE_CONDITION_CNT
     yaml_data = yaml.load(raw_data)