Improve type hinting for ContextDecorator and AsyncContextDecorator

[ncoghlan]

I recently ran into the problem described in https://stackoverflow.com/questions/62703400/python-how-to-type-hint-a-callable-with-wrapped, where MyPy complained when I attempted to access __wrapped__ on a function decorated with a ContextDecorator subclass.

The function signature is handled correctly (as per #4399), but the fact ContextDecorator.__call__ necessarily adds the __wrapped__ attribute to the returned wrapper is lost.

The functools stub does have some machinery for accurately typing functools.wraps itself, but there isn't anything that could be readily used to adjust the ContextDecorator.__call__ signature.

Ideally something like the following would be possible:

    def __call__(self, func: _F) -> _WrappedCallable[_F]: ...

However, I'm not sure how _WrappedCallable could be expressed as a generic type - the corresponding protocol in functools.pyi accepts the call parameters and return type as separate type variables.

[srittau]

Without giving it too much tought, using a separate ParamSpec/TypeVar could work:

_P = ParamSpec("_P")
_R = TypeVar("_R")

class _WrappedCallable(Generic[_P, _R]):
    __wrapped__: Callable[_P, _R]
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> _R: ...

class ContextDecorator:
    def __call__(self, func: Callable[_P, _R]) -> _WrappedCallable[_P, _R]: ...

This is basically the approach used for functools._Wrapped.

[ncoghlan]

#13416 implements the above suggestion.

[ncoghlan]

Iterating on the PR showed that there's actually a deeper problem with the current nominal signature here than not being able to access __wrapped__ without type checkers complaining: the existing signature claims something that isn't true.

Specifically, it claims that the returned value will have the same type as the input value. This is not accurate: the returned value will be a callable with the same call signature as the input value, but it will otherwise be a different type (so no extra attributes or methods that were defined on the original callable will be available).

[ncoghlan]

Working on the following PR to jax for this change, I ran into a type inferencing problem severe enough to convince me that this should be reverted until we have a way to fix it that doesn't lose the type overload details for the wrapped callable: #13416 (comment)

Discourse thread: https://discuss.python.org/t/defining-overload-preserving-signatures-for-wrapped-callables/78350

I'll post a reversion PR shortly (I would reopen the issue, but I don't have required permissions).

[ncoghlan]

@srittau As noted above, I think we should revert this fix until we have a way to do it that doesn't break return type inferencing: #13436

[bswck]

@ncoghlan Please reopen :)

[bswck]

Related: #13512