From 4fc27bb8180323182a571150527f522252240598 Mon Sep 17 00:00:00 2001 From: Milan Lukac Date: Wed, 17 Jul 2024 19:08:43 -0700 Subject: [PATCH] Scan context: support list keys in getter (#2135) --- soda/core/soda/scan.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/soda/core/soda/scan.py b/soda/core/soda/scan.py index e2495e486..074e81ebb 100644 --- a/soda/core/soda/scan.py +++ b/soda/core/soda/scan.py @@ -960,10 +960,22 @@ def get_all_checks_text(self) -> str | None: def has_soda_cloud_connection(self): return self._configuration.soda_cloud is not None - def scan_context_get(self, key: str, default: any = None) -> any: - return self.scan_context.get(key, default) + def scan_context_get(self, key: str | list[str], default: any = None) -> any: + + if type(key) == str: + return self.scan_context.get(key, default) + elif type(key) == list: + nested_dict = self.scan_context + for k in key: + if isinstance(nested_dict, dict) and k in nested_dict: + nested_dict = nested_dict[k] + else: + return default + return nested_dict + else: + raise TypeError("Key must be a string or a list of strings") - def scan_context_set(self, key: str | list, value: any, overwrite: bool = True): + def scan_context_set(self, key: str | list[str], value: any, overwrite: bool = True): dic = self.scan_context dic_key = key