How do I type hint a function that returns a zip object?

Question

I've got a function that takes an arbitrary amount of lists (or any iterables, for that matter) and sorts them as one. The code looks like this:

def sort_as_one(*args):
    return zip(*sorted(zip(*args)))
def main():
    list1 = [3, 1, 2, 4]
    list2 = ["a", "b", "d", "e"]
    list3 = [True, False, True, False]
    result = sort_as_one(list1, list2, list3)
    # <zip object at ...>
    print(result)
    list1, list2, list3 = result
    print(list1, list2, list3)
if __name__ == "__main__":
    main()

How can I accurately type hint the function output?

Answer 1

A zip object is an iterator - it follows the iterator protocol. Idiomatically, you would probably just typing hint it as such. In this case, you want to type hint it as a generic using a type variable:

import typing

T = typing.TypeVar("T")

def sort_as_one(*args: T) -> typing.Iterator[T]:
    return zip(*sorted(zip(*args)))

Note, if you are using variadic arguments, you have to only accept a single type. In this case, the best you can probably do in your case is use Any instead of T. But you should consider using only a function like the above in your code if you want to be able to use it with static type checkers.