31 lines
713 B
Python
31 lines
713 B
Python
import asyncio
|
|
|
|
|
|
class WaitGroup:
|
|
"""
|
|
asyncio wait group implementation (similar to sync.WaitGroup in go)
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._counter = 0
|
|
self._zero_event = asyncio.Event()
|
|
self._zero_event.set()
|
|
|
|
def add(self, delta: int = 1) -> None:
|
|
new_value = self._counter + delta
|
|
if new_value < 0:
|
|
raise ValueError("WaitGroup counter cannot go negative.")
|
|
|
|
self._counter = new_value
|
|
|
|
if self._counter == 0:
|
|
self._zero_event.set()
|
|
else:
|
|
self._zero_event.clear()
|
|
|
|
def done(self) -> None:
|
|
self.add(-1)
|
|
|
|
async def wait(self) -> None:
|
|
await self._zero_event.wait()
|