take items from list

take.py

from asyncio import gather, sleep
from aiostream import stream
from itertools import batched, islice
import random

class Item:
    def __init__(self, payload: str):
        self.payload = payload

    @property
    def is_valid(self):
        return self.payload.endswith("GOOD")

    def __hash__(self):
        return hash(self.payload)


items = [
        Item("no"),
        Item("veryGOOD"),
        Item("nope"),
        Item("soGOOD"),
        Item("nope2"),
        Item("nope3"),
        Item("isGOOD"),
        Item("zzzz"),
        Item("vGOOD"),
        Item("nope4"),
        Item("nope5"),
        Item("lastGOOD"),
        ]

async def evaluate(item: Item) -> bool:
    """
    Simulate talking to a third party
    """
    print(f"begin processing: {item.payload}")
    await sleep(random.random() * 1)
    print(f"done processing: {item.payload}")
    return item.is_valid


async def get_winners(items, criterion, n, worker_count=3):
    """
    in haskell you'd write:
      take n (filter criterion items)
    """
    promise_pairs = [
            (element, criterion(element))
            for element
            in items
            ]
    async def generate_winners():
        # promise_pairs is a list of (item, not-yet-computed-value-containing-its-validity)
        for batch in batched(promise_pairs, worker_count):
            inputs, promises = zip(*batch)
            bool_results = await gather(*promises)
            paired_with_answer = zip(inputs, bool_results)
            for (item, should_yield) in paired_with_answer:
                if should_yield:
                    yield item

    answer = (await stream.list(stream.take(generate_winners(), n)))
    for item, promise in promise_pairs:
        # ugh - bookkeeping - let's keep python happy
        promise.close()
    return answer

def reconcile(all_items: list[Item], ones_who_win: list[Item]):
    good_ones_yielded = set()
    winner_set = set(ones_who_win)
    for item in all_items:
        if item in winner_set:
            yield item
            good_ones_yielded.add(item)
        elif good_ones_yielded >= winner_set:
            yield item


async def main():
    answer = await get_winners(items, evaluate, 3)
    print("OK!")
    for x in reconcile(items, answer):
        print(x.payload)