PythonのEnumが遅い

みーが今作成しているPythonのソフトウェアが徐々にもっさりしてきたので，cProfileで原因を調べてみると，どうやらEnumがボトルネックになっていることが分かったので，Enumの速度についていろいろと検証してみました．

環境

• Python 3.8.1
• iMac Late 2015, 3.2GHz Core i5
• macOS Catalina

Enumの実験

EnumとしてIntFlagを利用し，ビットフラグを確認することを想定します．具体的に，黄色は赤の要素があるかを確認する'is_red()'関数をEnum, dictionary, intで実装して，速度を比較してみます．速度計測のためにtimeitを利用し，プロファイルを取得するためにcProfileを利用しました．

import cProfile
import timeit
import enum


class EnumColor(enum.IntFlag):
    # Color as Enum
    RED = enum.auto()
    GREEN = enum.auto()
    BLUE = enum.auto()
    YELLOW = RED | GREEN


assert EnumColor.RED == 1
assert EnumColor.GREEN == 2
assert EnumColor.BLUE == 4

# Color as dictionary
DictColor = {
    "RED": 1,
    "GREEN": 2,
    "BLUE": 4
}
DictColor["YELLOW"] = DictColor["RED"] | DictColor["BLUE"]

# Color as int
COLOR_RED = 1
COLOR_GREEN = 2
COLOR_BLUE = 4
COLOR_YELLOW = COLOR_GREEN | COLOR_RED


def is_enum_red(r):
    return r & EnumColor.RED


def check_enum_flag():
    t = timeit.timeit(
        "is_enum_red(EnumColor.YELLOW)",
        globals=globals(), number=100000
    )
    print(f"enum:{t}")


def is_dict_red(r):
    return r & DictColor["RED"]


def check_dict_flag():
    t = timeit.timeit(
        'is_dict_red(DictColor["YELLOW"])',
        globals=globals(), number=100000
    )
    print(f"dict:{t}")


def is_int_red(r):
    return r & COLOR_RED


def check_int_flag():
    t = timeit.timeit(
        "is_int_red(COLOR_YELLOW)",
        globals=globals(), number=100000
    )
    print(f"int:{t}")


with cProfile.Profile() as p:
    p.runcall(check_enum_flag)
p.dump_stats("enum.stats")

with cProfile.Profile() as p:
    p.runcall(check_dict_flag)
p.dump_stats("dict.stats")

with cProfile.Profile() as p:
    p.runcall(check_int_flag)
p.dump_stats("int.stats")

これを実行すると，以下のようになりました．

実装方法	時間 (秒)	比較
Enum	0.343	x10
dictionary	0.036	x1.16
int	0.031

profileをGraphVizで可視化しました．左から順に，Enum, dictionary, int で並べました．

gprof2dot -f pstats enum.stats > enum.dot
gprof2dot -f pstats dict.stats > dict.dot
gprof2dot -f pstats int.stats > int.dot

Enumの & 演算子を実行する時に，200000回 newが実行されており，このために約10倍遅くなっているのだと推測されます．

Lib/enum.pyを見てみます．

class EnumMetaのcall関数は以下のように実装されています．

def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
        if names is None:  # simple value lookup
            return cls.__new__(cls, value)
        return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)

Enumのメンバー変数を使ってビットフラグを確認する度に新たなEnumのインスタンスが生成されることになります．

Enumを使うと列挙する分かりやすい名前が付けられて便利ですが，何十万回と実行するような場合には不向きであることが分かりました．

ビットフラグを扱うならばEnumではなくintの変数を用いるのがよいと思います．ちなみに折衷案としてdictionaryを使った実装も試しましたが，1.16倍と意外に遅くないので，Enumを真似てdictionaryとして実装するのも悪くないと思います．