Python で ACLibrary の modint を再現するために C++ Template もどきを作って CRTP もどきで実装した話

経緯いきさつ

筆者は普段 C++ を使っていますが、今日はエイプリルフールなので Pythonista になりました。
なので、 Python の記事（？）を書きます。

C++ と Python で modint クラスを作る話

私の所属する競プロサークルで C++ と Python の競プロライブラリを作ることになりました。
とりあえず modint は必要だろうということになり、私がその担当になりました。

「みんな ACLibrary の modint に慣れてるだろう」と思った私は、 ACLibrary の atcoder::static_modint および atcoder::dynamic_modint と同じようなインターフェースで作ることにしました。

CRTP で modint を実装した話

まずは C++ 版。
static_modint と dynamic_modint はほとんど同じなので、2回実装するのは面倒です。
しかし、継承を利用して基底クラスに実装すると、計算結果が基底クラスのオブジェクトになり、不便です。

そこで、 CRTP (Curiously Recurring Template Pattern) を使います！
これは、基底クラスのテンプレート引数に派生クラスを渡す方法です。
これにより、基底クラスで派生クラスのオブジェクトを扱えるようになります。

// 基底クラス
// テンプレート引数 Modint は CRTP で渡される派生クラスの型
template <class Modint>
class modint_base {
    // 派生クラスのオブジェクトを受け取って
    // 派生クラスのオブジェクトを返す処理を
    // 基底クラスに書ける！
    [[nodiscard]]
    friend constexpr Modint operator+(const Modint& x, const Modint& y) noexcept {
        return std::move(Modint{ x } += y);
    }

    // 共通部分...
};

template <std::uint32_t Mod>
class static_modint : public modint_base<static_modint<Mod>> {
    // 非共通部分...
};

template <int ID>
class dynamic_modint : public modint_base<dynamic_modint<ID>> {
    // 非共通部分...
};

完成したものがこちらです。

Python でも CRTP する話

C++ 版 modint が完成したので、次は Python 版に取り掛かります。
一回 C++ で作ってるので、 Python 版は C++ 版とおおむね同じ実装にして、楽しようとしました。

というわけで、 Python で CRTP をします！（？）

Python は type で動的に型を取れるので、コンストラクタで型を取得してあげれば簡単に CRTP （？）ができますね！

class ModintBase:
    def __init__(self, value: int = 0) -> None:
        self._mod: int
        self._value = value % self._mod
        self._type = type(self)  # 派生クラスの型を取得

    # 派生クラスのオブジェクトを受け取って
    # 派生クラスのオブジェクトを返す処理を
    # 基底クラスに書ける！
    def __add__(self, y: Self | int) -> Self:
        if isinstance(y, self._type):
            y = y._value
        return self._type(self._value + y)

    # 共通部分...

class StaticModint(ModintBase):
    # 非共通部分...

class DynamicModint(ModintBase):
    # 非共通部分...

Python で C++ Template もどきを作る話

さて、ここで困ったことに気が付きました。
Python には C++ のようなテンプレートはないのです。
C++ と同じ実装にしたせいで、 Mod や ID を渡す方法がなくなってしまいました。

仕方がないので、 Python で C++ のテンプレートを再現します！（？）

Python にはいい感じの型アノテーションと、デコレータという素晴らしい機能があるので、それらを使ってテンプレートっぽくします！
関数と区別するために、実体化は [] でできるようにしました。

from typing import Callable, Self

# テンプレートと、テンプレートの実体化を管理するクラス
class _Template[T, U, *TArgs]:
    def __init__(self, func: Callable[[U, *TArgs], T]) -> None:
        self.__func = func  # 実体化用関数
        self.__cache: dict[tuple[U, *TArgs], T] = {}  # 実体化の保存用
        self.__doc__ = func.__doc__

    # [] で実体化
    def __getitem__(self, args: U | tuple[U, *TArgs]) -> T:
        if not isinstance(args, tuple):
            args = args,
        if args not in self.__cache:
            result = self.__func(*args)
            self.__cache[args] = result
            return result
        return self.__cache[args]

def _template[T, U, *TArgs](func: Callable[[U, *TArgs], T]) -> _Template[T, U, *TArgs]:
    return _Template(func)


class ModintBase:
    # 共通部分...


# @_template を付けた関数の引数がテンプレート引数になり、戻り値が実体化になる
@_template
def StaticModint(Mod: int):
    assert Mod > 0

    # テンプレート引数に対してクラスを生成する
    class StaticModint(ModintBase):
        _mod = Mod

        @staticmethod
        def mod() -> int:
            return Mod
    return StaticModint


@_template
def DynamicModint(_: int):
    class DynamicModint(ModintBase):
        _mod = 998244353

        @classmethod
        def mod(cls) -> int:
            return cls._mod

        @classmethod
        def set_mod(cls, m: int) -> None:
            assert m > 0
            cls._mod = m
    return DynamicModint


# 実体化
Modint998244353 = StaticModint[998244353]
Modint1000000007 = StaticModint[1000000007]
Modint = DynamicModint[-1]


if __name__ == '__main__':
    x = Modint998244353(500000000)
    y = StaticModint[998244353](500000000)
    print(x + y)  # 1755647
    print(type(x) == type(y))  # True

完成したものがこちらです。

おわり

いかがでしたか？
Python 版のインターフェースを C++ 版に無理やりそろえようとした結果、調子に乗って変なコード書いちゃった話でした。

よい子の皆さんは、後で読む人が困らないようなコードを書きましょう。