vibe coding enabled :)

Author: mhordynski

Diff for: packages/ragbits-core/src/ragbits/core/utils/dict_transformations.py

@@ -43,6 +43,179 @@ def flatten_dict(input_dict: dict[str, Any], parent_key: str = "", sep: str = ".
     return items
 
 
+def _parse_key(key: str) -> list[tuple[str, bool]]:
+    """Parse a key into parts, each part being (name, is_array_index)."""
+    parts = []
+    current = ""
+    i = 0
+    while i < len(key):
+        if key[i] == "[":
+            if current:
+                parts.append((current, False))  # Parent is not an array
+                current = ""
+            i += 1  # Skip [
+            start = i
+            while i < len(key) and key[i] != "]":
+                i += 1
+            parts.append((key[start:i], True))
+            i += 1  # Skip ]
+            if i < len(key) and key[i] == ".":
+                i += 1  # Skip .
+        elif key[i] == ".":
+            if current:
+                parts.append((current, False))
+                current = ""
+            i += 1
+        else:
+            current += key[i]
+            i += 1
+    if current:
+        parts.append((current, False))
+    return parts
+
+
+def _ensure_array(obj: dict[str, Any] | list[Any], key: str) -> list[Any]:
+    """Ensure that obj[key] is a list, creating it if necessary."""
+    if isinstance(obj, list):
+        return obj
+    if key not in obj or not isinstance(obj[key], list):
+        obj[key] = []
+    return obj[key]
+
+
+def _ensure_dict(obj: dict[str, Any] | list[Any], key: str) -> dict[str, Any]:
+    """Ensure that obj[key] is a dict, creating it if necessary."""
+    if isinstance(obj, list):
+        # Lists should be handled by the caller
+        raise TypeError("Cannot ensure dict in a list")
+    if key not in obj or not isinstance(obj[key], dict):
+        obj[key] = {}
+    return obj[key]
+
+
+DictOrList = dict[str, Any] | list[Any]
+
+
+def _handle_array_part(
+    current: DictOrList,
+    part: str,
+    parent_key: str | None = None,
+) -> DictOrList:
+    """Handle an array part in the key."""
+    idx = int(part)
+    if isinstance(current, list):
+        while len(current) <= idx:
+            current.append({})
+        return current[idx]
+    if parent_key is None:
+        raise ValueError(f"Array part '{part}' without parent key")
+    current_list = _ensure_array(current, parent_key)
+    while len(current_list) <= idx:
+        current_list.append({})
+    return current_list[idx]
+
+
+def _handle_dict_part(
+    current: DictOrList,
+    part: str,
+    next_is_array: bool,
+    array_idx: int | None = None,
+) -> DictOrList:
+    """Handle a dictionary part in the key."""
+    if isinstance(current, list):
+        if array_idx is None:
+            raise ValueError("Array index is required when current is a list")
+        while len(current) <= array_idx:
+            current.append({})
+        current = current[array_idx]
+        if not isinstance(current, dict):
+            current = {}
+            current[str(array_idx)] = current
+    if next_is_array:
+        return _ensure_array(current, part)
+    return _ensure_dict(current, part)
+
+
+def _handle_single_part(
+    new_dict: dict[str, Any],
+    first_part: str,
+    is_array: bool,
+    value: SimpleTypes,
+) -> None:
+    """Handle a single-part key."""
+    if is_array:
+        idx = int(first_part)
+        current = _ensure_array(new_dict, first_part)
+        while len(current) <= idx:
+            current.append(None)
+        current[idx] = value
+    else:
+        new_dict[first_part] = value
+
+
+def _handle_last_array_part(
+    current_obj: DictOrList,
+    last_part: str,
+    value: SimpleTypes,
+    parts: list[tuple[str, bool]],
+) -> None:
+    """Handle the last part of the key when it's an array index."""
+    idx = int(last_part)
+    if len(parts) == 1:
+        # Direct array access like "users[0]"
+        parent_key = parts[0][0]
+        current_obj = _ensure_array(current_obj, parent_key)
+    if isinstance(current_obj, list):
+        while len(current_obj) <= idx:
+            current_obj.append(None)
+        current_obj[idx] = value
+    else:
+        raise TypeError("Expected list but got dict")
+
+
+def _handle_last_dict_part(
+    current_obj: DictOrList,
+    last_part: str,
+    value: SimpleTypes,
+    parts: list[tuple[str, bool]],
+) -> None:
+    """Handle the last part of the key when it's a dictionary key."""
+    if isinstance(current_obj, list):
+        # We're in a list, so we need to ensure the current index has a dict
+        idx = int(parts[-2][0])  # Get the index from the previous part
+        while len(current_obj) <= idx:
+            current_obj.append({})
+        current_obj = current_obj[idx]
+        if not isinstance(current_obj, dict):
+            current_obj = {}
+            current_obj[str(idx)] = current_obj
+    if isinstance(current_obj, dict):
+        current_obj[last_part] = value
+    else:
+        raise TypeError("Expected dict but got list")
+
+
+def _set_value(current: dict[str, Any], parts: list[tuple[str, bool]], value: SimpleTypes) -> None:
+    """Set a value in the dictionary based on the parsed key parts."""
+    current_obj: DictOrList = current
+
+    # Handle all parts except the last one
+    for i, (part, is_array) in enumerate(parts[:-1]):
+        if is_array:
+            current_obj = _handle_array_part(current_obj, part, parts[i - 1][0] if i > 0 else None)
+        else:
+            next_is_array = i + 1 < len(parts) and parts[i + 1][1]
+            array_idx = int(parts[i][0]) if isinstance(current_obj, list) else None
+            current_obj = _handle_dict_part(current_obj, part, next_is_array, array_idx)
+
+    # Handle the last part
+    last_part, is_array = parts[-1]
+    if is_array:
+        _handle_last_array_part(current_obj, last_part, value, parts)
+    else:
+        _handle_last_dict_part(current_obj, last_part, value, parts)
+
+
 def unflatten_dict(input_dict: dict[str, Any]) -> dict[str, Any]:
     """
     Converts a flattened dictionary with dot notation and array notation into a nested structure.
@@ -73,110 +246,6 @@ def unflatten_dict(input_dict: dict[str, Any]) -> dict[str, Any]:
 
     new_dict: dict[str, Any] = {}
 
-    def _parse_key(key: str) -> list[tuple[str, bool]]:
-        """Parse a key into parts, each part being (name, is_array_index)."""
-        parts = []
-        current = ""
-        i = 0
-        while i < len(key):
-            if key[i] == "[":
-                if current:
-                    parts.append((current, False))  # Parent is not an array
-                    current = ""
-                i += 1  # Skip [
-                start = i
-                while i < len(key) and key[i] != "]":
-                    i += 1
-                parts.append((key[start:i], True))
-                i += 1  # Skip ]
-                if i < len(key) and key[i] == ".":
-                    i += 1  # Skip .
-            elif key[i] == ".":
-                if current:
-                    parts.append((current, False))
-                    current = ""
-                i += 1
-            else:
-                current += key[i]
-                i += 1
-        if current:
-            parts.append((current, False))
-        return parts
-
-    def _ensure_array(obj: dict[str, Any], key: str) -> list:
-        """Ensure that obj[key] is a list, creating it if necessary."""
-        if key not in obj:
-            obj[key] = []
-        elif not isinstance(obj[key], list):
-            obj[key] = []
-        return obj[key]
-
-    def _ensure_dict(obj: Any, key: str) -> dict:
-        """Ensure that obj[key] is a dict, creating it if necessary."""
-        if key not in obj:
-            obj[key] = {}
-        elif not isinstance(obj[key], dict):
-            obj[key] = {}
-        return obj[key]
-
-    def _set_value(current: dict[str, Any], parts: list[tuple[str, bool]], value: Any) -> None:
-        """Set a value in the dictionary based on the parsed key parts."""
-        for i, (part, is_array) in enumerate(parts[:-1]):
-            if is_array:
-                idx = int(part)
-                if i > 0:
-                    # Get the parent key and ensure it's an array
-                    parent_key = parts[i - 1][0]
-                    if isinstance(current, list):
-                        while len(current) <= idx:
-                            current.append({})
-                        current = current[idx]
-                    else:
-                        current = _ensure_array(current, parent_key)
-                        while len(current) <= idx:
-                            current.append({})
-                        current = current[idx]
-            else:
-                if i + 1 < len(parts) and parts[i + 1][1]:  # Next part is array
-                    if isinstance(current, list):
-                        # We're in a list, so we need to ensure the current index has a dict
-                        idx = int(parts[i][0])
-                        while len(current) <= idx:
-                            current.append({})
-                        current = current[idx]
-                        current = _ensure_array(current, part)
-                    else:
-                        current = _ensure_array(current, part)
-                else:
-                    if isinstance(current, list):
-                        # We're in a list, so we need to ensure the current index has a dict
-                        idx = int(parts[i][0])
-                        while len(current) <= idx:
-                            current.append({})
-                        current = current[idx]
-                        current = _ensure_dict(current, part)
-                    else:
-                        current = _ensure_dict(current, part)
-
-        last_part, is_array = parts[-1]
-        if is_array:
-            idx = int(last_part)
-            if len(parts) == 1:
-                # Direct array access like "users[0]"
-                parent_key = parts[0][0]
-                current = _ensure_array(current, parent_key)
-            while len(current) <= idx:
-                current.append(None)
-            current[idx] = value
-        else:
-            if isinstance(current, list):
-                # We're in a list, so we need to ensure the current index has a dict
-                idx = int(parts[-2][0])  # Get the index from the previous part
-                while len(current) <= idx:
-                    current.append({})
-                current = current[idx]
-            current[last_part] = value
-
     # Sort keys to ensure we process parents before children
     field_keys = sorted(input_dict.keys())
     for key in field_keys:
@@ -191,14 +260,7 @@ def _set_value(current: dict[str, Any], parts: list[tuple[str, bool]], value: An
 
         # Set the value
         if len(parts) == 1:
-            if is_array:
-                idx = int(first_part)
-                current = _ensure_array(new_dict, first_part)
-                while len(current) <= idx:
-                    current.append(None)
-                current[idx] = input_dict[key]
-            else:
-                new_dict[first_part] = input_dict[key]
+            _handle_single_part(new_dict, first_part, is_array, input_dict[key])
         else:
             _set_value(new_dict, parts, input_dict[key])
 

Diff for: packages/ragbits-core/tests/unit/utils/test_dict_transformations.py

@@ -1,12 +1,14 @@
+from typing import Any
+
 import pytest
 
-from ragbits.core.utils.dict_transformations import flatten_dict, unflatten_dict
+from ragbits.core.utils.dict_transformations import SimpleTypes, flatten_dict, unflatten_dict
 
 
 def test_flatten_dict_simple():
     """Test flattening a simple dictionary."""
-    input_dict = {"key1": "value1", "key2": "value2"}
-    expected = {"key1": "value1", "key2": "value2"}
+    input_dict: dict[str, Any] = {"key1": "value1", "key2": "value2"}
+    expected: dict[str, SimpleTypes] = {"key1": "value1", "key2": "value2"}
     assert flatten_dict(input_dict) == expected
 
 
@@ -46,11 +48,11 @@ def test_flatten_dict_empty():
 
 def test_flatten_dict_with_non_dict_values():
     """Test flattening a dictionary with various value types."""
-    input_dict = {
+    input_dict: dict[str, Any] = {
         "key1": "value1",
         "nested": {"subkey1": 42, "subkey2": [1, 2, 3], "subkey3": {"a": 1}, "subkey4": True, "subkey5": 3.14},
     }
-    expected = {
+    expected: dict[str, SimpleTypes] = {
         "key1": "value1",
         "nested.subkey1": 42,
         "nested.subkey2[0]": 1,
@@ -65,7 +67,7 @@ def test_flatten_dict_with_non_dict_values():
 
 def test_flatten_unflatten():
     """Test flattening and unflattening a dictionary."""
-    input_dict = {
+    input_dict: dict[str, Any] = {
         "key1": "value1",
         "nested": {
             "subkey1": "subvalue1",
@@ -154,8 +156,8 @@ def test_mixed_array_and_object():
 def test_unflatten_dict_notation_based_types():
     """Test that unflatten_dict uses notation to determine types, not heuristics."""
     # Test that numeric keys without array notation stay as dict
-    input_dict = {"0": "first", "1": "second", "2": "third"}
-    expected = {"0": "first", "1": "second", "2": "third"}
+    input_dict: dict[str, SimpleTypes] = {"0": "first", "1": "second", "2": "third"}
+    expected: dict[str, Any] = {"0": "first", "1": "second", "2": "third"}
     assert unflatten_dict(input_dict) == expected
 
     # Test that array notation creates lists
@@ -171,31 +173,18 @@ def test_unflatten_dict_notation_based_types():
         "nested[0].name": "John",
         "nested[0].age": 30,
         "nested[1].name": "Jane",
-        "nested[1].age": 25
+        "nested[1].age": 25,
     }
     expected = {
         "dict_key": "value",
         "list_key": ["first", "second"],
-        "nested": [
-            {"name": "John", "age": 30},
-            {"name": "Jane", "age": 25}
-        ]
+        "nested": [{"name": "John", "age": 30}, {"name": "Jane", "age": 25}],
     }
     assert unflatten_dict(input_dict) == expected
 
     # Test that numeric keys in nested dicts stay as dict
-    input_dict = {
-        "config.0.name": "first",
-        "config.1.name": "second",
-        "config.2.name": "third"
-    }
-    expected = {
-        "config": {
-            "0": {"name": "first"},
-            "1": {"name": "second"},
-            "2": {"name": "third"}
-        }
-    }
+    input_dict = {"config.0.name": "first", "config.1.name": "second", "config.2.name": "third"}
+    expected = {"config": {"0": {"name": "first"}, "1": {"name": "second"}, "2": {"name": "third"}}}
     assert unflatten_dict(input_dict) == expected