From 3246b5e25ab2983f075b8ac12d746026e72ea082 Mon Sep 17 00:00:00 2001 From: Eduard Trott Date: Wed, 21 Aug 2024 15:54:56 +0200 Subject: [PATCH] Support logical sorting extra functionality --- tests/integration/test_extra.py | 213 +++++++++++++++++++++++++++++++- yamx/extra.py | 102 ++++++++++++++- yamx/loader/preprocessor.py | 8 +- 3 files changed, 317 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_extra.py b/tests/integration/test_extra.py index 780e87f..ac586c2 100644 --- a/tests/integration/test_extra.py +++ b/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() diff --git a/yamx/extra.py b/yamx/extra.py index e891676..ffa9273 100644 --- a/yamx/extra.py +++ b/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 diff --git a/yamx/loader/preprocessor.py b/yamx/loader/preprocessor.py index 9399070..beec51e 100644 --- a/yamx/loader/preprocessor.py +++ b/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)