要約

競プロ向けに、oj-bundleやwebpack、AtCoder Libraryのexpander.pyのようなツールのPython版を作りました。このツールを使うことで、それまでコピー＆ペーストで運用していたスニペット集をパッケージとして管理することができます。

使い方

インストール

pip install git+https://github.com/bayashi-cl/expander

実行コマンド

python -m expander <source file> [-o <output file>] [-m <expand module names...>]

例:

python -m expander main.py -o out.py -m yourlib

出力先の指定がない場合は標準出力に出力されます。

サンプル

展開前のコード

import sys

from byslib.core import IINF, MOD, debug, sinput
from byslib.data.union_find import UnionFindTree


def main() -> None:
    n, q = map(int, sinput().split())
    uft = UnionFindTree(n + 1)
    for _ in range(q):
        l, r = map(int, sinput().split())
        uft.union(l - 1, r)
    print("Yes" if uft.same(0, n) else "No")


if __name__ == "__main__":
    sys.setrecursionlimit(10**6)
    main()

展開後のコード(AtCoder)(若干のネタバレを含みます。)

パフォーマンス

Library Checkerで検証しました。

Problem: Shortest Path

CPython	Submission	Time
expander	https://judge.yosupo.jp/submission/79094	2745 ms
pure	https://judge.yosupo.jp/submission/79093	2970 ms

PyPy	Submission	Time
expander	https://judge.yosupo.jp/submission/79096	913 ms
pure	https://judge.yosupo.jp/submission/79097	918 ms

ライブラリをそのまま貼り付けたときより若干速くなっています。これは関数などの定義がグローバルからモジュール内に移っためだと考えられます。

技術的な話

内部の実装の話をします。興味がある方はどうぞ。

文字列からのimport

オンラインジャッジに提出するためには、ソースコードを一つのファイルにまとめる必要があります。しかし、パッケージ内のコードをそのまま貼り付けるだけでは、名前の衝突などの問題が発生する可能性があります。

そこで本ツールでは、コードを文字列として埋め込み、そこから動的にモジュールを作成するという手法をとっています。(ac-library-python内の実装を参考にしました。)

例として、以下のモジュールa.b.cを展開する場合を考えます。

import sys
import a.b.d
from a.e import f

def function() -> None:
    """docstring"""
    pass

この場合の手順は次のとおりです。

提出コード内にa.b = types.ModuleType("a.b")で動的なモジュールを作成する。
モジュール内のimport文のうち、更に展開が必要なもの(この場合ではimport a.c と from a.e import f)を依存先として記録し、コメントアウトしたうえで提出コードに文字列として書き込む。
依存先をa.b.__dict__に登録する
exec関数でコードを読み込む。このとき第2引数にa.b.__dict__を渡すことでモジュールに関数などが登録される。

提出コードは以下のようになります。

a.b = types.ModuleType("a.b")
_code_a_b = """
import sys
# import a.c
# from a.e import f

def function() -> None:
    \"""docstring\""" # <- 必要に応じてエスケープ処理をする
    pass
"""
a.b.__dict__["a.c"] = a.c
a.b.__dict__["f"] = a.e.f
exec(_code_a_b, a.b.__dict__)

import文を探す

コードからimport文を見つけるにはAST(抽象構文木)を使います。ast.NodeVisitorを継承したクラスに、適切な名前のメゾットを実装することで、コード内の全てのimport文に対してそのメゾットを適用することができます。これにより、import文が何行目にあるか、どのモジュールから何をimportしているのかを取得することができます。

依存関係を解決する

モジュールの作成とコードの読み込みはモジュール同士の依存関係に応じた適切な順番で行う必要があります。

モジュール作成の依存関係は例えば a.bがaに依存する、というようなもので、依存関係はグラフ用語で言うところの木になります。また、コードの読み込みの依存関係はimport文によるもので、その関係はDAG(Directed acyclic graph)になります。

よって、モジュールの作成は最上位パッケージを根とした深さ優先探索の行きがけ順、コードの読み込みはimport文による依存関係を深さ優先探索したときの帰りがけ順に行います。

importされているモジュール一覧を取得する

あらかじめ読み込むべき全てのモジュールを列挙できれば見通しよく探索することができます。標準ライブラリであるmodulefinderを使うことで、コード内でimportされているモジュールの一覧とそのファイルの場所を調べる事ができます。

使用例

from modulefinder import ModuleFinder
finder = ModuleFinder()
finder.run_script("path/to/source.py")
finder.modules # <- ここに情報が辞書として格納される

ライセンスへの対処

展開するパッケージがMITなどで公開されていた場合、提出コード内にライセンスの条文を含める必要があります。パッケージのライセンス情報などはimportlib.metadataを使って取得することができます。しかし、ここで必要なのはパッケージ名であり、これはモジュール名と異なる可能性があります(sklearnがモジュール名、scikit-learnがパッケージ名、など)。モジュール名からパッケージ名を取得するのは不可能なようなので、パッケージを全探索して目的のモジュールを見つけています。

おわりに

自作ライブラリの管理に悩まれているpythonistaの方に使っていただければ幸いです。

テストを書いたり、実際にコンテストで使ったりして確認はしていますが、もしバグ等をを見つけた方はissueなどで報告していただけると助かります。