✨ Add map_leaves

Author: Kajiih

src/nested_dict_tools/__init__.py

@@ -10,6 +10,7 @@
     NestedMutableMappingNode,
     flatten_dict,
     get_deep,
+    map_leaves,
     set_deep,
     unflatten_dict,
 )
@@ -24,6 +25,7 @@
     "NestedMutableMappingNode",
     "flatten_dict",
     "get_deep",
+    "map_leaves",
     "set_deep",
     "unflatten_dict",
 ]

src/nested_dict_tools/core.py

@@ -19,7 +19,7 @@
 This code is licensed under the terms of the MIT license.
 """
 
-from collections.abc import Iterable, Iterator, Mapping, MutableMapping, Sequence
+from collections.abc import Callable, Iterable, Iterator, Mapping, MutableMapping, Sequence
 from typing import Any, Literal, cast, overload
 
 type NestedMapping[K, V] = Mapping[K, NestedMappingNode[K, V]]
@@ -205,3 +205,98 @@ def set_deep[K, V](d: NestedMutableMapping[K, Any], keys: Sequence[K], value: An
             sub_dict[key] = sub_dict = {}
 
     sub_dict[keys[-1]] = value
+
+
+@overload
+def map_leaves[K, V, W](
+    func: Callable[[V], W], nested_dict1: NestedMapping[K, V], /
+) -> NestedMapping[K, W]: ...
+
+
+@overload
+def map_leaves[K, V1, V2, W](
+    func: Callable[[V1, V2], W],
+    nested_dict1: NestedMapping[K, V1],
+    nested_dict2: NestedMapping[K, V2],
+    /,
+) -> NestedMapping[K, W]: ...
+
+
+@overload
+def map_leaves[K, V1, V2, V3, W](
+    func: Callable[[V1, V2, V3], W],
+    nested_dict1: NestedMapping[K, V1],
+    nested_dict2: NestedMapping[K, V2],
+    nested_dict3: NestedMapping[K, V3],
+    /,
+) -> NestedMapping[K, W]: ...
+
+
+@overload
+def map_leaves[K, V1, V2, V3, V4, W](
+    func: Callable[[V1, V2, V3, V4], W],
+    nested_dict1: NestedMapping[K, V1],
+    nested_dict2: NestedMapping[K, V2],
+    nested_dict3: NestedMapping[K, V3],
+    nested_dict4: NestedMapping[K, V4],
+    /,
+) -> NestedMapping[K, W]: ...
+
+
+@overload
+def map_leaves[K, V1, V2, V3, V4, V5, W](
+    func: Callable[[V1, V2, V3, V4, V5], W],
+    nested_dict1: NestedMapping[K, V1],
+    nested_dict2: NestedMapping[K, V2],
+    nested_dict3: NestedMapping[K, V3],
+    nested_dict4: NestedMapping[K, V4],
+    nested_dict5: NestedMapping[K, V5],
+    /,
+) -> NestedMapping[K, W]: ...
+
+
+@overload
+def map_leaves[K, W](
+    func: Callable[..., W],
+    nested_dict1: NestedMapping[K, Any],
+    nested_dict2: NestedMapping[K, Any],
+    nested_dict3: NestedMapping[K, Any],
+    nested_dict4: NestedMapping[K, Any],
+    nested_dict5: NestedMapping[K, Any],
+    /,
+    *nested_dicts: NestedMapping[K, Any],
+) -> NestedMapping[K, W]: ...
+
+
+def map_leaves[K, V, W](
+    func: Callable[..., W],
+    *nested_dicts: NestedMapping[K, V],
+) -> NestedMapping[K, W]:
+    """
+    Apply the function to every leaf (non-mapping values) of the nested dictionaries.
+
+    If multiple nested dictionaries are passed, performs element-wise operations on their corresponding values at each key.
+
+    Args:
+        func: Function to apply on the leaves.
+        *nested_dicts: Nested dictionaries on which to apply the function.
+
+    Return:
+        The result nested dictionary with mapped leaves.
+
+    >>> map_leaves(lambda x: x * 2, {"a": 1, "b": 2, "c": 3})
+    {'a': 2, 'b': 4, 'c': 6}
+
+    >>> map_leaves(lambda x, y: x + y, {"a": 1, "b": 2}, {"a": 3, "b": 4})
+    {'a': 4, 'b': 6}
+    """
+    dict_res: NestedMapping[K, W] = {}
+    dict1 = nested_dicts[0]
+    for key in dict1:
+        args = (d[key] for d in nested_dicts)
+        if isinstance(dict1[key], Mapping):
+            dict_res[key] = map_leaves(func, *cast(Iterator[NestedMapping[K, V]], args))
+        else:
+            dict_res[key] = func(*cast(Iterator[V], args))
+
+    return dict_res

tests/test_core.py

@@ -1,6 +1,7 @@
 """Tests for core module."""
 
 import math
+import operator
 import time
 
 import pytest
@@ -9,6 +10,7 @@
     KeySeparatorCollisionError,
     flatten_dict,
     get_deep,
+    map_leaves,
     set_deep,
     unflatten_dict,
 )
@@ -393,3 +395,107 @@ def test_set_deep_with_large_dictionary(self):
         for key in keys[:-1]:
             sub_dict = sub_dict[key]
         assert sub_dict[keys[-1]] == value
+
+
+class TestMapLeaves:
+    # Function correctly maps single-level dictionary with single input dictionary
+    def test_map_single_level_dict(self):
+        input_dict = {"a": 1, "b": 2, "c": 3}
+        expected = {"a": 2, "b": 4, "c": 6}
+        result = map_leaves(lambda x: x * 2, input_dict)
+        assert result == expected
+
+    # Empty input dictionary
+    def test_map_empty_dict(self):
+        input_dict = {}
+        expected = {}
+        result = map_leaves(lambda x: x * 2, input_dict)
+        assert result == expected
+
+    # Function correctly maps nested dictionary with multiple input dictionaries
+    def test_map_leaves_with_multiple_dicts(self):
+        # Define the input dictionaries
+        dict1 = {"a": 1, "b": {"c": 2, "d": 3}}
+        dict2 = {"a": 4, "b": {"c": 5, "d": 6}}
+
+        # Expected output after applying the function
+        expected_output = {"a": 5, "b": {"c": 7, "d": 9}}
+
+        # Call the map_leaves function and assert the result
+        result = map_leaves(operator.add, dict1, dict2)
+        assert result == expected_output
+
+    # Dictionary with mixed types (mappings and non-mappings) at same level
+    def test_map_leaves_with_mixed_types(self):
+        input_dict = {
+            "a": {"b": 1, "c": {"d": 2}},
+            "e": 3,
+            "f": {"g": 4, "h": {"i": 5}},
+        }
+        expected_output = {
+            "a": {"b": 2, "c": {"d": 4}},
+            "e": 6,
+            "f": {"g": 8, "h": {"i": 10}},
+        }
+
+        result = map_leaves(lambda x: x * 2, input_dict)
+        assert result == expected_output
+
+    # Dictionaries with different structures/missing keys
+    def test_map_leaves_with_different_structures(self):
+        dict1 = {"a": 1, "b": {"c": 2}}
+        dict2 = {"a": 3, "b": {"d": 4}}
+
+        with pytest.raises(KeyError):
+            map_leaves(operator.add, dict1, dict2)
+
+    # Deep recursion with many nested levels
+    def test_map_leaves_deep_recursion(self):
+        # Create a deeply nested dictionary
+        depth = 100
+        nested_dict = current_level = {}
+        for i in range(depth):
+            current_level[f"level_{i}"] = {}
+            current_level = current_level[f"level_{i}"]
+        current_level["value"] = 1
+
+        # Define a simple function to apply
+        def increment(x):
+            return x + 1
+
+        # Apply map_leaves with deep recursion
+        result = map_leaves(increment, nested_dict)
+
+        # Verify the result
+        current_level = result
+        for i in range(depth):
+            current_level = current_level[f"level_{i}"]
+        assert current_level["value"] == 2
+
+    # Non-commutative operations with multiple dictionaries
+    def test_non_commutative_operations(self):
+        dict1 = {"a": 1, "b": 2}
+        dict2 = {"a": 3, "b": 4}
+
+        expected_output = {"a": -2, "b": -2}
+        result = map_leaves(operator.sub, dict1, dict2)
+
+        assert result == expected_output
+
+    # Function returning different type than input values
+    def test_map_leaves_with_type_conversion(self):
+        # Define a function that changes the type of the input value
+        def to_string(x):
+            return str(x)
+
+        # Input nested dictionary with integer values
+        input_dict = {"a": 1, "b": {"c": 2, "d": 3}}
+
+        # Expected output where all integer values are converted to strings
+        expected_output = {"a": "1", "b": {"c": "2", "d": "3"}}
+
+        # Apply map_leaves with the to_string function
+        result = map_leaves(to_string, input_dict)
+
+        # Assert that the result matches the expected output
+        assert result == expected_output