75 lines
1.6 KiB
Python
75 lines
1.6 KiB
Python
from typing import AsyncIterator
|
|
|
|
import trio
|
|
import tractor
|
|
|
|
|
|
log = tractor.log.get_logger('multiportal')
|
|
|
|
|
|
async def stream_data(seed: int = 10) -> AsyncIterator[int]:
|
|
'''
|
|
Stream a finite sequence of integers.
|
|
|
|
'''
|
|
log.info('Starting stream task')
|
|
|
|
i: int
|
|
for i in range(seed):
|
|
yield i
|
|
await trio.sleep(0) # trigger scheduler
|
|
|
|
|
|
async def stream_from_portal(
|
|
portal: tractor.Portal,
|
|
consumed: list[int],
|
|
) -> None:
|
|
'''
|
|
Consume one stream and toggle each value in a shared list.
|
|
|
|
'''
|
|
stream: tractor.MsgStream
|
|
async with portal.open_stream_from(stream_data) as stream:
|
|
item: int
|
|
async for item in stream:
|
|
if item in consumed:
|
|
consumed.remove(item)
|
|
else:
|
|
consumed.append(item)
|
|
|
|
|
|
async def main() -> None:
|
|
'''
|
|
Consume two concurrent streams through one portal.
|
|
|
|
'''
|
|
an: tractor.ActorNursery
|
|
async with tractor.open_nursery(loglevel='info') as an:
|
|
|
|
portal: tractor.Portal = await an.start_actor(
|
|
'stream_boi',
|
|
enable_modules=[__name__],
|
|
)
|
|
|
|
consumed: list[int] = []
|
|
|
|
n: trio.Nursery
|
|
async with trio.open_nursery() as n:
|
|
for _ in range(2):
|
|
n.start_soon(
|
|
stream_from_portal,
|
|
portal,
|
|
consumed,
|
|
)
|
|
|
|
# both streaming consumer tasks have completed and so we
|
|
# should have nothing in our list thanks to single
|
|
# threadedness
|
|
assert not consumed
|
|
|
|
await an.cancel()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
trio.run(main)
|