zip для асинхронных итераторов в python

Question

У меня есть 2 итератора

async def iter1():
   for i in range(5):
       await asyncio.sleep(0.1)
       yield 1

async def iter2():
   for i in range(10):
       await asyncio.sleep(1)
       yield 2

Я хочу сделать функцию async_zip_stream(*args, default_value=0)

Она должна итерировать все переданные в нее итераторы одновременно, не синхронизируясь, выдавая таплы значений. Если какой то итератор в какой то момент не сгенерировал значение, нужно вернуть default value

ТО есть в данном случае ответом было бы

1, 0
1, 0
1, 0 
1, 0
1, 0 
0, 2
....

Не очень понимаю как такое можно сделать без синхронизации между итераторами

Answer 1

async def async_zip_stream(*iterators, default_value=0):
    tasks = [asyncio.create_task(iterator.__anext__()) for iterator in iterators]
    done = [False] * len(iterators)
    
    while not all(done):
        results = []
        
        for i, task in enumerate(tasks):
            if done[i]:
                results.append(default_value)
            elif task.done():
                try:
                    results.append(task.result())
                    tasks[i] = asyncio.create_task(iterators[i].__anext__())
                except StopAsyncIteration:
                    done[i] = True
                    results.append(default_value)
            else:
                results.append(default_value)
        
        yield tuple(results)
        await asyncio.sleep(0.01)  # Slight delay to allow other tasks to progress

БЛОГ НА HUSL

zip для асинхронных итераторов в python

Ответы (1 шт):