Commit da92c961 authored by Kenzo-Hugo Hillion's avatar Kenzo-Hugo Hillion
Browse files

add option to keep original source dict

parent fe154842
......@@ -4,7 +4,7 @@ logging.basicConfig(level=logging.INFO)
_LOGGER = logging.getLogger(__name__)
def extract_dict(source_dict, keys):
def extract_dict(source_dict, keys, keep_original_source=False):
"""
Extract a dict from a given dict based on a given set of keys
"""
......@@ -12,7 +12,8 @@ def extract_dict(source_dict, keys):
for key in keys:
try:
extracted_dict[key] = source_dict[key]
del source_dict[key]
if not keep_original_source:
del source_dict[key]
except KeyError:
_LOGGER.warning(f"[{key}] is not found in the source dict, extraction skipped for this key.")
return extracted_dict
......@@ -22,3 +22,12 @@ class TestExtractDict(TestCase):
test_extract_dict = extract_dict(source_dict, extract_keys)
self.assertDictEqual(source_dict, expected_source_dict)
self.assertDictEqual(test_extract_dict, expected_extract_dict)
def test_extract_dict_original_source(self):
source_dict = {'a': 1, 'b': 2}
extract_keys = ['b']
expected_source_dict = {'a': 1, 'b': 2}
expected_extract_dict = {'b': 2}
test_extract_dict = extract_dict(source_dict, extract_keys, keep_original_source=True)
self.assertDictEqual(source_dict, expected_source_dict)
self.assertDictEqual(test_extract_dict, expected_extract_dict)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment