dictキーを正規表現で検索したい

dict辞書は高速で非常に使い勝手が良い。
キーを正規表現で検索したい場合普通はこうすると思うが、多用する場合は面倒だ。

a = dict(abc=1, def=2) #-> {"abc": 1, "def": 2}
s = re.compile(".bc")
ret = []
for k in a:
    if s.search(k):
        ret.append(a[k])
print(ret) #-> [1]

# 内包表記で短くなってもせいぜいこんな感じだが、可読性が悪く、忘れたころにソースを読むと何をやりたかったのかよくわからん。
s = re.compile(".bc")
[a[k] for k in a if s.search(k)] #-> [1]

これが面倒なので標準のdictを拡張してこんな感じで使えるようにしてみた。

a = rdict(abc=1, def=2) #-> {"abc": 1, "def": 2}
a.search(".bc") #-> [1]

ソースはこう。

import re
from functools import lru_cache

@lru_cache(16) # re.compileは結構重めな処理なので、同じ正規表現パターンはキャッシュし、性能稼ぐ
def re_call(pattern, flags=0, call_re_func="search"):
    return re.compile(pattern, flags=flags).__getattribute__(call_re_func)

class rdict(dict):
    def _filter(self, _callable):
        return (k for k in self if _callable(k))

    def isin(self, key_or_function):
        if callable(key_or_function):
            return any(True for _ in self._filter(key_or_function))
        return dict.__contains__(self, key_or_function)

    def findall(self, key_or_function):
        if callable(key_or_function):
            return [dict.__getitem__(self, key) for key in self._filter(key_or_function)]
        return dict.__getitem__(self, key_or_function)

    def search(self, pattern, flags=0):
        return [dict.__getitem__(self,key) for key in self if re_call(pattern, flags, "search")(key)]

    def fullmatch(self, pattern, flags=0):
        return [dict.__getitem__(self,key) for key in self if re_call(pattern, flags, "fullmatch")(key)]

    def __setitem__(self, key_or_function, value):
        if callable(key_or_function):
            for key in self._filter(key_or_function):
                dict.__setitem__(self, key, value)
        else:
            return dict.__setitem__(self, key_or_function, value)

    def __delitem__(self, key_or_function):
        if callable(key_or_function):
            for key in list(self._filter(key_or_function)):
                dict.__delitem__(self, key)
        else:
            return dict.__delitem__(self, key_or_function)

その他にもいくつか関数を追加してるが使い方のイメージはこう。

正規表現パターンで一致するキーがあるかどうかを調べる


    >>> a.isin(re.compile(".b.*").search)
    True
    >>> a.isin(re.compile(".z.*").search)
    False

その他１、条件が真になる場合のキーを持つ場合の値を返す

    >>> a.findall(lambda x: len(x) == 3)
    [1, 2]

その他２、キーを範囲で検索し、値を返す
findallの引数はcallableなら、何でも動くので以下のような応用も可


    >>> from datetime import datetime
    >>> b = funcdict()
    >>> b[datetime(2020,1,1)] = "2020/01/01"
    >>> b[datetime(2020,2,1)] = "2020/02/01"
    >>> b[datetime(2020,3,1)] = "2020/03/01"

    >>> def between_0131_0202(x):
    ...    return datetime(2020,1,31) < x and x < datetime(2020,2,2)
    >>> b.findall(between_0131_0202)
    ['2020/02/01']

    >>> def less_0401(x):
    ...    return x < datetime(2020, 4, 1)
    >>> b.isin(less_0401)
    True

    >>> def grater_0401(x):
    ...    return x > datetime(2020, 4, 1)
    >>> b.isin(grater_0401)
    False

    >>> b.findall(less_0401)
    ['2020/01/01', '2020/02/01', '2020/03/01']

あとは、条件にマッチするキーの値を一括で変更する機能

    >>> b[less_0401] = "test"
    >>> b
    {datetime.datetime(2020, 1, 1, 0, 0): 'test',
     datetime.datetime(2020, 2, 1, 0, 0): 'test',
     datetime.datetime(2020, 3, 1, 0, 0): 'test'}

ついでに、条件にマッチするキーを一括で削除する機能

    >>> del b[between_0131_0202]
    >>> b
    {datetime.datetime(2020, 1, 1, 0, 0): 'test',
     datetime.datetime(2020, 3, 1, 0, 0): 'test'}