Pythonモジュールとしてのdigdagの型付け

## 概要

[Digdag](https://www.digdag.io/) はバッチ処理のワークフローを定義して定期的に実行するなどのタスク管理ツール[^1]．

[^1]: つまりは cron をより高級にしたようなもの．

公式ドキュメント：

https://docs.digdag.io/

具体的な処理はシェルコマンドや Python や Ruby などで記述されたスクリプトを呼び出すことで行うが，本記事では Python ライブラリとしての `digdag` の型付けについてまとめる．

## digdag の使い方

Digdag（アプリケーション）で Python を使用する際の基本的な処理は以下の2つ．

- Digdag からスクリプトを呼び出す
  - Python関数の引数は，`*.dig` ファイル内で定義された変数や組み込み変数で同名のものが自動的に渡される
- Python からデータを渡す
  - `digdag.env.store({'key': value})` で辞書に値を追加する

他に子タスクを生成する機能などがある（参考: https://docs.digdag.io/python_api.html ）[^2]．

[^2]: ソースを見るにまだ他にも関数があるのだがAPIリファレンスに記載されていない．引数の条件なども無く，見ればわかるやろで流されてる感じがあるのでもうちょっとちゃんと書いてほしい．

特に後者では `import digdag` として `digdag` モジュールを参照するのだが，linterが効かないので設定を行う．

## digdag モジュールの型ヒント

`digdag` モジュールはパッケージマネージャからインストールするのではなく、Digdag アプリから呼び出されたときに限り動的にモジュールがランタイム環境に生やされるという <font color="lightgray">邪悪な</font> 実装がされているので静的解析が効かない．
このため，`digdag` モジュールの型付けを行うにはスタブファイル（ロジックを省略して型情報のみを記載したファイル）を作成する必要がある．

[ソースコード](https://github.com/treasure-data/digdag/blob/master/digdag-standards/src/main/resources/digdag/standards/py/runner.py)を確認すると以下のようなクラス構造になっていることがわかるので，この通りにスタブを作成することで型チェックを通せるようになる．

```python:stubs/digdag.pyi
from collections import OrderedDict
from typing import Any, Dict


class Env:
    params: Any
    subtask_config: OrderedDict
    export_params: Dict[Any, Any]
    store_params: Dict[Any, Any]
    state_params: Dict[Any, Any]
    subtask_index: int

    def __init__(self, digdag_env_mod: Any) -> None:
        ...

    def add_subtask(self, function: Any = None, **params: Any) -> None:
        ...
 
    def export(self, params: Dict[Any, Any] = {}, **kwds: Any) -> None:
        ...

    def set_state(self, params: Dict[Any, Any] = {}, **kwds: Any) -> None:
        ...

    def store(self, params: Dict[Any, Any] = {}, **kwds: Any) -> None:
        ...

env: Env
```

上記は実装コードに [stubgen](https://mypy.readthedocs.io/en/stable/stubgen.html) を適用し，適当に整えたものである．
ここではパラメーターの型に制限がなく `Any` となっているが，実際にはデータのやり取りをJSON化して行うため，辞書はキーが `str` であるなどの制約をつけたほうがより正確な型チェックを行える．

```python:stubs/digdag.pyi 改
from collections import OrderedDict
from typing import Any, Dict, TypeAlias

DictStrToAny: TypeAlias = Dict[str, Any]

class Env:
    params: Any
    subtask_config: OrderedDict
    export_params: DictStrToAny
    store_params: DictStrToAny
    state_params: DictStrToAny
    subtask_index: int

    def __init__(self, digdag_env_mod: Any) -> None:
        ...

    def add_subtask(self, function: Any = None, **params: Any) -> None:
        ...
 
    def export(self, params: DictStrToAny = {}, **kwds: Any) -> None:
        ...

    def set_state(self, params: DictStrToAny = {}, **kwds: Any) -> None:
        ...

    def store(self, params: DictStrToAny = {}, **kwds: Any) -> None:
        ...

env: Env
```

valueの方は入れ子の辞書となることもあり得，ややこしくなるためここでは `Any` のままとした．また，より受け入れを広くして `collections.abc.Mapping` を使ったほうが適当かもしれない．
このあたりは各自で調整していただきたい．

### mypyの型チェック設定

mypyで使用するには，`stubs/digdag.pyi`に上記のスタブファイルを記述し，

```ini: mypy.ini
[mypy]
mypy_path = ./stubs
```

などのようにして指定する．
2つ目に示した，パラメータを `Dict[str, Any]` で受け取るバージョンで型チェックが通過することを確認している．

### VSCodeでの設定

環境設定でパスを通せば Pylance などからも認識され，`digdag` モジュールの型情報を認識できるようになる．（ただし，Python 環境にモジュールがインストールされているわけではないので，インポート解決はできないようである）

```json:.vscode/settings.json
{
    ...,
    "python.analysis.extraPaths": [
        "./stubs"
    ]
}
```

スタブ作成前：未知の状態で、シンタックスハイライトも効いていない状態から

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2760163/0c726565-d02b-4431-7e41-94e71a615cd8.png)

スタブ作成後：メソッドであることが認識され、インターフェースが表示されるようになった。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2760163/b4e3ab7b-4d92-e641-c3f0-97a7bbaa5454.png)


## 関連記事

https://qiita.com/hiroysato/items/d0fe5e2d88c267413a82

https://zenn.dev/fagai/articles/9ff7a008171ddd

https://qiita.com/HavenSpring/items/92ad70c33f73ab613d08