Implement issue #92

Author: pszulczewski

jdiff/extract_data.py

@@ -9,6 +9,7 @@
     jmespath_refkey_parser,
     associate_key_of_my_value,
     keys_values_zipper,
+    multi_reference_keys,
 )
 
 
@@ -43,6 +44,14 @@ def extract_data_from_json(data: Union[Mapping, List], path: str = "*", exclude:
         # return if path is not specified
         return data
 
+    # Multi ref_key
+    if len(re.findall(r"\$.*?\$", path)) > 1:
+        clean_path = path.replace("$", "")
+        values = jmespath.search(f"{clean_path}{' | []' * (path.count('*') - 1)}", data)
+        paired_key_value = associate_key_of_my_value(clean_path, values)
+        list_of_reference_keys = multi_reference_keys(path, data)
+        return keys_values_zipper(list_of_reference_keys, paired_key_value)
+
     values = jmespath.search(jmespath_value_parser(path), data)
 
     if values is None:

jdiff/utils/jmespath_parsers.py

@@ -8,6 +8,8 @@
 import re
 from typing import Mapping, List, Union
 
+import jmespath
+
 
 def jmespath_value_parser(path: str):
     """
@@ -120,3 +122,40 @@ def keys_values_zipper(list_of_reference_keys: List, wanted_value_with_key: List
         final_result.append({my_key: wanted_value_with_key[my_index]})
 
     return final_result
+
+
+def multi_reference_keys(jmspath, data):
+    """Build a list of concatenated reference keys.
+
+    Args:
+        jmspath: "$*$.peers.$*$.*.ipv4.[accepted_prefixes]"
+        data: tests/mock/napalm_get_bgp_neighbors/multi_vrf.json
+
+    Returns:
+        ["global.10.1.0.0", "global.10.2.0.0", "global.10.64.207.255", "global.7.7.7.7", "vpn.10.1.0.0", "vpn.10.2.0.0"]
+    """
+    ref_key_regex = re.compile(r"\$.*?\$")
+    mapping = []
+    split_path = jmspath.split(".")
+
+    ref_key_index = -1  # -1 as the starting value, so it will match split path list indexes
+    for index, element in enumerate(split_path):
+        if ref_key_regex.search(element):
+            ref_key_index += 1
+            key_path = (
+                ".".join(split_path[:index]).replace("$", "") or "@"
+            )  # @ is for top keys, as they are stripped with "*"
+            flat_path = f"{key_path}{' | []' * key_path.count('*')}"  # | [] to flatten the data, nesting level is eq to "*" count
+            sub_data = jmespath.search(flat_path, data)  # extract sub-data with up to the ref key
+            if isinstance(sub_data, dict):
+                keys = list(sub_data.keys())
+            elif isinstance(sub_data, list):
+                keys = []
+                for parent, children in zip(
+                    mapping[ref_key_index - 1], sub_data
+                ):  # refer to previous keys as they are already present in mapping
+                    keys.extend(f"{parent}.{child}" for child in children.keys())  # concatenate keys
+            else:
+                raise ValueError("Ref key anchor must return either a dict or a list.")
+            mapping.append(keys)
+    return mapping[-1]  # return last element as it as all previous ref_keys concatenated.

tests/test_get_value.py

@@ -44,6 +44,33 @@ def test_jmspath_return_none(jmspath):
         "vpn.peers.$*$.*.ipv6.[accepted_prefixes]",
         [{"10.1.0.0": {"accepted_prefixes": 1000}}, {"10.2.0.0": {"accepted_prefixes": 1000}}],
     ),
+    (
+        "$*$.peers.$*$.*.ipv4.[accepted_prefixes]",
+        [
+            {"global.10.1.0.0": {"accepted_prefixes": 1000}},
+            {"global.10.2.0.0": {"accepted_prefixes": 1000}},
+            {"global.10.64.207.255": {"accepted_prefixes": 1000}},
+            {"global.7.7.7.7": {"accepted_prefixes": 1000}},
+            {"vpn.10.1.0.0": {"accepted_prefixes": 1000}},
+            {"vpn.10.2.0.0": {"accepted_prefixes": 1000}},
+        ],
+    ),
+    (
+        "$*$.peers.$*$.*.ipv6.[accepted_prefixes]",
+        [
+            {"global.10.1.0.0": {"accepted_prefixes": 1000}},
+            {"global.10.2.0.0": {"accepted_prefixes": 1000}},
+            {"global.10.64.207.255": {"accepted_prefixes": 1000}},
+            {"global.7.7.7.7": {"accepted_prefixes": 1000}},
+            {"vpn.10.1.0.0": {"accepted_prefixes": 1000}},
+            {"vpn.10.2.0.0": {"accepted_prefixes": 1000}},
+        ],
+    ),
+    pytest.param(
+        "$*$.peers.$*$.address_family.$*$.[accepted_prefixes]",
+        "",
+        marks=pytest.mark.xfail(reason="Jmespath issue - path returns empty list."),
+    ),
 ]
 
 

tests/test_jmespath_parsers.py

@@ -5,8 +5,9 @@
     jmespath_refkey_parser,
     keys_values_zipper,
     associate_key_of_my_value,
+    multi_reference_keys,
 )
-from .utility import ASSERT_FAIL_MESSAGE
+from .utility import load_json_file, ASSERT_FAIL_MESSAGE
 
 
 value_parser_case_1 = (
@@ -112,3 +113,41 @@ def test_keys_zipper(ref_keys, wanted_values, expected_output):
 def test_keys_association(path, wanted_values, expected_output):
     output = associate_key_of_my_value(path, wanted_values)
     assert expected_output == output, ASSERT_FAIL_MESSAGE.format(output=output, expected_output=expected_output)
+
+
+multi_ref_key_case_1 = (
+    "$*$.peers.$*$.*.ipv4.[accepted_prefixes]",
+    ["global.10.1.0.0", "global.10.2.0.0", "global.10.64.207.255", "global.7.7.7.7", "vpn.10.1.0.0", "vpn.10.2.0.0"],
+)
+
+
+multi_ref_key_case_2 = (
+    "$*$.peers.$*$.address_family.$*$.[accepted_prefixes]",
+    [
+        "global.10.1.0.0.ipv4",
+        "global.10.1.0.0.ipv6",
+        "global.10.2.0.0.ipv4",
+        "global.10.2.0.0.ipv6",
+        "global.10.64.207.255.ipv4",
+        "global.10.64.207.255.ipv6",
+        "global.7.7.7.7.ipv4",
+        "global.7.7.7.7.ipv6",
+        "vpn.10.1.0.0.ipv4",
+        "vpn.10.1.0.0.ipv6",
+        "vpn.10.2.0.0.ipv4",
+        "vpn.10.2.0.0.ipv6",
+    ],
+)
+
+
+multi_ref_key_test_cases = [
+    multi_ref_key_case_1,
+    multi_ref_key_case_2,
+]
+
+
+@pytest.mark.parametrize("path, expected_output", multi_ref_key_test_cases)
+def test_multi_ref_key(path, expected_output):
+    data = load_json_file("napalm_get_bgp_neighbors", "multi_vrf.json")
+    output = multi_reference_keys(path, data)
+    assert expected_output == output, ASSERT_FAIL_MESSAGE.format(output=output, expected_output=expected_output)