関数アノテーションを軽量化しました

この記事は KLab 2020 Advent Calendar の 12/2 分になります。 qiita.com

最近の Python に対する改善を紹介します。私が設計、コードレビューまでしましたが、実装は他のコントリビューターにしていただきました。 (プルリクエストはこちら)

背景として、Python 3.10 からは from __future__ import annotations がデフォルト化され、アノテーション部分は実行時に評価されずにただの文字列になります。( PEP 563 を参照してください。)

>>> def add(a: int, b: int) -> int:
...     return a+b
...
>>> add.__annotations__
{'a': 'int', 'b': 'int', 'return': 'int'}

アノテーションが実行時に評価されないということは、コンパイル時にアノテーションがすべて計算可能ということです。そこで Python 3.10 からさらなる最適化を導入しました。

まずは Python 3.9 で関数アノテーションがどう実装されていたかをおさらいします。

# a.py
from __future__ import annotations

def add(a: int, b: int) -> int:
    return a + b
$ python3 -m dis a.py
(snip)
  3          12 LOAD_CONST               2 ('int')
             14 LOAD_CONST               2 ('int')
             16 LOAD_CONST               2 ('int')
             18 LOAD_CONST               3 (('a', 'b', 'return'))
             20 BUILD_CONST_KEY_MAP      3
             22 LOAD_CONST               4 (<code object add at 0x7f7c9efac870, file "a.py", line 3>)
             24 LOAD_CONST               5 ('add')
             26 MAKE_FUNCTION            4 (annotations)
             28 STORE_NAME               2 (add)

4つのLOAD_CONST命令と1つのBUILD_CONST_KEY_MAP命令でdictを作って、それをMAKE_FUNCTIONに渡しています。この dict は function オブジェクトに格納されます。

これは関数の定義時に実行される命令なのでプログラムの実行速度への影響は軽微ですが、ループの中で inner function を定義している場合には毎回実行されますし、そうでなくても pyc ファイルを import する時間に多少影響を与える可能性があります。

これが Python 3.10 ではこうなります。

  3          12 LOAD_CONST               9 (('a', 'int', 'b', 'int', 'return', 'int'))
             14 LOAD_CONST               6 (<code object add at 0x7f9e932a7be0, file "a.py", line 3>)
             16 LOAD_CONST               7 ('add')
             18 MAKE_FUNCTION            4 (annotations)
             20 STORE_NAME               2 (add)

アノテーションが1つのタプルにまとめられ、このタプルが dict に変換されずそのまま MAKE_FUNCTION に渡され、関数オブジェクトに保存されるようになりました。これでアノテーション付きの関数を作る速度が数倍速くなります。 https://bugs.python.org/issue42202#msg381320 からマイクロベンチマーク結果を1つだけピックアップします。

def f(a: int, /, b: int, *, c: int, **d: int) -> None: pass

Python 3.9.0
5000000 loops, best of 5: 326 nsec per loop

Python 3.10.0a2+ (from __future__ import annotations がデフォルト化された状態)
5000000 loops, best of 5: 264 nsec per loop

Python 3.10.0a2+ with compact representation
5000000 loops, best of 5: 87.1 nsec per loop

さらにメモリ使用量にも影響があります。このタプルは func.__annotations__ が最初にアクセスされた時に dict に変換されるのですが、静的型チェックやコード補完機能、ドキュメントの目的で付けられた関数アノテーションは実行時には利用されないことが多いので、アクセスされないままならずっとタプルのままです。これでメモリ使用量は4割以下に減らせます。

>>> sys.getsizeof({"a":"int","b":"int","return":"int"})
232
>>> sys.getsizeof(("a","int","b","int","return","int"))
88

おまけに、同じソースファイルの中にある LOAD_CONST で読まれるタプルは、中身が同じであれば1つのインスタンスを使いまわします。 (Python 3.7までは1つのcodeオブジェクト(1つの関数など)内でしか使い回さなかったのですが、これも私が改良していました。 https://github.com/python/cpython/commit/c2e1607a51d7a17f143b5a34e8cff7c6fc58a091)

コードジェネレータに生成されたファイルなどで数千の同じシグネチャの関数がある場合は、今までは関数の数だけ dict オブジェクトを作っていたのが、 Python 3.10 ではたった1つのタプルで済むようになります。これでより気軽に関数アノテーションを使えるようになります。

将来の野望としては、関数アノテーションだけでなく docstring も、実際に利用されてからロードする遅延ロードが実装できたら良いなと思っています。しかしこれには pyc ファイルのフォーマットを完全に刷新する必要があり、しかもモジュールをロードしたあとに pyc ファイルが上書きされたらどうするのかという問題に対する解決策をまだ思いついていないので、 近い将来に実現するのは難しいです。いいアイデアを思いついたら教えて下さい。

このブログに乗せているコードは引用を除き CC0 1.0 で提供します。