xor_bytesの高速化

websocketのフレームのマスク処理や暗号化処理などでバイナリデータ同士のxorが必要になることがあります。

Pythonのbytes型はxorを提供していないので自分で実装しないといけないのですが、この時 bytes([x ^ y for (x,y) in zip(data, mask)]) のような内包表記で行うことが多いようです。

https://sourcegraph.com/search?q=context:global+lang:Python+xor_bytes

実は Python 3 には int.from_bytes() と int.to_bytes() があり、うまく使えば内包表記よりも10倍以上高速化できます。

def xor_bytes_list(a, b):
    return bytes([(aa ^ bb) for (aa, bb) in zip(a, b, strict=True)])


def xor_bytes_generator(a, b):
    return bytes((aa ^ bb) for (aa, bb) in zip(a, b, strict=True))


def xor_bytes_via_int(a, b):
    if len(a) != len(b):
        raise ValueError(f"a and b must have same length; {len(a)=} {len(b)=}")
    aa = int.from_bytes(a)
    bb = int.from_bytes(b)
    return (aa ^ bb).to_bytes(len(a))

# MacBook Pro (M1 Pro) で計測
# 入力は256バイト

list: Mean +- std dev: 7.59 us +- 0.12 us
generator: Mean +- std dev: 10.1 us +- 0.3 us
int: Mean +- std dev: 685 ns +- 7 ns

bytes.to_bytes() の引数にもとの長さを指定してやるのがコツです。 エンディアンも指定できる（デフォルトは 'big'）ので、うまく使えば長さが違うときに前か後ろを0パディングするような動作も実装できるはずです。

追記: エンディアンはPython 3.11から big がデフォルトになったので、3.10まででも動くコードを書く場合は "big" を指定しましょう。

追記2: int を使った方法はPyPyで遅くなってしまいました。PyPyでもCPythonでも速くなる方法として bytes(map(operator.xor, a, b)) があります。