Properly type a callable returning a Generator expression

I have the following snippet of code

from contextlib import _GeneratorContextManager, contextmanager

GoodWrapperType = Callable[[int, str], _GeneratorContextManager[None]]
BadWrapperType = Callable[[int, str], Generator[None, None, None]]

def wrapper() -> GoodWrapperType:

    @contextmanager
    def inner(some_int: int, some_str: str) -> Generator[None, None, None]:
        # Do things with some_int, some_str
        yield

    return inner

I want to do this in the context of a pytest testing suite, where the wrapper is getting some fixture injected into the inner.

I have above the GoodWrapperType, and the BadWrapperType. Pylance is telling me that I can't assign the BadWrapperType as the return type of my wrapper. I have found a solution with the GoodWrapperType, using _GeneratorContextManager, but since it's prefixed with an _, I am not supposed to be importing it.

Is there a better way ? What's the proper way of doing this ?

I was wondering if the fact that there's no straight-forward solution (that I've found anyway), might be a hint that I shouldn't do this in Python.

Solution

wrapper doesn't return callable returning a generator function, it returns a callable returning a context manager, which is a class with __enter__ and __exit__ methods.

@contextmanager decorator builds this class from the generator function it is applied to, but the end result is not a generator.

So we want to type inner as returning a generator function, but wrapper as returning a callable which returns the context manager that is the result of applying the decorator.

Here is a version which type-checks:

from contextlib import contextmanager, AbstractContextManager
from typing import Callable, Iterator, ContextManager


WrapperType = Callable[[int, str], AbstractContextManager[None]]


def wrapper() -> WrapperType:

    @contextmanager
    def inner(some_int: int, some_str: str) -> Iterator[None]:
        # Do things with some_int, some_str
        yield

    return inner

https://mypy-play.net/?mypy=latest&python=3.11&flags=strict&gist=6b2b3360a4916ec654a84ac84fcea544

If we're not using the send and return features of the generator we can simplify its type to just Iterator[<yield type>] as above.

AbstractContextManager[<enter return type>] is a generic type for context managers. (Note: pre Python 3.9 you should use from typing import ContextManager instead)

Since our inner function yields nothing, the type param in both cases is None.

However contextmanager returns a "context decorator", a dual-purpose class that can function as either a context manager or a decorator. Typing with AbstractContextManager will hide the decorator-ness of the result from downstream type checks.

If you need that property then your original approach with _GeneratorContextManager is probably the best way. (There is a public contextlib.ContextDecorator base class but mypy lacks intersection types so there's no convenient way to type something as both an AbstractContextManager & ContextDecorator using only the publicly exported types).