python
`[Async]ContextDecorator.__call__` return type should depend on the underlying context manager's `_ExitT_co`
[!Important] Please see https://discuss.python.org/t/support-suppressing-generator-context-managers/103615 first!
This is how ContextDecorator.__call__ is implemented in the Python standard library:
def __call__(self, func):
@wraps(func)
def inner(*args, **kwds):
with self._recreate_cm():
return func(*args, **kwds)
return inner
In typeshed, ContextDecorator.__call__ currently has a return type of the original callable (func).
However, if the return value of self._recreate_cm().__call__ is truthy, the return type of ContextDecorator.__call__ can as well be None, like in the silly example below:
from collections.abc import Generator
from contextlib import contextmanager, suppress
from typing import reveal_type
@contextmanager
def zero_division_to_none() -> Generator[None]:
with suppress(ZeroDivisionError):
yield
@zero_division_to_none()
def div(a: int, b: int) -> float:
return a / b
reveal_type(div(1, 0))
# mypy, pyright: Revealed type is "builtins.float"
# Runtime type is 'NoneType'
My first shot was to patch the proper __call__ functions to "extend" the return type with _ThrottledExcReturnT (a new type parameter of ContextDecorator), constrained by Never and None, with the default type Never for backward compatibility. To me, there's no way to declare the relationship symbolically, so it would have to be typed manually in classes implementing the context manager protocol.
After looking at it now, ContextDecorator has an orig base of AbstractContextManager[_T_co, bool | None], therefore it logically implies that ContextDecorator.__call__ always returns an optional value (despite it not being represented by the code).
That is incorrect, because this div always returns a float:
@contextmanager
def zero_division_propagate() -> Generator[None]:
yield
@zero_division_propagate()
def div(a: int, b: int) -> float:
return a / b
and this div returns float | None:
@contextmanager
def zero_division_to_none() -> Generator[None]:
with suppress(ZeroDivisionError):
yield
@zero_division_to_none()
def div(a: int, b: int) -> float:
return a / b
The preservation of the original (potentially non-optional) return type can only be made if the context manager doesn't suppress the wrapped function's exception.
Ideally it would be specifiable whether a context manager swallows exceptions, since the type checker can't really see that. Something like @contextmanager(swallows_exceptions=True)? I don't like the idea of changing this API though.
I got it working with the suggested patch:
from collections.abc import Generator
from contextlib import contextmanager, suppress
from typing import TYPE_CHECKING, overload, reveal_type
@contextmanager
def zero_division_to_none() -> Generator[None]:
with suppress(ZeroDivisionError):
yield
if TYPE_CHECKING:
@overload
@zero_division_to_none()
def foobar_decorated(a: int, b: None) -> float:
...
@overload
@zero_division_to_none()
def foobar_decorated(a: int, b: int) -> int:
...
@zero_division_to_none()
def foobar_decorated(a: int, b: int | None) -> int | float:
if b is None:
return 0.0
return a // b
reveal_type(foobar_decorated)
# [runtime] Runtime type is 'function'
# [pyright] Type of "foobar_decorated" is "Overload[(a: int, b: None) -> (float | None), (a: int, b: int) -> (int | None)]"
# [mypy] Revealed type is "Overload(def (a: builtins.int, b: None) -> Union[builtins.float, None], def (a: builtins.int, b: builtins.int) -> Union[builtins.int, None])"
reveal_type(foobar_decorated(1, 0))
# [runtime] Runtime type is 'NoneType'
# [pyright] Type of "foobar_decorated(1, 0)" is "int | None"
# [mypy] Revealed type is "Union[builtins.int, None]"