Zenn
🔧

PydanticモデルからサブコマンドベースのCLIツールを生成するパッケージを作った

に公開
Python
機械学習
CLI
Pydantic
tech

PydanticモデルからサブコマンドベースのCLIツールを生成するpydantic-autocliというパッケージを作ったので紹介します。

endaaman/pydantic-autocli

気に入ったならリポジトリにStarを残していっていただければ大変嬉しく思います！

モチベーション

コマンドラインの引数の定義、面倒ですよね。PythonのargparseのAPI、変で使いにくいですよね。Pydanticで全部定義できればいいのにって思いませんか？

それを作りました。

インストール

pip install pydantic-autocli

最小限の実装

from pydantic import BaseModel
from pydantic_autocli import AutoCLI

class MyCLI(AutoCLI):
    class CustomArgs(BaseModel):
        # 必須パラメータ
        required_value: int
        # オプショナルパラメータ
        optional_value: int = 123
        # 配列パラメータ
        names: list[str] = []
        # フラグ
        flag: bool = False

    def run_simple(self, args: CustomArgs):
        print(f"Required: {args.required_value}")
        print(f"Optional: {args.optional_value}")
        print(f"Names: {args.names}")
        print(f"Flag: {args.flag}")

if __name__ == "__main__":
    cli = MyCLI()
    cli.run()

上記のように「型を指定する」ことで、関数内でDIのように自動的にインスタンス化されます。run_simpleのsimpleの部分がサブコマンド名として使えるようになり、以下のように実行できます。

# 必須パラメータを指定
python script.py simple --required-value 42 --flag

# すべてのパラメータを指定
python script.py simple --required-value 42 --names John Jane Bob

引数の型はstr int floatとその配列、boolをフラグとして使えます。

共通引数と詳細な設定

from pydantic import Field
from pydantic_autocli import AutoCLI, param

class MyCLI(AutoCLI):
    # 全コマンドで共有される引数
    class CommonArgs(AutoCLI.CommonArgs):
        verbose: bool = param(False, l="--verbose", s="-v", description="詳細な出力を有効化")
        seed: int = Field(42, json_schema_extra={"l": "--seed"})

    class AdvancedCommandArgs(CommonArgs):
        # ファイル名は--file-nameとして指定
        file_name: str = param(..., l="--name", pattern=r"^[a-zA-Z]+\.(txt|json|yaml)$")
        # 選択肢を制限
        mode: str = param("read", l="--mode", choices=["read", "write", "append"])
        # `Field`も直に使える
        wait: float = Field(0.5, json_schema_extra={"l": "--wait", "s": "-w"})

    async def run_advanced_command(self, args):
        print(f"File name: {args.file_name}")
        print(f"Mode: {args.mode}")
        
        if args.verbose:
            print("Verbose mode enabled")


if __name__ == "__main__":
    cli = MyCLI()
    cli.run()

こちらは型アノテーションを指定していませんが、「名前を合わせる」ことで同様に自動的にインスタンス化されます。ルールは run_advanced_commandのように二語なら、snake_caseがPascalCase変換され、AdvancedCommandにArgsが結合したAdvancedCommandArgsが対応します。

こちらは以下のように実行できます。

python script.py advanced-command --file-name data.txt --mode write --wait 1.5 -v

その他の細かい仕様としては、

  • param 関数を使って引数の詳細を定義します。
    • l="--long",s='-s'、のようにlong/shortの引数を定義できます。
    • Fieldの構文糖衣なので、Fieldも直接使えます。
    • 名前付き引数はすべて自動的にFieldに引き継がれるので、patternやltやgtなどすべて使えます。
  • CommonArgsは共通の引数としてフォールバックとして機能します。例のように--seedなどを設定するといいでしょう。
  • pre_common という関数が共通の初期化処理として使えます。

詳細は https://github.com/endaaman/pydantic-autocli のREADMEをご覧ください。まだ作ったばっかりなのでドキュメントの修正や補足のPRも大歓迎です。

開発エピソード

普段は機械学習を使った研究をしているのですが、

$ python experiment.py train --model resnet34 --batch-size 32

のように条件を引数にして実験を管理していました。Pydanticインスタンスをargparseに落とし込むアイデアはいろいろありますが、

  • サブコマンドに特化
  • argparseの挙動を完全に隠蔽

するようなライブラリはありませんでした。こういった背景から自作したcli.pyを使いまわしていました。

ただ毎回コピペするのもよろしくないし、ちゃんとパッケージ化してPyPIに置こうと思っていたところに、Cursorの学生1年無料のキャンペーンがあったので、これを使って今回パッケージ化作業を行いました。

まずは元からあったcli.pyを持ってきて、READMEに仕様をまとめます。その仕様を元にUnitテストを書かせて、内部を練り込みました。元は名前ベースの解決しかできなかったんですが、このタイミングでアノテーションで解決する機能も追加しました。型の解析周りは自分で書くには結構大変ですが、Sonnet3.7ならこれくらいは余裕のようですね。実装の詳細を読み解いてすらいませんが、Unitテストが通っているのでOKと思うことにしました。半日足らずで使えるものをPyPIにアップロードできました。

ちゃんとしたCLIツールを作る用というよりは、実験管理やバッチ処理をサブコマンド化するのにちょうどいい感じだと思うので、ぜひともご活用ください。

Discussion