【Python】型安全に、省コードで、ディスクリプタを実装する

[以前書いた記事](https://qiita.com/junkmd/items/2d7fb2d5254fa8d12e5b)で、「インスタンス属性としてアクセスしたときの振る舞い」しか定義していないディスクリプタの[`__get__`メソッド](https://docs.python.org/ja/3/reference/datamodel.html#object.__get__)に「クラス属性としてアクセスしたときの振る舞い」を追加するデコレータについて、[クラス属性としてアクセスした際の型解釈がうまくされない](https://qiita.com/junkmd/items/2d7fb2d5254fa8d12e5b#%E3%83%9C%E3%83%84%E6%A1%88)という問題があったのですが、解決したので共有します。

使い方や結果はdocstringのとおりです。

```py
from collections.abc import Callable
import functools
from typing import Any, overload, TYPE_CHECKING, TypeVar


_This = TypeVar("_This")
_Inst = TypeVar("_Inst")
_Own = None | type[_Inst]
_R = TypeVar("_R")


# overloadされた関数を表現できる型アノテーション方法は存在しないため、返り値の型は型チェッカーに推論させる
def add_cls_attribute_behavior(dunder_get: Callable[[_This, _Inst, Any], _R]):
    """ディスクリプタの`__get__`メソッドへのデコレーター。

    インスタンス属性としてアクセスする振る舞いしか実装されていない`__get__`に対してデコレートすると、
    クラス属性としてアクセスしたときにディスクリプタインスタンス自身を返す振る舞いを追加する。

        >>> from typing import Protocol, TypeVar
        >>> _T = TypeVar('_T')
        >>> class _Component(Protocol[_T]):
        ...     wrp: _T
        ...
        >>> class UndecoratedDescriptor:
        ...     def __get__(self, instance, owner=None):
        ...         return instance
        ...
        >>> class DecoratedDescriptor:
        ...     @add_cls_attribute_behavior
        ...     def __get__(self, instance, owner=None):
        ...         return instance
        ...
        >>> class MyComponent:
        ...     undecorated_descriptor = UndecoratedDescriptor()
        ...     decorated_descriptor = DecoratedDescriptor()
        ...
        >>> my_comp = MyComponent()
        >>> assert my_comp.undecorated_descriptor is my_comp
        >>> assert MyComponent.undecorated_descriptor is None
        >>> assert my_comp.decorated_descriptor is my_comp
        >>> assert isinstance(MyComponent.decorated_descriptor, DecoratedDescriptor)

    `__get__`の引数に型アノテーションを付けることで、意図しないクラスにアサインされたディスクリプタ
    へアクセスしようとしたときに、静的型チェッカーがエラーを送出するように型制約をつけることができる。

        ```py
        from typing import Protocol, reveal_type, TypeVar

        _T = TypeVar("_T")

        class _Component(Protocol[_T]):
            wrp: _T

        class MyDescriptor:
            @add_cls_attribute_behavior
            def __get__(self, instance: _Component[int], owner=None):
                return instance

        class ComponentA:
            wrp: int

            my_descriptor = MyDescriptor()

        class ComponentB:
            wrp: str

            my_descriptor = MyDescriptor()

        reveal_type(ComponentA.my_descriptor)  # OK
        #           ^^^^^^^^^^^^^^^^^^^^^^^^
        #           Type of "ComponentA.my_descriptor" is "MyDescriptor"
        #
        reveal_type(ComponentA().my_descriptor)  # OK
        #           ^^^^^^^^^^^^^^^^^^^^^^^^^^
        #           Type of "ComponentA().my_descriptor" is "_Component[int]"
        #
        reveal_type(ComponentB.my_descriptor)  # NG
        #           ^^^^^^^^^^^^^^^^^^^^^^^^ Type of "ComponentB.my_descriptor" is "Unknown"
        #                      ^^^^^^^^^^^^^ (reportAttributeAccessIssue)
        # Cannot access member "my_descriptor" for type "type[ComponentB]"
        #   Failed to call method "__get__" for descriptor class "MyDescriptor"
        #
        reveal_type(ComponentB().my_descriptor)  # NG
        #           ^^^^^^^^^^^^^^^^^^^^^^^^^^ Type of "ComponentB().my_descriptor" is "Unknown"
        #                        ^^^^^^^^^^^^^ (reportAttributeAccessIssue)
        # Cannot access member "my_descriptor" for type "ComponentB"
        #   Failed to call method "__get__" for descriptor class "MyDescriptor"
        #
        ```
    """
    if TYPE_CHECKING:
        # fmt: off
        @overload
        def __get__(self: _This, instance: _Inst, owner: _Own[_Inst] = ..., /) -> _R: ...  # noqa
        @overload
        def __get__(self: _This, instance: None, owner: _Own[_Inst] = ..., /) -> _This: ...  # noqa
        # 実行時の実装は`else`以下が提供するが、空(`...`のみ)でもいいので実装部分を型定義しないと、
        # `@overload`がデコレートされている側に`"__get__" is marked as overload, but no
        # implementation is provided`が表示されたり、デコレートした`__get__`に`"__get__" is marked as
        # overload, but no implementation is provided`が送出される
        def __get__(self, instance, owner=None, /) -> Any: ...  # noqa
        # fmt: on
    else:
        # 実行時の実装。`functools.wraps`と`overload`が共存するとうまく型解釈されなかったので
        # 型定義と実行時の実装を分ける方法に至った。
        @functools.wraps(dunder_get)  # noqa
        def __get__(this, instance, owner=None):
            if instance is None:
                return this
            return dunder_get(this, instance, owner)

    return __get__
```

またこの実装でも、[以前の記事同様に](https://qiita.com/junkmd/items/2d7fb2d5254fa8d12e5b#%E5%AE%9F%E8%A3%85)ディスクリプタがクロージャを返す場合はそのdocstringをツールチップから参照することが可能です。

```py
from typing import Protocol, TypeVar

_T = TypeVar("_T")

class _Component(Protocol[_T]):
    wrp: _T

class MoveMethod:
    """class docstring"""
    @add_cls_attribute_behavior
    def __get__(self, instance: _Component[int], owner: type[_Component[int]] | None = None, /):
        """dunder_get method docstring"""
        def func(x: float, y: float, /) -> None:
            """callable docstring"""
            ...

        return func

class Sample:
    wrp: int

    Move = MoveMethod()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/364180/55ef1c10-afdd-45a8-81ce-0cb1e0e6a0c2.png)