How to write a meaningful typehint for a function that returns a filter?

Question

Assume the following function

def get_list_of_nums_and_square_nums_with_last_digit_9() -> List[Tuple[int, int]]:
    sqr = lambda n: n * n
    is_last_digit_9 = lambda n: n % 10 == 9
    is_last_digit_of_square_num_9 = lambda t: (lambda n, n_squared: is_last_digit_9(n_squared))(*t)

    nums_and_square_nums = map(lambda n: (n, sqr(n)), range(10))
    return list(filter(is_last_digit_of_square_num_9, nums_and_square_nums))

which would return

>>> get_list_of_nums_and_square_nums_with_last_digit_9()
[(3, 9), (7, 49)]

. Now I want an equivalent function, but returning a iterator - or more specifically a filter:

def get_nums_and_square_nums_with_last_digit_9():
    sqr = lambda n: n * n
    is_last_digit_9 = lambda n: n % 10 == 9
    is_last_digit_of_square_num_9 = lambda t: (lambda n, n_squared: is_last_digit_9(n_squared))(*t)

    nums_and_square_nums = map(lambda n: (n, sqr(n)), range(10))
    return filter(is_last_digit_of_square_num_9, nums_and_square_nums)

Which type hints can I use so that we can clearly see the structure of the result, when the resulting filter is applied? Is Iterator[Tuple[int, int]] correct? Or is there a more precise way to do it? Somehow it should be possible to determine the elements' structure Tuple[int, int] from the typehint.

Answer 1

The easy way to verify whether a given type annotation matches the implementation/usage of a function is to add the type annotation to the code and run a type checker on it.

from typing import Iterator

def get_nums_and_square_nums_with_last_digit_9() -> Iterator[tuple[int, int]]:
    sqr = lambda n: n * n
    is_last_digit_9 = lambda n: n % 10 == 9
    is_last_digit_of_square_num_9 = lambda t: (
        lambda n, n_squared: is_last_digit_9(n_squared)
    )(*t)

    nums_and_square_nums = map(lambda n: (n, sqr(n)), range(10))
    return filter(is_last_digit_of_square_num_9, nums_and_square_nums)

Running mypy on the above reports no errors, so it should be fine. We can verify the returned type with reveal_type:

reveal_type(next(get_nums_and_square_nums_with_last_digit_9()))  # tuple[int, int]

If we used an incorrect annotation, e.g. using list when a filter is returned:

def get_nums_and_square_nums_with_last_digit_9() -> list[tuple[int, int]]:

then mypy will report:

error: Incompatible return value type (got "filter[Tuple[int, Any]]", expected "List[Tuple[int, int]]")