Pythonでデコレーターを使う際に引数情報やdocstring情報などがエディタやLint上で失われないようにするための小ネタです。

何が問題なのか

たとえば以下のような関数があったとします。

def sample_func(a: int, b: str = '') -> int:
    """
    Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do
    eiusmod tempor incididunt ut labore.

    Parameters
    ----------
    a : int
        Et dolore magna aliqua.
    b : str, optional
        Ut enim ad minim veniam, quis nostrud exercitation ullamco
        laboris

    Returns
    -------
    c : int
        Nisi ut aliquip ex ea commodo consequat.
    """
    return a * 2

これをVS Codeなどのエディタ上でPylanceなどの拡張機能を入れた状態で呼び出せば以下のように関数の引数情報やdocstringなどが表示されます。また、引数の型などに対するチェックも（チェックを有効にしていれば）リアルタイムに実行されます。

この関数だけではないですが任意の関数に設定するデコレーターを追加したとします。適当ですが大体以下のような記述になるかと思います。以下のデコレーター用の関数のサンプルのように通常、元の関数の引数内容は任意のものを受け付けられるように*argや**kwargsでの可変長の引数を設定することが多いと思います。

import functools
from typing import Any, Callable


def sample_decorator(f: Callable) -> Callable:
    """
    Duis aute irure dolor in reprehenderit in voluptate velit.

    Parameters
    ----------
    f : Callable
        対象の関数もしくはメソッド。

    Returns
    -------
    f : Callable
        デコレーター設定後の関数もしくはメソッド。
    """

    @functools.wraps(f)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        """
        Excepteur sint occaecat cupidatat non proident.

        Returns
        -------
        returned_value : Any
            関数もしくはメソッド実行後の返却値。
        """
        returned_value: Any = f(*args, **kwargs)
        return returned_value

    return wrapper

このデコレーターを関数に設定してみます。

@sample_decorator
def sample_func(a: int, b: str = '') -> int:
    """
    Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do
    eiusmod tempor incididunt ut labore.

    Parameters
    ----------
    a : int
        Et dolore magna aliqua.
    b : str, optional
        Ut enim ad minim veniam, quis nostrud exercitation ullamco
        laboris

    Returns
    -------
    c : int
        Nisi ut aliquip ex ea commodo consequat.
    """
    return a * 2

デコレーターを設定した関数を参照してみます。するとfunctoolsの機能でdocstring内容は保持されているものの引数のsignatureの情報が消えて可変長の引数の表示になってしまいました。返却値なども型情報が消えています（(...) -> Unknownといったような表示になっています）。

この状態だと引数情報が見れないだけでなく、Pylanceでのリアルタイムの型チェック（やCI/CDでのPyrightなども）でエラーで引っかからなくなります。例えばcという存在しない引数を指定してもエラーになりません。

sample_func(a=10, c=30)

型チェックでは引っかからないものの、この関数呼び出しを実行するとエラーになります。

    ...
    returned_value: Any = f(*args, **kwargs)
TypeError: sample_func() got an unexpected keyword argument 'c'

まあ単体テストなどをちゃんと書いていれば大半のケースではデプロイ前に気づきますが、実行時だけでなくエディタ上でもリアルタイムにエラーを表示してほしいところではあります。

デコレーターを関数から外すとこのような存在しない引数が指定されていた場合は（型チェックなどを有効化していれば）エラーがリアルタイムに表示されます。

これを解決してちゃんと引数情報を表示してくれて型チェックも走るようにしたいところです。

解決方法

いくつか解決方法はありますが、有力な解決策の一つとしてCallableやProtocolなどによる型アノテーションの代わりにジェネリックの型アノテーションを引数の関数（もしくはメソッド）と返却値へ設定するという方法があります。

TypeVarを使い、bound（ジェネリックの型に設定する制約設定）にCallableを指定します。今回はFという変数名にしました。

これをデコレーター用の関数の引数と返却値に設定します。

from typing import Any, Callable
from typing import TypeVar

F = TypeVar('F', bound=Callable)


def sample_decorator(f: F) -> F:
    ...

この場合型アノテーションとしては「引数のfはCallableである」という条件と「引数と返却値のCallableの引数構造は同じである」という条件になります。

ただし返却値の記述部分ではエラーになってしまいます。これは実際に返却されているのはwrapper関数なので引数構造が可変長の引数（*argsなど）になっている一方で引数に指定されたCallableは可変長の引数構造になるとは限らない（むしろ大半のケースでは別の構造になる）ため発生します。

...
    @functools.wraps(f)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        """
        Excepteur sint occaecat cupidatat non proident.

        Returns
        -------
        returned_value : Any
            関数もしくはメソッド実行後の返却値。
        """
        returned_value: Any = f(*args, **kwargs)
        return returned_value

    return wrapper

これは避けられないので普通にtype: ignoreを付けてしまいましょう（デコレーターの関数の定義の1か所のみなので）。これでPylanceやmypyで引っかからなくなります。

    return wrapper  # type: ignore

実際にデコレーター設定がされた関数で、存在しない引数が指定されているとエディタ上でエラーが表示されるようになりました。

sample_func(a=10, c=30)

また、関数参照時にもUnknownなどにならずに引数や返却値の構造の情報が表示されることが確認できます。

引数を受け付けるデコレーターの場合

引数を受け付ける形のデコレーターが必要なケースも結構あると思います。そのような場合は3重に関数を入れ子にしたりして対応しますが、そういったケースでも同様にジェネリックの指定で対応ができます。

3重の入れ子となるのでtype: ignoreの指定が1か所増えて2か所で指定する必要が出てきます。

例えば以下のようなコードになります。

import functools
from typing import Any, Callable
from typing import TypeVar

F = TypeVar('F', bound=Callable)


def sample_decorator(any_option: bool = True) -> F:

    def wrapper(f: F) -> F:
        """
        Duis aute irure dolor in reprehenderit in voluptate velit.

        Parameters
        ----------
        f : Callable
            対象の関数もしくはメソッド。

        Returns
        -------
        f : Callable
            デコレーター設定後の関数もしくはメソッド。
        """

        @functools.wraps(f)
        def inner_wrapper(*args: Any, **kwargs: Any) -> Any:
            """
            Excepteur sint occaecat cupidatat non proident.

            Returns
            -------
            returned_value : Any
                関数もしくはメソッド実行後の返却値。
            """
            returned_value: Any = f(*args, **kwargs)
            return returned_value

        return inner_wrapper  # type: ignore

    return wrapper  # type: ignore


@sample_decorator(any_option=True)
def sample_func(a: int, b: str = '') -> int:
    """
    Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do
    eiusmod tempor incididunt ut labore.

    Parameters
    ----------
    a : int
        Et dolore magna aliqua.
    b : str, optional
        Ut enim ad minim veniam, quis nostrud exercitation ullamco
        laboris

    Returns
    -------
    c : int
        Nisi ut aliquip ex ea commodo consequat.
    """
    return a * 2


sample_func(a=10, c=30)

このケースでも関数呼び出し部分でc引数の箇所で通常通り型のエラーが発生するのと、関数の引数構造も同様に表示されます。

また、Pylance（Pyright）などを使っている場合にはデコレーターの一番外側の以下の部分でもジェネリック関係の警告が出ます。

...
def sample_decorator(any_option: bool = True) -> F:
...

これに関しては気になる場合にはデコレーターの関数の前に# pyright: reportInvalidTypeVarUse=falseといったコメントを付けておくと無視できます。

該当のモジュール内のみ、且つそのコメント行以降が無視されるという挙動になるので、他の実装分はこのエラーや警告も無視せず正しくチェックして欲しいといった場合にはデコレーター定義用のモジュールを他と分けて、且つデコレーター用の関数もモジュール内の最後の方に定義しておくと良いかなと思います。

サードパーティーのライブラリでジェネリックの型アノテーションがされていない場合は・・・

サードパーティーのライブラリのデコレーター、特に古くからあるライブラリなどではこの辺の型アノテーションがされていなかったりで使用時に引数情報が消えてしまったりしてしまうケースがあります。

そういった場合には自分でラッパー的にデコレーター用の関数を用意して、対象のライブラリのデコレーター用の関数を内部で呼び出して使うことで回避することができます（まだそういった使い方をしていないので認識を間違えていなければ・・・）。

対象のライブラリのデコレーターが2重の関数になっていて関数を引数に取るものであれば関数を指定し、もしくはオプションを取る3重の関数構造になっている場合にはその関数をオプション指定しつつ呼び出せば対応ができます。

以下のコードではsample_decorator_1関数がライブラリのデコレーター、sample_decorator_2の関数が自前のラッパーとしてのデコレーターの関数と想定して対応しています。最後のデコレーター設定がされた関数の呼び出しでも引数情報などが失われていないことが確認できます。

import functools
from typing import Any, Callable
from typing import TypeVar

F = TypeVar('F', bound=Callable)


def sample_decorator_1(f: Callable) -> Callable:

    @functools.wraps(f)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        returned_value: Any = f(*args, **kwargs)
        return returned_value

    return wrapper


def sample_decorator_2(f: F) -> F:
    """
    Duis aute irure dolor in reprehenderit in voluptate velit.

    Parameters
    ----------
    f : Callable
        対象の関数もしくはメソッド。

    Returns
    -------
    f : Callable
        デコレーター設定後の関数もしくはメソッド。
    """

    @functools.wraps(f)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        """
        Excepteur sint occaecat cupidatat non proident.

        Returns
        -------
        returned_value : Any
            関数もしくはメソッド実行後の返却値。
        """
        returned_value: Any = sample_decorator_1(f=f)
        return returned_value

    return wrapper  # type: ignore


@sample_decorator_2
def sample_func(a: int, b: str = '') -> int:
    """
    Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do
    eiusmod tempor incididunt ut labore.

    Parameters
    ----------
    a : int
        Et dolore magna aliqua.
    b : str, optional
        Ut enim ad minim veniam, quis nostrud exercitation ullamco
        laboris

    Returns
    -------
    c : int
        Nisi ut aliquip ex ea commodo consequat.
    """
    return a * 2


sample_func(a=10)

スタブファイルを使う方法でも一応対応ができる

他の対策方法の1つとして、mypyなどによるスタブファイル生成によって引数情報のみのスタブファイルを生成することでデコレーター設定がされていても引数情報を保持する・・・という対応もあります。

ただしこちらはコードを変えたらスタブファイルも更新しないと古い構造のままエディタやLintなどが認識してしまうので同期の自動化していないと頻繁に更新するプロジェクト内のコードだと苦しい感じではあります。

また、mypyのスタブファイル生成のコマンドなどだとdocstringが消えたりするのでその辺は追加でdocstringを追加設定するなどの対応が必要になります。この辺は以前記事にしたのでそちらも必要に応じてご確認ください。

参考資料・関連資料