AI・機械学習チームの北川(@kitagry)です。 この記事はAI・機械学習チームブログリレー1日目の記事です。
この記事では、gokartの型チェックをmypy pluginからpyrightやpyrefly, tyなどのモダンな型チェッカーへ移行するために、SpotifyのluigiやFacebookのpyreflyといったOSSにコントリビュートした経緯を紹介します。
Pythonの__new__やDescriptor Protocol(__get__/__set__)など、型システムの深い部分にも触れていきます。
結果的に社内のプロダクトについてはmypyで1分10秒ほどかかっていましたが、pyright移行によって11秒ほどの高速化を実現しました。
- gokartと型の話
- mypy pluginの限界とpyright移行への動機
- luigi.StrParameter()をstrとして扱いたい：newアプローチ
- 文脈で型を使い分ける：get/setアプローチ
- pyrightへの移行、そして今後
- まとめ
- We are hiring
gokartと型の話
gokartはエムスリーがOSSとして公開しているPythonの機械学習パイプラインツールです。Spotifyが開発したluigiをベースにしており、タスクの依存関係を定義してパイプラインを構築できます。
gokartでは次のようにタスクを定義します。
class MyTask(gokart.TaskOnKart): param: str = luigi.StrParameter() def run(self): print(self.param) # strとして使える
luigi.StrParameter() はluigiのパラメータオブジェクトですが、タスクのインスタンスから self.paramでアクセスすると文字列として取得できます。これはluigiの内部でうまく変換されているためです。
Pythonの型アノテーション（param: str）を書くことで、self.param を str として扱えるようにしたいわけですが、実際にはluigi.StrParameter() を代入しているため、型チェッカーからすると「str に luigi.StrParameter()を代入している」と見えてしまい、型エラーになります。
mypy pluginの限界とpyright移行への動機
以前の記事でmypy pluginを使ってgokartを型安全にした話を紹介しました。
mypy pluginによって MyTask(name=1) のような型エラーを検出できるようになりましたが、mypy以外の後続のpyrightやpyrefly, tyといった型チェッカーでは動作しません。
実際の開発ではLSPを利用することが多いと思いますが、Language Serverとして開発されている上記のpyrightなどでは常にエラーが出ている状態になります。 CIの速度向上やClaude Codeなどとの相性の問題から移行を検討していましたが、mypy pluginに依存している限り移行できない状態でした。
そこで「mypy pluginなしでgokartを型安全にできないか」を検討し始めました。
luigi.StrParameter()をstrとして扱いたい：newアプローチ
__new__ アプローチのPRをluigiに送る
mypy pluginなしで型チェックを実現するアプローチとして、まず .pyiスタブファイルで__new__を使う方法を検討しました。
__new__メソッドはインスタンス生成時に呼ばれるメソッドで、その返り値の型が「このクラスのインスタンスの型」として型チェッカーに伝わります。これを利用して、次のように定義します。
class StrParameter: def __new__(cls, **kwargs) -> str: ....
こうすることで、luigi.StrParameter()はstrを、luigi.IntParameter()はintを返すものとして型チェッカーに認識されます。
ただし、実際の挙動とは異なるので.pyiファイルを利用して、型チェッカーに対して「このモジュールの型はこうなっている」と伝えつつも、実際の挙動には手を加えないアプローチをしました。
これによって、次のようなコードが型チェッカーに怒られないようになります。
s: str = luigi.StrParameter(default='a')
次に@dataclass_transformを利用することによって、引数の型チェックも行えるようにしました。
@dataclass_transform はPEP 681で定義されたデコレータで、型チェッカーに「このクラスはdataclassのように振る舞う」と伝えることができます。luigiの Task クラスに @dataclass_transform(field_specifiers=(Parameter,)) を付与することで、次のようなコードが型チェッカーに正しく認識されます。
class MyTask(gokart.TaskOnKart): param: str = luigi.StrParameter() MyTask(param="hello") # ✅ 型チェッカーがコンストラクタを認識できる
このPRはマージされ、luigi v3.7.0としてリリースされました。「これで解決だ！」と思ったのですが、実際に動かしてみると問題が残っていました。
pyrefly, tyのバグを踏む
luigiのPRがマージされたので、実際にpyrightとpyreflyで動作確認をしてみました。pyrightでは問題なく動いたのですが、pyreflyで次のようなエラーが出てしまいました。
.pyiでは Unpack と継承した TypedDict を組み合わせて各Parameterのキーワード引数を定義しています。
class _BaseParameterKwargs(TypedDict, Generic[T], total=False): default: T | None class _IntParameterKwargs(_BaseParameterKwargs[int], total=False): ... def __new__(cls, **kwargs: Unpack[_IntParameterKwargs]) -> int: ...
ところがpyreflyは継承したTypedDictのTypeVarを正しく解決できておらず、default: T | NoneのTがintに置換されずにエラーになっていました。
これはpyrefly側のバグだと判断し、PRを送りました。
こちらはマージされ、pyreflyでも正しく動作するようになりました。
またtyでも同様に__new__の返り値型がコンストラクタ呼び出しで使われないバグがあったのでPRを送りましたが、こちらは別のPRで対応されたためクローズになりました。
newアプローチの落とし穴：Parameter単体で使えない
luigiのPRもマージされ、pyreflyのバグも修正され、「ついに完成だ！」と思ったのですが、gokart側のテストコードが落ちてしまいました。
原因は __new__ アプローチの根本的な問題にありました。__new__ で返り値の型を変えるということは、型チェッカーからすると「このクラスのインスタンスはその型だ」と認識されます。つまり次のようになります。
# __new__アプローチでは... TaskInstanceParameter() # → TaskOnKartのインスタンスとして認識される # なのでこういうコードは型エラーになる gokart.TaskInstanceParameter().serialize(original) # ❌ TaskOnKartにserializeなんてない gokart.TaskInstanceParameter().parse(s) # ❌ TaskOnKartにparseなんてない
luigi.StrParameter()をタスクのパラメータとして使う分には問題ないのですが、gokartではParameterを独自に拡張してparseや serializeを直接呼ぶケースがありました。このようなParameterを「オブジェクトとして扱う」コードが全て型エラーになってしまいます。
この問題は__new__アプローチでは根本的に解決できません。「luigi.StrParameter()はstrだ」と型チェッカーに嘘をつく以上、Parameterとして扱うことができなくなるのは避けられないためです。
このようにタスクのパラメータとしてはstr型でそれ以外ではluigi.StrParameter()として扱う方法がないかを探していました。
文脈で型を使い分ける：get/setアプローチ
__get__/__set__ アプローチのPRをluigiに送る
__new__ アプローチの問題は「luigi.Parameter() は常に str だ」と型チェッカーに嘘をついていたことでした。必要なのは「タスクのインスタンスからアクセスするときは str、Parameter 単体として扱うときは Parameter」という文脈による使い分けです。
これはPythonのDescriptor Protocol（__get__/__set__）で実現できます。
class Parameter(Generic[T]): @overload def __get__(self, instance: None, owner: Any) -> "Parameter[T]": ... # クラスからアクセス → Parameter @overload def __get__(self, instance: Any, owner: Any) -> T: ... # インスタンスからアクセス → T def __set__(self, instance: Any, value: T): ... # 代入時の型はT
instance が None（クラスからのアクセス）かどうかで返り値の型が変わるため、型チェッカーは文脈に応じて正しい型を認識できます。
luigi.Parameter() # → Parameter[str]（クラスとして扱える） task_instance.param # → str（strとして扱える）
この仕様についてある程度理解した後、__get__/__set__ を実装したPRをluigiに送りました。
ただし、実装上の注意点が1つありました。
__set__ を定義すると Parameter がdata descriptorに昇格し、Pythonの属性参照の優先順位が変わります。
通常luigiはタスクの値を instance.__dict__ に直接書き込みますが、data descriptorはそれより先に参照されてしまいます。
そのため __get__ でも明示的に instance.__dict__ を参照するよう実装しています。
def __get__(self, instance: Any, owner: Any) -> Any: if instance is None: return self return instance.__dict__[self._attribute_name] # __dict__を明示的に参照 def __set_name__(self, owner, name): self._attribute_name = name def __set__(self, instance: Any, value: T): instance.__dict__[self._attribute_name] = value
またこの変更によって.pyiスタブファイルも不要になったため削除し、実装コード自体に型情報を持たせるシンプルな構成になりました。このPRもマージされました。
ただし、__new__ アプローチと比べて書き方が変わるというデメリットがあります。
# __new__アプローチ（旧） class MyTask(gokart.TaskOnKart): s: str = luigi.StrParameter() # strとアノテーションできた # __get__/__set__アプローチ（新） class MyTask(gokart.TaskOnKart): s: luigi.Parameter[str] = luigi.StrParameter() # Parameter[str]と書く必要がある
mypy pluginアプローチや__new__アプローチではParameterをstrに見せかけていたため s: str と書けましたが、__get__/__set__アプローチではParameterはあくまでParameterオブジェクトなので、型アノテーションもluigi.Parameter[str]と書く必要があります。
既存コードの書き換えが必要になる点は痛いところですが、ParameterをParameterとして正直に扱えるようになった分、型システムとしては健全です。
pyrightへの移行、そして今後
luigiへのコントリビュートによってpyright移行への道が開けましたが、gokart自体にはまだpyrightに対応できていない部分が多く、pyright移行は道半ばです。
一方で、gokartを使っている社内プロダクトではすでにpyrightへの移行を進めており、mypyで1分10秒かかっていた型チェックがpyrightでは11秒ほどで終わるようになりました。
gokart自体のpyright完全対応は引き続き進めていく予定です。
まとめ
今回はgokartの型安全性を高めるために、luigiやpyreflyといったOSSにコントリビュートした話を紹介しました。
型チェッカーの対応が不十分なときの選択肢として、mypy pluginのように自分たちのツール側で吸収する方法もありますが、upstreamのライブラリ自体を修正するという選択肢もあります。upstream修正はライブラリの全ユーザーにとってのメリットになりますし、自分たちのメンテナンスコストも下がります。
また、今回の試行錯誤を通じて __new__、Descriptor Protocol、@dataclass_transform といったPythonの型システムの深い部分への理解が深まりました。「なぜ型エラーになるのか」「型チェッカーはどう解釈しているのか」を理解することが、正しいアプローチを選ぶ上で重要だと改めて感じました。
gokart自体のpyright完全対応はまだですが、引き続きコントリビュートを続けていく予定です。興味がある方はぜひgokartのリポジトリを覗いてみてください！コントリビュートもお待ちしています！
We are hiring
弊社ではPythonの型によって安全なプロダクトを作成できるようなエンジニアを募集しています。 興味がある方は次のリンクから応募をお待ちしています。
