freqtrade_origin/freqtrade/rpc/api_server/ws/message_stream.py

33 lines
899 B
Python
Raw Normal View History

2022-11-15 03:27:45 +00:00
import asyncio
import time
2022-11-15 03:27:45 +00:00
class MessageStream:
"""
A message stream for consumers to subscribe to,
and for producers to publish to.
"""
2024-05-12 14:51:11 +00:00
2022-11-15 03:27:45 +00:00
def __init__(self):
self._loop = asyncio.get_running_loop()
self._waiter = self._loop.create_future()
def publish(self, message):
"""
Publish a message to this MessageStream
:param message: The message to publish
"""
2022-11-15 03:27:45 +00:00
waiter, self._waiter = self._waiter, self._loop.create_future()
waiter.set_result((message, time.time(), self._waiter))
2022-11-15 03:27:45 +00:00
async def __aiter__(self):
"""
Iterate over the messages in the message stream
"""
2022-11-15 03:27:45 +00:00
waiter = self._waiter
while True:
# Shield the future from being cancelled by a task waiting on it
message, ts, waiter = await asyncio.shield(waiter)
yield message, ts