Python nested dict diff

diff.py

"""
Hacky shell tool for finding dict deep diffs

a = {
    'b': {
        'c': 2, 
        'd': 3
    }, 
    'e': 2
}
b = {
    'b': {
        'c': 2, 
        'd': 4  # only diff
    }, 
    'e': 2
}

>>> flatten(a)
{'b_c': 2, 'b_d': 3, 'e': 2}

>>> diff(a, b)
{'b_d': '3 -> 4'}
"""

def flatten(data):
    if type(data) not in (dict, list, set):
        return {'': data}
    if type(data) == set:
        data = list(data)
    return_data = {}
    if type(data) == list:
        for i, d in enumerate(data):
            for k2, v2 in flatten(d).items():
                if k2:
                    return_data[f'{i}_{k2}'] = v2
                else:
                    return_data[f'{i}'] = v2
        return return_data
    for k, v in data.items():
        for k2, v2 in flatten(v).items():
            if k2:
                return_data[f'{k}_{k2}'] = v2
            else:
                return_data[f'{k}'] = v2
    return return_data

def diff(a, b):
    a = flatten(a)
    b = flatten(b)
    old = {}
    new = {}
    for k, v in a.items():
        if b.get(k) == v:
            continue
        old[k] = v
        new[k] = b.get(k)
    for k, v in b.items():
        if a.get(k) == v:
            continue
        old[k] = a.get(k)
        new[k] = v
    all_keys = set(list(old.keys()) + list(new.keys()))
    diff = {}
    for k in all_keys:
        diff[k] = f'{old.get(k)} -> {new.get(k)}'

    return diff