partition_by_type.py

from typing import Any, Iterable, Tuple

def partition_by_type(iterable: Iterable[Any], *args: Tuple[type], with_rest=False):
    """Partitions iterable by types.
    Creates as many iterators as type tuples have been given and optionally
    returns another iterator for all non-matching instances."""
    iterator = iter(iterable)
    types_map = {
        types: collections.deque()
        for types in args
    }
    if with_rest:
        types_map[None] = collections.deque()

    def types_gen(my_queue):
        while True:
            if not my_queue:
                try:
                    item = next(iterator)
                except StopIteration:
                    return
                is_processed = False
                for types in args:
                    if isinstance(item, types):
                        types_map[types].append(item)
                        is_processed = True
                if not is_processed and with_rest:
                    types_map[None].append(item)
            if my_queue:
                yield my_queue.popleft()
    return tuple(types_gen(queue) for queue in types_map.values())