How to Compress a Generator with a Boolean Mask in Python
Filters items from a generator based on a parallel boolean mask, yielding only the items where the mask is True.
Python code
11 linesdef compress(generator, mask):
for item, keep in zip(generator, mask):
if keep:
yield item
if __name__ == "__main__":
data = [1, 2, 3, 4, 5]
mask = [True, False, True, False, True]
result = list(compress(iter(data), mask))
print(result)Output
[1, 3, 5]How it works
The zip function pairs each item from the generator with the corresponding boolean from the mask, stopping when either is exhausted. The if keep condition only yields items where the mask value is truthy. This lazy generator avoids building an intermediate list of all items, making it memory-efficient for large or infinite sequences. Using iter(data) is redundant here since lists are already iterable, but it clarifies that the function accepts any iterable.
Common mistakes
- Forgetting to convert the generator result to a list before printing it.
- Assuming the mask and generator must be the same length; zip stops at the shorter one.
- Using a truthy check on non-boolean values, which may filter unintentionally.
- Not realizing that the function returns a generator, so it must be consumed.
Variations
- Use `itertools.compress(data, mask)` from the standard library for a built-in equivalent.
- Use a list comprehension: `[item for item, keep in zip(data, mask) if keep]`.
Real-world use cases
- Filtering sensor readings where a companion boolean array indicates valid measurements.
- Selecting certain log entries based on a pre-computed mask of relevance flags.
- Applying a subset of items from a large stream of API responses matching a user's filter criteria.
Sponsored
More from Comprehensions & generators
- Batch Rows in Chunks with a Generator in Python easy
- Build a Generator Pipeline in Python: Filter Then Map medium
- Build a lazy generator to read file lines in Python easy
- Chunk an Iterable into Batches with a Generator in Python easy
- Convert Data in Python with Comprehensions and Generators easy
- Count Data in Python with Comprehensions and Generators easy
Keep learning
Related tutorials and quizzes for this topic.