Improve [Async]ContextDecorator type hinting

[ncoghlan]

Updated annotations allow access to __wrapped__ on decorated callables without complaints from typecheckers.

They also correctly indicate that the returned type is NOT the same as the input type.
Instead, it is a different callable type with the same call signature.

Closes #13403

[ncoghlan]

MyPy is not a fan:

stubs/decorator/decorator.pyi:69: error: Signature of "__call__" incompatible with supertype "ContextDecorator"  [override]
stubs/decorator/decorator.pyi:69: note:      Superclass:
stubs/decorator/decorator.pyi:69: note:          def [_P`-1, _R] __call__(self, func: Callable[_P, _R]) -> _WrappedCallable[_P, _R]
stubs/decorator/decorator.pyi:69: note:      Subclass:
stubs/decorator/decorator.pyi:69: note:          def [_C: Callable[..., Any]] __call__(self, func: _C) -> _C

[srittau]

That was my bad. Callable is a special form, not a class, but the following could work:

Suggested change

	class _WrappedCallable(Callable[_P, _R]):
	class _WrappedCallable(Generic[_P, _R]):

[ncoghlan]

I'll switch back to Generic to show the jax failure that version reported in the mypy_primer comment (mypy doesn't seem to consider _WrappedCallable[_P, _R] as conforming to _Callable[_P, _R])

[ncoghlan]

Although, thinking about it, the old signature here was genuinely wrong, since it claimed the returned callable type would be the same type as the input (and hence have any extra attributes, etc, that a callable subclass defined), rather than just returning a new callable with a compatible calling signature.

The newly introduced mypy complaints in jax reflect that it had passed along that same erroneous claim.

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

jax (https://github.com/google/jax)
+ jax/_src/api.py:2685: error: Incompatible return value type (got "_WrappedCallable[[VarArg(Any), KwArg(Any)], Any]", expected "F")  [return-value]
+ jax/experimental/jax2tf/jax2tf.py:3062: error: Argument 1 has incompatible type "_GeneratorContextManager[NameStack, None, None]"; expected "Callable[[partial[Sequence[Any]]], partial[Sequence[Any]]]"  [arg-type]
+ jax/experimental/jax2tf/jax2tf.py:3062: note: "_GeneratorContextManager[NameStack, None, None].__call__" has type "Callable[[Arg(Callable[_P, _R], 'func')], _WrappedCallable[_P, _R]]"

cki-lib (https://gitlab.com/cki-project/cki-lib)
- tests/test_metrics.py:75: error: Call to untyped function "dummy_function" in typed context  [no-untyped-call]

[ncoghlan]

These are the new type error reports that this change introduces in jax:

jax (https://github.com/google/jax)
+ jax/_src/api.py:2685: error: Incompatible return value type (got "_WrappedCallable[[VarArg(Any), KwArg(Any)], Any]", expected "F")  [return-value]
+ jax/experimental/jax2tf/jax2tf.py:3062: error: Argument 1 has incompatible type "_GeneratorContextManager[NameStack, None, None]"; expected "Callable[[partial[Sequence[Any]]], partial[Sequence[Any]]]"  [arg-type]
+ jax/experimental/jax2tf/jax2tf.py:3062: note: "_GeneratorContextManager[NameStack, None, None].__call__" has type "Callable[[Arg(Callable[_P, _R], 'func')], _WrappedCallable[_P, _R]]"

The first one is for a generic function with a F -> F signature along the same lines as the existing ContextDecorator.__call__ signature, which overpromises when it comes to the available functionality on the returned type. With this patch, MyPy now complains about that case (while neither typing nor typing_extensions provide a public WrappedCallable protocol to use to type it fully, it can at least be more correctly typed as accepting and emitting Callable[P, R] rather than using a single typevar).

The second one is related, but with slightly different symptoms, as it arises due to mypy inferring a narrow type for a sequence populated with partial objects, and now correctly detecting that the subsequent sequence transformation changes the type of the callables in the sequence. The fix would be to either use different names for the two sequences, or else declare the broader type for the first sequence.

[srittau]

Why did you remove __call__ here? The decorator library explicitly adds it and it's not part of the inheritance tree:

https://github.com/micheles/decorator/blob/519cc713878df4bfb768edfaa65e4bf2d875e421/src/decorator.py#L314-L318

[ncoghlan]

Because the stub tests failed complaining that the __call__ definition here no longer conformed to the parent class declaration of ContextDecorator.__call__ (it was the same issue as in jax: the return type isn't the same as the input type, it just has the same call signature).

Deleting the stub in the subclass means it inherits the ContextDecorator.__call__ signature, which is also correct for decorator.ContextManager.

I didn't check if the nominal signature of decorator.decorate itself might need updating, though.

[ncoghlan]

I edited the PR description to include that additional information about the change fixing a genuine bug in the old signature definition.

[srittau]

Whoops, I missed some of the base classes of _GeneratorContextManager when checking (twice!).

[srittau]

Thanks!

[ncoghlan]

I submitted a follow up PR to jax here: jax-ml/jax#25994