Subtype Generator with fixed type variables

Say I have a program with several generator functions that all have the return type Generator[<type>, <type2>, None] and I want to make an alias SimpleGenerator[<type>, <type2>] that is then expanded to the previous.

So for example

def my_generator() -> Generator[int, str, None]:
    ...

could be written

def my_generator() -> SimpleGenerator[int, str]:
    ...

I imagine the code might look something like (obviously incorrect python code)

T1 = TypeVar('T1')
T1 = TypeVar('T2')
SimpleGenerator['T1', 'T2'] = Generator['T1', 'T2', None]

But I’m not finding an easy way of writing it, and if going all the way of defining a class that inherits from Generator:

T1 = TypeVar("T1")
T2 = TypeVar("T2")
class SimpleGenerator(Generator["T1", "T2", None]):
    ...

def my_generator() -> SimpleGenerator[int, str]:
    ...

Then both mypy and pyright will not accept that as a return type for generators, saying

The return type of a generator function should be “Generator” or one of its supertypes

and

Return type of generator function must be “Generator” or “Iterable”

respectively. Is this possible to do?

—

(Originally I wanted Generator[<type>, None, None] and the pyright error message helped me realize I could then just use Iterable[<type>] as the return type. But my curiousity remains so I modified the question.)

Answer

Your generator is not a class, but a method, and you annotate it with your type. So you should do something like this:

from typing import TypeVar, Generic, 
Generator

T = TypeVar('T')

MyGenerator = Generator[T, None, None]

def my_generator() -> MyGenerator[str]:
    yield 'this'
    yield 'is'
    yield 'working'

print(' '.join(my_generator()))

Output:

python3.10 -m mypy test.py
Success: no issues found in 1 source file

python test.py
this is working

Advertisement

Answer