Recursive dictionary merge in Python

dict_merge.py

# Recursive dictionary merge
# Copyright (C) 2016 Paul Durivage <[email protected]>
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

def dict_merge(base_dct, merge_dct):
    """ Recursive dict merge. Inspired by :meth:``dict.update()``, instead of
    updating only top-level keys, dict_merge recurses down into dicts nested
    to an arbitrary depth, updating keys. The ``merge_dct`` is merged into
    ``dct``.
    
    :param base_dct: dict onto which the merge is executed
    :param merge_dct: dct merged into dct
    :return: None
    """
    base_dct.update({
        key: dict_merge(rtn_dct[key], merge_dct[key])
        if isinstance(base_dct.get(key), dict) and isinstance(merge_dct[key], dict)
        else merge_dct[key]
        for key in merge_dct.keys()
    })

def dict_fmerge(base_dct, merge_dct, add_keys=True):
    """
    Recursive dict merge.

    Args:
        base_dct (dict) onto which the merge is executed
        merge_dct (dict): base_dct merged into base_dct
        add_keys (bool): whether to add new keys

    Returns:
        dict: updated dict
    """
    rtn_dct = base_dct.copy()
    if add_keys is False:
        merge_dct = {key: merge_dct[key] for key in set(rtn_dct).intersection(set(merge_dct))}

    rtn_dct.update({
        key: dict_fmerge(rtn_dct[key], merge_dct[key], add_keys=add_keys)
        if isinstance(rtn_dct.get(key), dict) and isinstance(merge_dct[key], dict)
        else merge_dct[key]
        for key in merge_dct.keys()
    })
    return rtn_dct


from unittest import TestCase

class DictMergeTestCase(TestCase):
    def test_merges_dicts(self):
        a = {
            "a": 1,
            "b": {"b1": 2, "b2": 3,},
        }
        b = {
            "a": 1,
            "b": {"b1": 4,},
        }

        assert dict_merge(a, b)["a"] == 1
        assert dict_merge(a, b)["b"]["b2"] == 3
        assert dict_merge(a, b)["b"]["b1"] == 4

    def test_inserts_new_keys(self):
        """Will it insert new keys by default?"""
        a = {
            "a": 1,
            "b": {"b1": 2, "b2": 3,},
        }
        b = {
            "a": 1,
            "b": {"b1": 4, "b3": 5},
            "c": 6,
        }

        assert dict_merge(a, b)["a"] == 1
        assert dict_merge(a, b)["b"]["b2"] == 3
        assert dict_merge(a, b)["b"]["b1"] == 4
        assert dict_merge(a, b)["b"]["b3"] == 5
        assert dict_merge(a, b)["c"] == 6

    def test_does_not_insert_new_keys(self):
        """Will it avoid inserting new keys when required?"""
        a = {
            "a": 1,
            "b": {"b1": 2, "b2": 3,},
        }
        b = {
            "a": 1,
            "b": {"b1": 4, "b3": 5,},
            "c": 6,
        }

        assert dict_merge(a, b, add_keys=False)["a"] == 1
        assert dict_merge(a, b, add_keys=False)["b"]["b2"] == 3
        assert dict_merge(a, b, add_keys=False)["b"]["b1"] == 4
        try:
            assert dict_merge(a, b, add_keys=False)["b"]["b3"] == 5
        except KeyError:
            pass
        else:
            raise Exception("New keys added when they should not be")

        try:
            assert dict_merge(a, b, add_keys=False)["b"]["b3"] == 6
        except KeyError:
            pass
        else:
            raise Exception("New keys added when they should not be")

    def test_ME(self):
        """Will it avoid inserting new keys when required?"""
        dict_of_listsA = {
            "LIST": [1, 2, 3],
            "DICT_SET": {
                "SET1": {"one", "two", "three"},
                "SET2": {"four", "five", "six"},
                "KeyA": "data",
            },
        }
        dict_of_listsB = {
            "LIST": ["one", "two", "three"],
            "DICT_SET": {
                "SET1": {1, "22", 3},
                "SET2": {"four", "five", "six"},
                "KeyB": "data",
            },
        }
        print(dict_merge(dict_of_listsA, dict_of_listsB, add_keys=True))

        assert (
            dict_merge(dict_of_listsA, dict_of_listsB, add_keys=False)["DICT_SET"].get(
                "KeyB"
            )
            is None
        )
        assert dict_merge(dict_of_listsA, dict_of_listsB, add_keys=True)["LIST"] == [
            "one",
            "two",
            "three",
        ]
        assert dict_merge(dict_of_listsA, dict_of_listsB, add_keys=True)["DICT_SET"][
            "SET1"
        ] == {1, "22", 3}
        assert (
            dict_merge(dict_of_listsA, dict_of_listsB, add_keys=True)["DICT_SET"][
                "KeyA"
            ]
            == "data"
        )
        assert (
            dict_merge(dict_of_listsA, dict_of_listsB, add_keys=True)["DICT_SET"][
                "KeyB"
            ]
            == "data"
        )
        assert (
            dict_merge(dict_of_listsA, dict_of_listsB, add_keys=False)["DICT_SET"].get(
                "KeyB"
            )
            is None
        )