関数をcontextmanager, decoratorの両方に対応させることはできるか?

Question

関数の実行時間を計る@stopwatchデコレータと、withブロック単位で処理にかかった時間を計るtimerコンテキストマネジャがあります。
どちらがデコレータでどちらがコンテキストマネジャか分からなくなるので名前を統一したいのですが、御存知の通りPythonでは関数のオーバーロードはできません。
文脈によって内部で条件分岐できないかと考えてみたのですが、うまい方法が思いつきません。
どなたかお力を貸していただけないでしょうか?

[Wandbox]三へ( へ՞ਊ ՞)へ ﾊｯﾊｯ

import time
from contextlib import contextmanager
from functools import wraps
from typing import Callable


def stopwatch(func: Callable) -> Callable:
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        elapsed_time = time.time() - start
        print(f"{func.__name__}: {elapsed_time}")
        return result

    return wrapper


@contextmanager
def timer(label: str):
    start = time.time()
    yield
    elapsed_time = time.time() - start
    print(f"{label}: {elapsed_time}")


@stopwatch
def sleep():
    time.sleep(0.1)


def main():
    sleep()

    with timer("sleep_with"):
        time.sleep(0.1)


if __name__ == "__main__":
    main()

Answer 1

contextlib.ContextDecorator がそのためのものかと

コンテキストマネージャをデコレータとしても使用できるようにする基底クラスです。

Comments

[退会済みユーザー]

きちんと用意されていたのですね。ContextDecoratorを使った書き方も考えてみます。
ありがとうございました。

Answer 2

デコレータは 別の関数を返す関数 でしか御座いませんが、 with 文で使用するコンテキストマネージャの実体は __enter__ 及び __exit__ という 二つの dunder methods を持つオブジェクト です。関数を @contextlib.contextmanager で wrap するのは、あくまでこのコンテキストマネージャ型オブジェクトを 関数から生成することができるデコレータ によって生成している (正確には、それを生成するファクトリ関数を生成している) に過ぎません。

そのため、どうしても文脈によって使われ方を判断したいのであれば、あるオブジェクトが デコレータとして使われているのか、それともコンテキストマネージャとして使われているのかは、呼び出される dunder method によって区別することが可能 かと存じます。具体的には、 @contextlib.contextmanager を使わず、 Callable 扱いで __call__ が呼ばれた場合はデコレータとして、コンテキストマネージャ扱いで __enter__ が呼ばれた場合はコンテキストマネージャとして動作するようなオブジェクトを作ればよいわけですから、次のような実装が求めるものに近しいのではないかと思われます。

import time
from functools import wraps
from typing import Callable


class stopwatch:
    def __init__(self, label=None):
        self._label = label

    def __call__(self, func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = time.time()
            result = func(*args, **kwargs)
            elapsed_time = time.time() - start
            print(f"{func.__name__}: {elapsed_time}")
            return result

        return wrapper

    def __enter__(self):
        self._start = time.time()

    def __exit__(self, exc_type, exc_value, traceback):
        elapsed_time = time.time() - self._start
        print(f"{self._label}: {elapsed_time}")


@stopwatch()
def sleep():
    time.sleep(0.1)


def main():
    sleep()

    with stopwatch("sleep_with"):
        time.sleep(0.1)


if __name__ == "__main__":
    main()

但し、上記コードで @stopwatch() としている通り、 あくまで区別が可能となるのはオブジェクトを生成した後 なので、デコレータとして使うときも新規のオブジェクトを生成する必要性が出て参ります。もしもデコレータは毎回生成せずに静的な Callable として、コンテキストマネージャは label 付で毎回生成する型として実装することも求めるとなると、その場合はデコレータ呼出しとコンテキストマネージャ生成がどちらも関数呼び出し (またはオブジェクト生成) の文脈となってしまい、区別することは出来ないのではないかと思われます。その意味では、厳密には御質問の「関数 を文脈から両方に対応させる」という要件自体は (リフレクションを用いて頑張るのでなければ) 不可能と言えるかもしれません。

そのことを考えれば、文脈を踏まえて、といった考え方では無く、単純に どちらも関数の呼び出しであり、引数が異なるのだから、引数で判別すればよい と考えた方が、実装は単純で分かりやすくなるかもしれません。デコレータとして使われる場合、引数に Callable が渡されますから、その場合はデコレータとして、そうでない場合はコンテキストマネージャのファクトリとして振舞うような分岐を入れれば、これでも概ね求める要件は満たせるように思います。

import time
from contextlib import contextmanager
from functools import wraps
from typing import Callable


def stopwatch(func_or_label):
    def decorator_stopwatch(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = time.time()
            result = func(*args, **kwargs)
            elapsed_time = time.time() - start
            print(f"{func.__name__}: {elapsed_time}")
            return result

        return wrapper

    @contextmanager
    def contextmanager_stopwatch(label: str):
        start = time.time()
        yield
        elapsed_time = time.time() - start
        print(f"{label}: {elapsed_time}")

    if callable(func_or_label):
        return decorator_stopwatch(func_or_label)
    else:
        return contextmanager_stopwatch(func_or_label)

Comments

[退会済みユーザー]

解答のコードに至るまでの考え方まで詳しく教えていただいて、ありがとうございます。
勉強になりました。

質問とは多少逸れるのですが、内部でyieldするcontextmanager_stopwatch()に、もし返り値のtype hintsを付けるとすれば何になるのでしょうか?
with文内は何でもありうるので、詳細なしのGeneratorでしょうか?