Add type hints to function that accepts any iterable and returns zipped values

Question

I have the following working code to solve Advent of Code 2021 day 1 (requires Python 3.10 due to itertools.pairwise), but I am stuck on how to correctly add type hints.

from itertools import pairwise, tee
from typing import Iterator, TypeVar, Tuple, Generator

_T_co = TypeVar("_T_co", covariant=True)

def tripletwise(iterable: Iterator[_T_co]) -> zip[Tuple[_T_co, _T_co, _T_co]]:
    """tripletwise('ABCDEFG') --> ABC BCD CDE DEF EFG"""
    a, b, c = tee(iterable, 3)
    next(b, None)
    next(c, None)
    next(c, None)
    return zip(a, b, c)


def sonar_sweep(sea_floor_depths: Generator[int, None, None], part: int = 1) -> int:
    if part == 1:
        return sum(b > a for a, b in pairwise(sea_floor_depths))
    return sonar_sweep((a + b + c for a, b, c in tripletwise(sea_floor_depths)), part=1)


def sea_floor_depths(report: str) -> Generator[int, None, None]:
    return map(int, report.splitlines())


def main(part: int = 1) -> int:
    with open("2021/data/day01.txt") as f:
        report = f.read()
    return sonar_sweep(sea_floor_depths(report), part)


if __name__ == "__main__":
    report = """199
200
208
210
200
207
240
269
260
263"""
    assert sonar_sweep(sea_floor_depths(report), part=1) == 7
    assert sonar_sweep(sea_floor_depths(report), part=2) == 5
    assert list(tripletwise("ABCDEFG")) == [
        ("A", "B", "C"),
        ("B", "C", "D"),
        ("C", "D", "E"),
        ("D", "E", "F"),
        ("E", "F", "G"),
    ]
    assert list(tripletwise(sea_floor_depths(report))) == [
        (199, 200, 208),
        (200, 208, 210),
        (208, 210, 200),
        (210, 200, 207),
        (200, 207, 240),
        (207, 240, 269),
        (240, 269, 260),
        (269, 260, 263),
    ]

    print(main())
    print(main(part=2))

Currently, when I run mypy on this, I get:

2021/day01.py:22: error: Incompatible return value type (got "map[int]", expected "Generator[int, None, None]")
2021/day01.py:44: error: Argument 1 to "tripletwise" has incompatible type "str"; expected "Iterator[<nothing>]"

I want to indicate with the type hints that tripletwise() can accept any iterable (for example both the string "ABCDEFG" or the output of sea_floor_depths(report)). I'm also unsure about the correct return type of sea_floor_depths().

Answer 1

As for the first error, the map builtin is only documented to return an iterator, so the most specific type hint that you can use is Iterator[int]. Generator is a subtype of Iterator that only arises when using generator expressions or generator functions, which your function does not.

For the second error, the issue is that strs are not themselves iterators, but rather only iterables, so the appropriate type hint would be Iterable[_T_co].