async def test_concatmap(assert_run, event_loop):
# Concurrent run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.concatmap(lambda x: stream.range(x, x+2, interval=5))
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [1, 1, 3, 5, 5]
# Sequential run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.concatmap(
lambda x: stream.range(x, x+2, interval=5),
task_limit=1)
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [5, 1, 5, 1, 5]
# Limited run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.concatmap(
lambda x: stream.range(x, x+2, interval=5),
task_limit=2)
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [1, 4, 1, 4, 5]
# Make sure item arrive as soon as possible
with event_loop.assert_cleanup():
xs = stream.just(2)
ys = xs | pipe.concatmap(lambda x: stream.range(x, x+4, interval=1))
zs = ys | pipe.timeout(2) # Sould NOT raise
await assert_run(zs, [2, 3, 4, 5])
assert event_loop.steps == [1, 1, 1]
async def test_flatmap(assert_run, event_loop):
# Concurrent run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.flatmap(lambda x: stream.range(x, x+2, interval=5))
await assert_run(ys, [0, 2, 4, 1, 3, 5])
assert event_loop.steps == [1, 1, 3, 1, 1]
# Sequential run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.flatmap(
lambda x: stream.range(x, x+2, interval=5),
task_limit=1)
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [5, 1, 5, 1, 5]
# Limited run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.flatmap(
lambda x: stream.range(x, x+2, interval=5),
task_limit=2)
await assert_run(ys, [0, 2, 1, 3, 4, 5])
assert event_loop.steps == [1, 4, 1, 5]
async def test_slice(assert_run, event_loop):
slice = stream.select.slice
with event_loop.assert_cleanup():
xs = (stream.range(10, 20)
| add_resource.pipe(1)
| slice.pipe(2))
await assert_run(xs, [10, 11])
with event_loop.assert_cleanup():
xs = (stream.range(10, 20)
| add_resource.pipe(1)
| slice.pipe(8, None))
await assert_run(xs, [18, 19])
with event_loop.assert_cleanup():
xs = (stream.range(10, 20)
| add_resource.pipe(1)
| slice.pipe(-3, -1))
await assert_run(xs, [17, 18])
with event_loop.assert_cleanup():
xs = (stream.range(10, 20)
| add_resource.pipe(1)
| slice.pipe(-5, -1, 2))
await assert_run(xs, [15, 17])
with pytest.raises(ValueError):
xs = stream.range(10, 20) | slice.pipe(5, 1, -1)
with pytest.raises(ValueError):
xs = stream.range(10, 20) | slice.pipe(-8, 8)
async def test_list(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | pipe.list()
await assert_run(xs, [[0, 1, 2, 3, 4]])
with event_loop.assert_cleanup():
xs = stream.range(0) | add_resource.pipe(1) | pipe.list()
await assert_run(xs, [[]])
async def test_skiplast(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(10) | add_resource.pipe(1) | pipe.skiplast(8)
await assert_run(xs, [0, 1])
with event_loop.assert_cleanup():
xs = stream.range(10) | add_resource.pipe(1) | pipe.skiplast(0)
await assert_run(xs, list(range(10)))
async def test_take(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.count() | add_resource.pipe(1) | pipe.take(3)
await assert_run(xs, [0, 1, 2])
with event_loop.assert_cleanup():
xs = stream.count() | add_resource.pipe(1) | pipe.take(0)
await assert_run(xs, [])
async def test_chain(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) + stream.range(5, 10)
await assert_run(xs, list(range(10)))
with event_loop.assert_cleanup():
xs = stream.range(10, 15) | add_resource.pipe(1)
xs += stream.range(15, 20) | add_resource.pipe(1)
await assert_run(xs, list(range(10, 20)))
async def test_timeout(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(3) | pipe.timeout(5)
await assert_run(xs, [0, 1, 2])
assert event_loop.steps == []
with event_loop.assert_cleanup():
xs = stream.range(3) + stream.never()
ys = xs | pipe.timeout(1)
await assert_run(ys, [0, 1, 2], asyncio.TimeoutError())
assert event_loop.steps == [1]
async def test_merge(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
ys = stream.range(0, 5, 2, interval=2) | pipe.merge(xs)
await assert_run(ys, [0, 1, 2, 3, 4])
assert event_loop.steps == [1, 1, 1, 1]
with event_loop.assert_cleanup():
xs = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
ys = stream.range(0, 5, 2, interval=2) | pipe.merge(xs)
await assert_run(ys[:3], [0, 1, 2])
assert event_loop.steps == [1, 1]
async def test_reduce(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | pipe.reduce(min)
await assert_run(xs, [0])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | pipe.reduce(max)
await assert_run(xs, [4])
with event_loop.assert_cleanup():
xs = stream.range(0) | add_resource.pipe(1) | pipe.reduce(max)
await assert_run(xs, [], IndexError("Index out of range"))
async def test_ziplatest(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(0, 5, 2, interval=2)
ys = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
zs = stream.ziplatest(xs, ys, default='▲')
await assert_run(zs, [(0, '▲'), (0, 1), (2, 1), (2, 3), (4, 3)])
assert event_loop.steps == [1, 1, 1, 1]
with event_loop.assert_cleanup():
xs = stream.range(0, 5, 2, interval=2)
ys = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
zs = stream.ziplatest(xs, ys, partial=False)
await assert_run(zs, [(0, 1), (2, 1), (2, 3), (4, 3)])
assert event_loop.steps == [1, 1, 1, 1]
async def test_switchmap(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(0, 30, 10, interval=3)
ys = xs | pipe.switchmap(lambda x: stream.range(x, x+5, interval=1))
await assert_run(ys, [0, 1, 2, 10, 11, 12, 20, 21, 22, 23, 24])
assert event_loop.steps == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
# Test cleanup procedure
with event_loop.assert_cleanup():
xs = stream.range(0, 5, interval=1)
ys = xs | pipe.switchmap(lambda x: stream.range(x, x+2, interval=2))
await assert_run(ys[:3], [0, 1, 2])
assert event_loop.steps == [1, 1]
async def test_action(assert_run, event_loop):
with event_loop.assert_cleanup():
lst = []
xs = stream.range(3) | add_resource.pipe(1) | pipe.action(lst.append)
await assert_run(xs, [0, 1, 2])
assert lst == [0, 1, 2]
with event_loop.assert_cleanup():
queue = asyncio.Queue()
xs = stream.range(3) | add_resource.pipe(1) | pipe.action(queue.put)
await assert_run(xs, [0, 1, 2])
assert queue.get_nowait() == 0
assert queue.get_nowait() == 1
assert queue.get_nowait() == 2
async def test_print(assert_run, event_loop):
with event_loop.assert_cleanup():
f = io.StringIO()
xs = stream.range(3) | add_resource.pipe(1) | pipe.print(file=f)
await assert_run(xs, [0, 1, 2])
assert f.getvalue() == '0\n1\n2\n'
with event_loop.assert_cleanup():
f = io.StringIO()
xs = (
stream.range(3)
| add_resource.pipe(1)
| pipe.print('{:.1f}', end='|', file=f)
)
await assert_run(xs, [0, 1, 2])
assert f.getvalue() == '0.0|1.0|2.0|'
async def test_dropwhile(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = (stream.range(1, 10)
| add_resource.pipe(1)
| pipe.dropwhile(lambda x: x < 7))
await assert_run(xs, [7, 8, 9])
async def afunc(x):
await asyncio.sleep(1)
return x < 7
with event_loop.assert_cleanup():
xs = (stream.range(1, 10)
| add_resource.pipe(1)
| pipe.dropwhile(afunc))
await assert_run(xs, [7, 8, 9])
assert event_loop.steps == [1] * 8
async def test_streamcontext(event_loop):
with event_loop.assert_cleanup():
xs = stream.range(3) | add_resource.pipe(1)
async with streamcontext(xs) as streamer:
it = iter(range(3))
async for item in streamer:
assert item == next(it)
assert event_loop.steps == [1]
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
async with xs.stream() as streamer:
it = iter(range(5))
async for item in streamer:
assert item == next(it)
assert event_loop.steps == [1]
async def test_filter(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = (stream.range(1, 10)
| add_resource.pipe(1)
| pipe.filter(lambda x: x in [4, 7, 8]))
await assert_run(xs, [4, 7, 8])
async def afunc(x):
await asyncio.sleep(1)
return x in [3, 6, 9]
with event_loop.assert_cleanup():
xs = (stream.range(1, 10)
| add_resource.pipe(1)
| pipe.filter(afunc))
await assert_run(xs, [3, 6, 9])
assert event_loop.steps == [1]*10
async def test_dropwhile(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = (stream.range(1, 10)
| add_resource.pipe(1)
| pipe.dropwhile(lambda x: x < 7))
await assert_run(xs, [7, 8, 9])
async def afunc(x):
await asyncio.sleep(1)
return x < 7
with event_loop.assert_cleanup():
xs = (stream.range(1, 10)
| add_resource.pipe(1)
| pipe.dropwhile(afunc))
await assert_run(xs, [7, 8, 9])
assert event_loop.steps == [1]*8
async def test_streamcontext(event_loop):
with event_loop.assert_cleanup():
xs = stream.range(3) | add_resource.pipe(1)
async with streamcontext(xs) as streamer:
it = iter(range(3))
async for item in streamer:
assert item == next(it)
assert event_loop.steps == [1]
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
async with xs.stream() as streamer:
it = iter(range(5))
async for item in streamer:
assert item == next(it)
assert event_loop.steps == [1]
async def test_concatmap(assert_run, event_loop):
# Concurrent run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.concatmap(lambda x: stream.range(x, x + 2, interval=5))
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [1, 1, 3, 5, 5]
# Sequential run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.concatmap(
lambda x: stream.range(x, x + 2, interval=5), task_limit=1
)
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [5, 1, 5, 1, 5]
# Limited run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.concatmap(
lambda x: stream.range(x, x + 2, interval=5), task_limit=2
)
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [1, 4, 1, 4, 5]
# Make sure item arrive as soon as possible
with event_loop.assert_cleanup():
xs = stream.just(2)
ys = xs | pipe.concatmap(lambda x: stream.range(x, x + 4, interval=1))
zs = ys | pipe.timeout(2) # Sould NOT raise
await assert_run(zs, [2, 3, 4, 5])
assert event_loop.steps == [1, 1, 1]
# An exception might get discarded if the result can be produced before the
# processing of the exception is required
with event_loop.assert_cleanup():
xs = stream.iterate([True, False])
ys = xs | pipe.concatmap(
lambda x: stream.range(0, 3, interval=1)
if x
else stream.throw(ZeroDivisionError)
)
zs = ys | pipe.take(3)
await assert_run(zs, [0, 1, 2])
assert event_loop.steps == [1, 1]
async def test_chunks(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(3, interval=1) | pipe.chunks(3)
await assert_run(xs, [[0, 1, 2]])
with event_loop.assert_cleanup():
xs = stream.range(4, interval=1) | pipe.chunks(3)
await assert_run(xs, [[0, 1, 2], [3]])
with event_loop.assert_cleanup():
xs = stream.range(5, interval=1) | pipe.chunks(3)
await assert_run(xs, [[0, 1, 2], [3, 4]])
with event_loop.assert_cleanup():
xs = (stream.count(interval=1)
| add_resource.pipe(1)
| pipe.chunks(3))
await assert_run(xs[:1], [[0, 1, 2]])
async def test_merge(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
ys = stream.range(0, 5, 2, interval=2) | pipe.merge(xs)
await assert_run(ys, [0, 1, 2, 3, 4])
assert event_loop.steps == [1, 1, 1, 1]
with event_loop.assert_cleanup():
xs = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
ys = stream.range(0, 5, 2, interval=2) | pipe.merge(xs)
await assert_run(ys[:3], [0, 1, 2])
assert event_loop.steps == [1, 1]
with event_loop.assert_cleanup():
xs = stream.just(1) + stream.never()
ys = xs | pipe.merge(xs) | pipe.timeout(1)
await assert_run(ys, [1, 1], asyncio.TimeoutError())
assert event_loop.steps == [1]
async def test_cycle(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.empty() | pipe.cycle() | pipe.timeout(1)
await assert_run(xs, [], asyncio.TimeoutError())
with event_loop.assert_cleanup():
xs = (
stream.empty()
| add_resource.pipe(1)
| pipe.cycle()
| pipe.timeout(1)
)
await assert_run(xs, [], asyncio.TimeoutError())
with event_loop.assert_cleanup():
xs = stream.just(1) | add_resource.pipe(1) | pipe.cycle()
await assert_run(xs[:5], [1]*5)
assert event_loop.steps == [1]*5
async def test_chunks(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(3, interval=1) | pipe.chunks(3)
await assert_run(xs, [[0, 1, 2]])
with event_loop.assert_cleanup():
xs = stream.range(4, interval=1) | pipe.chunks(3)
await assert_run(xs, [[0, 1, 2], [3]])
with event_loop.assert_cleanup():
xs = stream.range(5, interval=1) | pipe.chunks(3)
await assert_run(xs, [[0, 1, 2], [3, 4]])
with event_loop.assert_cleanup():
xs = (stream.count(interval=1)
| add_resource.pipe(1)
| pipe.chunks(3))
await assert_run(xs[:1], [[0, 1, 2]])
async def test_cycle(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.empty() | pipe.cycle() | pipe.timeout(1)
await assert_run(xs, [], asyncio.TimeoutError())
with event_loop.assert_cleanup():
xs = (
stream.empty()
| add_resource.pipe(1)
| pipe.cycle()
| pipe.timeout(1)
)
await assert_run(xs, [], asyncio.TimeoutError())
with event_loop.assert_cleanup():
xs = stream.just(1) | add_resource.pipe(1) | pipe.cycle()
await assert_run(xs[:5], [1]*5)
assert event_loop.steps == [1]*5
async def test_item(assert_run, event_loop):
item = stream.select.item
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(2)
await assert_run(xs, [2])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(-2)
await assert_run(xs, [3])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(10)
exception = IndexError('Index out of range', )
await assert_run(xs, [], exception)
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(-10)
exception = IndexError('Index out of range', )
await assert_run(xs, [], exception)
async def test_item(assert_run, event_loop):
item = stream.select.item
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(2)
await assert_run(xs, [2])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(-2)
await assert_run(xs, [3])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(10)
exception = IndexError('Index out of range',)
await assert_run(xs, [], exception)
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | item.pipe(-10)
exception = IndexError('Index out of range',)
await assert_run(xs, [], exception)
async def test_starmap(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5)
ys = stream.range(5)
zs = xs | pipe.zip(ys) | pipe.starmap(lambda x, y: x + y)
expected = [x * 2 for x in range(5)]
await assert_run(zs, expected)
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = stream.range(1, 4)
zs = xs | pipe.zip(ys) | pipe.starmap(asyncio.sleep)
await assert_run(zs, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = stream.range(1, 4)
zs = xs | pipe.zip(ys) | pipe.starmap(asyncio.sleep, task_limit=1)
await assert_run(zs, [1, 2, 3])
assert event_loop.steps == [1, 2, 3]
async def test_starmap(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5)
ys = stream.range(5)
zs = xs | pipe.zip(ys) | pipe.starmap(lambda x, y: x+y)
expected = [x*2 for x in range(5)]
await assert_run(zs, expected)
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = stream.range(1, 4)
zs = xs | pipe.zip(ys) | pipe.starmap(asyncio.sleep)
await assert_run(zs, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = stream.range(1, 4)
zs = xs | pipe.zip(ys) | pipe.starmap(asyncio.sleep, task_limit=1)
await assert_run(zs, [1, 2, 3])
assert event_loop.steps == [1, 2, 3]
async def test_getitem(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | pipe.getitem(2)
await assert_run(xs, [2])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
await assert_run(xs[2], [2])
with event_loop.assert_cleanup():
s = slice(1, 3)
xs = stream.range(5) | add_resource.pipe(1) | pipe.getitem(s)
await assert_run(xs, [1, 2])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
await assert_run(xs[1:3], [1, 2])
with event_loop.assert_cleanup():
s = slice(1, 5, 2)
xs = stream.range(5) | add_resource.pipe(1) | pipe.getitem(s)
await assert_run(xs, [1, 3])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
await assert_run(xs[1:5:2], [1, 3])
with pytest.raises(TypeError):
xs = stream.range(5)[None]
async def test_getitem(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | pipe.getitem(2)
await assert_run(xs, [2])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
await assert_run(xs[2], [2])
with event_loop.assert_cleanup():
s = slice(1, 3)
xs = stream.range(5) | add_resource.pipe(1) | pipe.getitem(s)
await assert_run(xs, [1, 2])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
await assert_run(xs[1:3], [1, 2])
with event_loop.assert_cleanup():
s = slice(1, 5, 2)
xs = stream.range(5) | add_resource.pipe(1) | pipe.getitem(s)
await assert_run(xs, [1, 3])
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1)
await assert_run(xs[1:5:2], [1, 3])
with pytest.raises(TypeError):
xs = stream.range(5)[None]
async def test_aitercontext(event_loop):
async def agen():
for x in range(5):
await asyncio.sleep(1)
yield x
with event_loop.assert_cleanup():
async with aitercontext(agen()) as safe_gen:
it = iter(range(5))
async for item in safe_gen:
assert item == next(it)
assert event_loop.steps == [1] * 5
with pytest.raises(RuntimeError):
await anext(safe_gen)
with pytest.raises(RuntimeError):
async with safe_gen:
pass
with pytest.raises(RuntimeError):
await safe_gen.athrow(ValueError())
with event_loop.assert_cleanup():
async with aitercontext(agen()) as safe_gen:
assert await safe_gen.__anext__() == 0
with pytest.raises(ZeroDivisionError):
await safe_gen.athrow(ZeroDivisionError())
async for item in safe_gen:
assert False # No more items
safe_gen = aitercontext(agen())
with pytest.warns(UserWarning):
await anext(safe_gen)
with pytest.raises(TypeError):
AsyncIteratorContext(None)
with pytest.raises(TypeError):
AsyncIteratorContext(safe_gen)
async def test_aggregate(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | pipe.accumulate()
await assert_run(xs, [0, 1, 3, 6, 10])
with event_loop.assert_cleanup():
xs = (stream.range(2, 4)
| add_resource.pipe(1)
| pipe.accumulate(func=operator.mul, initializer=2))
await assert_run(xs, [2, 4, 12])
with event_loop.assert_cleanup():
xs = stream.range(0) | add_resource.pipe(1) | pipe.accumulate()
await assert_run(xs, [])
async def sleepmax(x, y):
return await asyncio.sleep(1, result=max(x, y))
with event_loop.assert_cleanup():
xs = stream.range(3) | add_resource.pipe(1) | pipe.accumulate(sleepmax)
await assert_run(xs, [0, 1, 2])
assert event_loop.steps == [1] * 3
async def test_aggregate(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) | add_resource.pipe(1) | pipe.accumulate()
await assert_run(xs, [0, 1, 3, 6, 10])
with event_loop.assert_cleanup():
xs = (stream.range(2, 4)
| add_resource.pipe(1)
| pipe.accumulate(func=operator.mul, initializer=2))
await assert_run(xs, [2, 4, 12])
with event_loop.assert_cleanup():
xs = stream.range(0) | add_resource.pipe(1) | pipe.accumulate()
await assert_run(xs, [])
async def sleepmax(x, y):
return await asyncio.sleep(1, result=max(x, y))
with event_loop.assert_cleanup():
xs = stream.range(3) | add_resource.pipe(1) | pipe.accumulate(sleepmax)
await assert_run(xs, [0, 1, 2])
assert event_loop.steps == [1]*3
async def test_flatmap(assert_run, event_loop):
# Concurrent run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.flatmap(lambda x: stream.range(x, x + 2, interval=5))
await assert_run(ys, [0, 2, 4, 1, 3, 5])
assert event_loop.steps == [1, 1, 3, 1, 1]
# Sequential run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.flatmap(lambda x: stream.range(x, x + 2, interval=5),
task_limit=1)
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [5, 1, 5, 1, 5]
# Limited run
with event_loop.assert_cleanup():
xs = stream.range(0, 6, 2, interval=1)
ys = xs | pipe.flatmap(lambda x: stream.range(x, x + 2, interval=5),
task_limit=2)
await assert_run(ys, [0, 2, 1, 3, 4, 5])
assert event_loop.steps == [1, 4, 1, 5]
async def test_map(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(5) | pipe.map(lambda x: x**2)
expected = [x**2 for x in range(5)]
await assert_run(xs, expected)
with event_loop.assert_cleanup():
xs = stream.range(5)
ys = xs | pipe.map(lambda x, y: x + y, xs)
expected = [x * 2 for x in range(5)]
await assert_run(ys, expected)
with event_loop.assert_cleanup():
xs = stream.range(1, 4) | pipe.map(asyncio.sleep)
expected = [None] * 3
await assert_run(xs, expected)
assert event_loop.steps == [1, 2, 3]
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = xs | pipe.map(asyncio.sleep, xs)
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 2, 3]
event_loop.steps.clear()
async def test_merge(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
ys = stream.range(0, 5, 2, interval=2) | pipe.merge(xs)
await assert_run(ys, [0, 1, 2, 3, 4])
assert event_loop.steps == [1, 1, 1, 1]
with event_loop.assert_cleanup():
xs = stream.range(1, 5, 2, interval=2) | pipe.delay(1)
ys = stream.range(0, 5, 2, interval=2) | pipe.merge(xs)
await assert_run(ys[:3], [0, 1, 2])
assert event_loop.steps == [1, 1]
with event_loop.assert_cleanup():
xs = stream.just(1) + stream.never()
ys = xs | pipe.merge(xs) | pipe.timeout(1)
await assert_run(ys, [1, 1], asyncio.TimeoutError())
assert event_loop.steps == [1]
# Reproduce issue #65
with event_loop.assert_cleanup():
xs = stream.iterate([1, 2])
ys = stream.iterate([3, 4])
zs = stream.merge(xs, ys) | pipe.take(3)
await assert_run(zs, [1, 2, 3])
with event_loop.assert_cleanup():
xs = stream.iterate([1, 2, 3])
ys = stream.throw(ZeroDivisionError)
zs = stream.merge(xs, ys) | pipe.delay(1) | pipe.take(3)
await assert_run(zs, [1, 2, 3])
# Silencing of a CancelledError
async def agen1():
if False:
yield
try:
await asyncio.sleep(2)
except asyncio.CancelledError:
return
async def agen2():
yield 1
with event_loop.assert_cleanup():
xs = stream.merge(agen1(), agen2()) | pipe.delay(1) | pipe.take(1)
await assert_run(xs, [1])
async def test_map(assert_run, event_loop):
# Synchronous/simple
with event_loop.assert_cleanup():
xs = stream.range(5) | pipe.map(lambda x: x**2)
expected = [x**2 for x in range(5)]
await assert_run(xs, expected)
# Synchronous/multiple
with event_loop.assert_cleanup():
xs = stream.range(5)
ys = xs | pipe.map(lambda x, y: x + y, xs)
expected = [x * 2 for x in range(5)]
await assert_run(ys, expected)
# Asynchronous/simple/concurrent
with event_loop.assert_cleanup():
xs = stream.range(1, 4) | pipe.map(asyncio.sleep)
expected = [None] * 3
await assert_run(xs, expected)
assert event_loop.steps == [1, 1, 1]
# Asynchronous/simple/sequential
with event_loop.assert_cleanup():
xs = stream.range(1, 4) | pipe.map(asyncio.sleep, task_limit=1)
expected = [None] * 3
await assert_run(xs, expected)
assert event_loop.steps == [1, 2, 3]
# Asynchronous/multiple/concurrent
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = xs | pipe.map(asyncio.sleep, xs)
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
# Asynchronous/multiple/sequential
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = xs | pipe.map(asyncio.sleep, xs, task_limit=1)
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 2, 3]
# As completed
with event_loop.assert_cleanup():
xs = stream.iterate([2, 4, 1, 3, 5])
ys = xs | pipe.map(asyncio.sleep, xs, ordered=False)
await assert_run(ys, [1, 2, 3, 4, 5])
assert event_loop.steps == [1, 1, 1, 1, 1]
# Invalid argument
with pytest.raises(ValueError):
await (stream.range(1, 4) | pipe.map(asyncio.sleep, task_limit=0))
# Break
with event_loop.assert_cleanup():
xs = stream.count(1)
ys = xs | pipe.map(asyncio.sleep, xs, task_limit=10)
await assert_run(ys[:3], [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
# Stuck
with event_loop.assert_cleanup():
xs = stream.count(1)
ys = xs | pipe.map(asyncio.sleep, xs, task_limit=1) | pipe.timeout(5)
await assert_run(ys, [1, 2, 3, 4], asyncio.TimeoutError())
# Force await
with event_loop.assert_cleanup():
xs = stream.iterate([1, 2, 3])
ys = xs | pipe.map(async_(lambda x: asyncio.sleep(x, x)))
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
# Map await_
with event_loop.assert_cleanup():
xs = stream.iterate(map(lambda x: asyncio.sleep(x, x), [1, 2, 3]))
ys = xs | pipe.map(await_)
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
async def test_flatmap(assert_run, event_loop):
with event_loop.assert_cleanup():
xs = stream.range(0, 3, interval=1)
ys = xs | pipe.flatmap(lambda x: stream.range(x, x + 6, 3, interval=3))
await assert_run(ys, [0, 1, 2, 3, 4, 5])
assert event_loop.steps == [1, 1, 1, 1, 1]
async def test_exitstack(event_loop):
@async_context_manager
async def acontext(arg):
await asyncio.sleep(1)
try:
yield arg
finally:
await asyncio.sleep(1)
with event_loop.assert_cleanup():
async with AsyncExitStack() as stack:
for i in range(5):
value = await stack.enter_context(acontext(i))
assert value == i
assert event_loop.steps == [1] * 10
with event_loop.assert_cleanup():
async with AsyncExitStack() as stack:
for i in range(5):
value = await stack.enter_context(acontext(i))
assert value == i
await stack.aclose()
assert event_loop.steps == [1] * 10
with event_loop.assert_cleanup():
async with AsyncExitStack() as stack:
for i in range(5):
value = await stack.enter_context(acontext(i))
assert value == i
new_stack = stack.pop_all()
assert event_loop.steps == [1] * 5
async with new_stack:
pass
assert event_loop.steps == [1] * 10
with event_loop.assert_cleanup():
async with AsyncExitStack() as stack:
@stack.callback
async def cleanup():
await asyncio.sleep(1)
for i in range(5):
value = await stack.enter_context(acontext(i))
assert value == i
assert event_loop.steps == [1] * 11
with event_loop.assert_cleanup():
async with AsyncExitStack() as stack:
for i in range(5):
context = acontext(i)
assert await context.__aenter__() == i
stack.push(context)
assert event_loop.steps == [1] * 10
with event_loop.assert_cleanup():
with pytest.raises(KeyError):
async with AsyncExitStack() as stack:
for i in range(5):
value = await stack.enter_context(acontext(i))
assert value == i
raise KeyError('test')
assert event_loop.steps == [1] * 10
async def test_map(assert_run, event_loop):
# Synchronous/simple
with event_loop.assert_cleanup():
xs = stream.range(5) | pipe.map(lambda x: x**2)
expected = [x**2 for x in range(5)]
await assert_run(xs, expected)
# Synchronous/multiple
with event_loop.assert_cleanup():
xs = stream.range(5)
ys = xs | pipe.map(lambda x, y: x+y, xs)
expected = [x*2 for x in range(5)]
await assert_run(ys, expected)
# Asynchronous/simple/concurrent
with event_loop.assert_cleanup():
xs = stream.range(1, 4) | pipe.map(asyncio.sleep)
expected = [None] * 3
await assert_run(xs, expected)
assert event_loop.steps == [1, 1, 1]
# Asynchronous/simple/sequential
with event_loop.assert_cleanup():
xs = stream.range(1, 4) | pipe.map(asyncio.sleep, task_limit=1)
expected = [None] * 3
await assert_run(xs, expected)
assert event_loop.steps == [1, 2, 3]
# Asynchronous/multiple/concurrent
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = xs | pipe.map(asyncio.sleep, xs)
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
# Asynchronous/multiple/sequential
with event_loop.assert_cleanup():
xs = stream.range(1, 4)
ys = xs | pipe.map(asyncio.sleep, xs, task_limit=1)
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 2, 3]
# As completed
with event_loop.assert_cleanup():
xs = stream.iterate([2, 4, 1, 3, 5])
ys = xs | pipe.map(asyncio.sleep, xs, ordered=False)
await assert_run(ys, [1, 2, 3, 4, 5])
assert event_loop.steps == [1, 1, 1, 1, 1]
# Invalid argument
with pytest.raises(ValueError):
await (stream.range(1, 4) | pipe.map(asyncio.sleep, task_limit=0))
# Break
with event_loop.assert_cleanup():
xs = stream.count(1)
ys = xs | pipe.map(asyncio.sleep, xs, task_limit=10)
await assert_run(ys[:3], [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
# Stuck
with event_loop.assert_cleanup():
xs = stream.count(1)
ys = xs | pipe.map(asyncio.sleep, xs, task_limit=1) | pipe.timeout(5)
await assert_run(ys, [1, 2, 3, 4], asyncio.TimeoutError())
# Force await
with event_loop.assert_cleanup():
xs = stream.iterate([1, 2, 3])
ys = xs | pipe.map(async_(lambda x: asyncio.sleep(x, x)))
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
# Map await_
with event_loop.assert_cleanup():
xs = stream.iterate(map(lambda x: asyncio.sleep(x, x), [1, 2, 3]))
ys = xs | pipe.map(await_)
await assert_run(ys, [1, 2, 3])
assert event_loop.steps == [1, 1, 1]
