utamaro’s blog

誰かの役に立つ情報を発信するブログ

指定したキーをdictから再帰的に検索して削除する方法

Djangoを使っているなかで、apiのテストをする際に特定のkeyを取り除きたいということがありました。

例えば、レスポンスにcreated_atが含まれている場合、レコードが作成された日が入ると、assertDictEqualでは期待値と一致しないためエラーとなってしまいます。

その際にcreated_atを取り除くことが目的です。

これによって、一つ一つのkeyを取り出して期待値と一致するか検証する必要がなくなり、assertDictEqualを使って検証することができるようになります。

※ 値を取り除くよりも、freezegunを利用してdatetime.nowにパッチするのが良いと思われます。

コード

remove_keysを実行することでdictから指定したキーを取り除けます。

また、valueがlistやdictであった場合も対応できます。

def remove_keys(d: Dict, _remove_keys: List[str] = None):
    if _remove_keys is None:
        raise ValueError('keys must be required.')
    for k in list(d.keys()):
        v = d[k]
        if k in _remove_keys:
            del d[k]
        if isinstance(v, dict):
            remove_keys(v, _remove_keys)
        elif isinstance(v, list):
            for vd in v:
                remove_keys(vd, _remove_keys)

使い方

sample = {
    'id': 1,
    'name': 'sample',
    'description': 'sample',
    'created_at': int(datetime.now().timestamp()),
    'modified_at': int(datetime.now().timestamp()),
    'obj': {
        'id': 1,
        'name': 'sample',
        'description': 'sample',
        'created_at': int(datetime.now().timestamp()),
        'modified_at': int(datetime.now().timestamp()),
    },
    'list_obj': [
        {
            'id': 1,
            'name': 'sample',
            'description': 'sample',
            'created_at': int(datetime.now().timestamp()),
            'modified_at': int(datetime.now().timestamp()),
        },
        {
            'id': 2,
            'name': 'sample',
            'description': 'sample',
            'created_at': int(datetime.now().timestamp()),
            'modified_at': int(datetime.now().timestamp()),
        },
    ]
}
remove_keys(sample, ['created_at', 'modified_at'])
assert sample == {
    'id': 1,
    'name': 'sample',
    'description': 'sample',
    'obj': {
        'id': 1,
        'name': 'sample',
        'description': 'sample',
    },
    'list_obj': [
        {
            'id': 1,
            'name': 'sample',
            'description': 'sample',
        },
        {
            'id': 2,
            'name': 'sample',
            'description': 'sample',
        },
    ]
}