Files
langgraph/libs/langgraph/langgraph/utils/aio.py
T
2024-09-17 09:31:48 -07:00

36 lines
1.2 KiB
Python

import asyncio
import sys
PY_310 = sys.version_info >= (3, 10)
class Queue(asyncio.Queue):
async def wait(self):
"""If queue is empty, wait until an item is available.
Copied from Queue.get(), removing the call to .get_nowait(),
ie. this doesn't consume the item, just waits for it.
"""
while self.empty():
if PY_310:
getter = self._get_loop().create_future()
else:
getter = self._loop.create_future()
self._getters.append(getter)
try:
await getter
except:
getter.cancel() # Just in case getter is not done yet.
try:
# Clean self._getters from canceled getters.
self._getters.remove(getter)
except ValueError:
# The getter could be removed from self._getters by a
# previous put_nowait call.
pass
if not self.empty() and not getter.cancelled():
# We were woken up by put_nowait(), but can't take
# the call. Wake up the next in line.
self._wakeup_next(self._getters)
raise