Is there a better way to zip iterables with mask?

Question

I need to zip iterables with a personal mask function or just mask lists [True, False, ...]

from collections.abc import Callable

def zip_mask(a, b, mask):
    iter_a, iter_b = iter(a), iter(b)
    if isinstance(mask, Callable):
        assert len(a) == (valid := sum(mask(i) for i in b)), f'have mismatch length with `{len(a)}`; `{valid}`'
        for i in b:
            yield next(iter_a) if mask(i) else None, next(iter_b)
    else:
        assert (select := set(mask)) == {1, False}, f'bad input mask selections {select}'
        assert len(b) == len(mask) > len(a), f'bad input shapes'
        assert len(a) == (valid := sum(mask)), f'have mismatch length with `{len(a)}`; `{valid}`'
        for m in mask:
            yield next(iter_a) if m else None, next(iter_b)


# (None, 4) (None, 5) (None, 6) (1, 7) (2, 8) (3, 9)
print(*zip_mask([1, 2, 3], [4, 5, 6, 7, 8, 9], lambda x: x >= 7))

Hopefully there's a better way to do this.

Answer 1

from itertools import compress, repeat

def zip_mask(a, b, mask):
    iter_b = iter(b)
    if callable(mask):
        filtered_b = compress(iter_b, (mask(i) for i in iter_b))
    else:
        filtered_b = compress(iter_b, mask)

    filtered_b_list = list(filtered_b)  # Convert to list for length calculation

    # Ensure both parts are iterables for concatenation
    yield from zip(list(repeat(None, len(b) - len(filtered_b_list))) + a, b)
    
result = list(zip_mask([1, 2, 3], [4, 5, 6, 7, 8, 9], lambda x: x >= 7))
print(' '.join(map(str, result)))