From fc3b9b12dffaad49e243372dddc83f73aeb58ba8 Mon Sep 17 00:00:00 2001
From: Benjamin Wohlwend <beni@elastic.co>
Date: Wed, 13 Oct 2021 15:05:01 +0200
Subject: [PATCH] added utility function to get an item from a nested dict
 (#1358)

this was initially part of #1321
---
 elasticapm/utils/__init__.py | 10 ++++++++++
 tests/utils/tests.py         | 20 ++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/elasticapm/utils/__init__.py b/elasticapm/utils/__init__.py
index 12084c478..804ea57a8 100644
--- a/elasticapm/utils/__init__.py
+++ b/elasticapm/utils/__init__.py
@@ -193,3 +193,13 @@ def starmatch_to_regex(pattern):
         else:
             res.append(re.escape(c))
     return re.compile(r"(?:%s)\Z" % "".join(res), options)
+
+
+def nested_key(d: dict, *args):
+    for arg in args:
+        try:
+            d = d[arg]
+        except (TypeError, KeyError):
+            d = None
+            break
+    return d
diff --git a/tests/utils/tests.py b/tests/utils/tests.py
index 2271da14b..ae132279d 100644
--- a/tests/utils/tests.py
+++ b/tests/utils/tests.py
@@ -37,6 +37,7 @@
 from elasticapm.utils import (
     get_name_from_func,
     get_url_dict,
+    nested_key,
     read_pem_file,
     sanitize_url,
     starmatch_to_regex,
@@ -240,3 +241,22 @@ def test_read_pem_file_chain():
     with open(os.path.join(os.path.dirname(__file__), "..", "ca", "chain.crt"), mode="rb") as f:
         result = read_pem_file(f)
         assert result.endswith(b"\xc8\xae")
+
+
+@pytest.mark.parametrize(
+    "data,key,expected",
+    [
+        (None, "x", None),
+        ({}, "x", None),
+        ({"x": 1}, "x", 1),
+        ({"x": {"y": 1}}, "x.y", 1),
+        ({"x": 1}, "x.y", None),
+        ({"x": {"y": {}}}, "x.y.z", None),
+    ],
+)
+def test_nested_key(data, key, expected):
+    r = nested_key(data, *key.split("."))
+    if expected is None:
+        assert r is expected
+    else:
+        assert r == expected