def _read_sensor( # pylint: disable=too-many-arguments
self, source_uuid: UUID, sid: int, unit: str, topic: str,
callback_config: AdvancedCallbackConfiguration
) -> AsyncGenerator[DataEvent, None]:
monitor_stream = (
stream.repeat(self.device, interval=1)
| pipe.map(
async_(lambda sensor: sensor.get_callback_configuration(sid)))
| pipe.map(lambda current_config: None
if current_config == callback_config else self.device)
| pipe.filter(lambda sensor: sensor is not None)
| pipe.action(lambda sensor: logging.getLogger(__name__).info(
"Resetting callback config for %s", sensor))
| pipe.action(
async_(lambda sensor: sensor.set_callback_configuration(
sid, *callback_config)))
| pipe.filter(lambda x: False))
return stream.merge(
stream.just(monitor_stream),
stream.iterate(self.device.read_events(sids=(sid, )))
| pipe.map(lambda item: DataEvent(sender=source_uuid,
topic=topic,
value=item.payload,
sid=item.sid,
unit=str(unit))),
)
async def euclidean_norm_handler(reader, writer):
# Define lambdas
strip = lambda x: x.decode().strip()
nonempty = lambda x: x != ''
square = lambda x: x ** 2
write_cursor = lambda x: writer.write(b'> ')
square_root = lambda x: x ** 0.5
# Create awaitable
handle_request = (
stream.iterate(reader)
| pipe.print('string: {}')
| pipe.map(strip)
| pipe.takewhile(nonempty)
| pipe.map(float)
| pipe.map(square)
| pipe.print('square: {:.2f}')
| pipe.action(write_cursor)
| pipe.accumulate(initializer=0)
| pipe.map(square_root)
| pipe.print('norm -> {:.2f}')
)
# Loop over norm computations
while not reader.at_eof():
writer.write(INSTRUCTIONS.encode())
try:
result = await handle_request
except ValueError:
writer.write(ERROR.encode())
else:
writer.write(RESULT.format(result).encode())
async def euclidean_norm_handler(reader, writer):
# Define lambdas
strip = lambda x: x.decode().strip()
nonempty = lambda x: x != ''
square = lambda x: x**2
write_cursor = lambda x: writer.write(b'> ')
square_root = lambda x: x**0.5
# Create awaitable
handle_request = (stream.iterate(reader)
| pipe.print('string: {}')
| pipe.map(strip)
| pipe.takewhile(nonempty)
| pipe.map(float)
| pipe.map(square)
| pipe.print('square: {:.2f}')
| pipe.action(write_cursor)
| pipe.accumulate(initializer=0)
| pipe.map(square_root)
| pipe.print('norm -> {:.2f}'))
# Loop over norm computations
while not reader.at_eof():
writer.write(INSTRUCTIONS.encode())
try:
result = await handle_request
except ValueError:
writer.write(ERROR.encode())
else:
writer.write(RESULT.format(result).encode())
async def run(self) -> None:
"""
The main task, that reads data from the sensors and pushes it onto the event_bus.
"""
# Generate the UUIDs of new sensors
sensor_stream = stream.chain(
stream.iterate(
iterate_safely(f"{self.__topic}/get",
f"{self.__topic}/status_update")),
stream.iterate(event_bus.subscribe(f"{self.__topic}/add_host")),
) | pipe.flatmap(
lambda item: stream.chain(
(stream.call(event_bus.call, f"{self.__topic}/get_config", item
) | catch.pipe(TopicNotRegisteredError)),
stream.iterate(
event_bus.subscribe(f"nodes/by_uuid/{item}/update")),
)
| pipe.until(lambda config: config is None)
| pipe.map(lambda config: config if self._is_config_valid(
self.__node_id, config) else None)
| pipe.map(self._create_transport)
| pipe.switchmap(lambda transport: stream.empty() if transport is
None else stream.iterate(transport.stream_data()))
| pipe.action(lambda data: event_bus.publish("wamp/publish", data)
))
await sensor_stream
async def run(
subscription_kinds: List[SubscriptionKind],
topic: str,
begin_timestamp=None,
interval_sec=None,
num_msgs=10,
on_message: Optional[Callable[[Any], Any]] = None,
):
def log_message(msg):
if not isinstance(msg, list):
#LOGGER.info("Status: %s", json.dumps(msg))
pass
else:
LOGGER.debug("Received message: %s", json.dumps(msg))
def maybe_callback(m):
if on_message:
on_message(m)
# the continuous version takes messages up until a time point
if interval_sec:
async with KrakenClient() as client:
for kind in subscription_kinds:
await subscribe(client, kind, PAIRS, SUBSCRIPTION_DEPTH)
await (client.stream(reconnect_count=WEBSOCKET_RECONNECT_COUNT)
| pipe.action(log_message)
| pipe.filter(lambda msg: isinstance(msg, list))
| pipe.map(json.dumps)
#| pipe.print()
| pipe.takewhile(
lambda x: time.time() - begin_timestamp < interval_sec)
| pipe.action(maybe_callback))
client.disconnect()
# the discrete version takes 100 distinct messages
else:
async with KrakenClient() as client:
for kind in subscription_kinds:
await subscribe(client, kind, PAIRS, SUBSCRIPTION_DEPTH)
await (client.stream(reconnect_count=WEBSOCKET_RECONNECT_COUNT)
| pipe.action(log_message)
| pipe.filter(lambda msg: isinstance(msg, list))
| pipe.map(json.dumps)
#| pipe.print()
| pipe.take(num_msgs)
| pipe.action(maybe_callback))
print("disconnecting")
client.disconnect()
def stream_data(self) -> AsyncGenerator[DataEvent, None]:
"""
Generate the initial configuration of the sensor, configure it, and finally stream the data from the sensor.
If there is a configuration update, reconfigure the sensor and start streaming again.
Returns
-------
AsyncGenerator of DataEvent
The data from the device
"""
# Generates the first configuration
# Query the database and if it does not have a config for the sensor, wait until there is one
data_stream = (
stream.chain(
stream.call(
call_safely, "db_labnode_sensors/get_config", "db_labnode_sensors/status_update", self.__uuid
)
| pipe.takewhile(lambda config: config is not None),
stream.iterate(event_bus.subscribe(f"nodes/by_uuid/{self.__uuid}/update")),
)
| pipe.action(
lambda config: logging.getLogger(__name__).info(
"Got new configuration for: %s",
self._device,
)
)
| pipe.map(self._create_config)
| pipe.switchmap(
lambda config: stream.empty()
if config is None or not config["enabled"]
else (self._configure_and_stream(config))
)
)
return data_stream
async def main():
# Create a counting stream with a 0.2 seconds interval
xs = stream.count(interval=0.2)
# Operators can be piped using '|'
ys = xs | pipe.map(lambda x: x**2)
# Streams can be sliced
zs = ys[1:10:2]
# Use a stream context for proper resource management
async with zs.stream() as streamer:
# Asynchronous iteration
async for z in streamer:
# Print 1, 9, 25, 49 and 81
print('->', z)
# Streams can be awaited and return the last value
print('9² = ', await zs)
# Streams can run several times
print('9² = ', await zs)
# Streams can be concatenated
one_two_three = stream.just(1) + stream.range(2, 4)
# Print [1, 2, 3]
print(await stream.list(one_two_three))
async def _get_hexes_at_resolution(self, res: int) -> Set[H3HexCell]:
idx_map_fn = h3.h3_to_parent if res < self.res else h3.h3_to_children
results = await (stream.iterate(self.cells)
| aiopipe.map(partial(idx_map_fn, res=11))
| aiopipe.reduce(lambda a, b: a | b))
return results
async def main(args):
app_id = get_app_id(args)
if app_id is None:
logger.critical('No APP_ID found')
return False
# Loading modules
logger.info('Loading modules')
modules = ModuleCache(args.modules)
assert modules, 'No modules loaded'
# Loading token
with closing(appstorage.create('vk-dump', 'the-island.ru')) as db:
token = await get_or_create_token(args, app_id, db, modules.scope)
async with TokenSession(token) as session:
api = API(session)
rate = TokenLimiter(args.parallel)
async with aiohttp.ClientSession() as session:
if not os.path.isdir(args.root):
logger.info('Creating root dump directory %s', args.root)
os.makedirs(args.root)
downloader = Downloader(args.root, session, rate)
xs = (
stream.merge(*[m.get_items(api, rate) for m in modules])
| pipe.map(downloader.download) # pylint: disable=E1101
)
await xs
logger.info('Done')
def _stream_data(self, transport):
config_stream = (stream.chain(
stream.call(
call_safely,
f"{self.__database_topic}/get_config",
f"{self.__database_topic}/status_update",
transport.uuid,
),
stream.iterate(
event_bus.subscribe(f"nodes/by_uuid/{transport.uuid}/add")),
)
| pipe.map(lambda config: self._create_device(
transport, config))
| pipe.starmap(lambda config, device: stream.empty(
) if device is None else device.stream_data(config))
| pipe.switch()
| context.pipe(
transport,
on_enter=lambda: logging.getLogger(__name__).info(
"Connected to %s at %s (%s).", transport.name,
transport.uri, transport.label),
on_exit=lambda: logging.getLogger(__name__).info(
"Disconnected from %s at %s (%s).", transport.
name, transport.uri, transport.label),
))
return config_stream
def stream_data(self,
config: dict[str, Any]) -> AsyncGenerator[DataEvent, None]:
"""
Stream the data from the sensor.
Parameters
----------
config: dict
A dictionary containing the sensor configuration.
Returns
-------
AsyncGenerator
The asynchronous stream.
"""
data_stream = (stream.chain(
stream.just(config),
stream.iterate(
event_bus.subscribe(f"nodes/by_uuid/{self.__uuid}/update")))
|
pipe.action(lambda _: logging.getLogger(__name__).info(
"Got new configuration for: %s", self) if config is
not None else logging.getLogger(__name__).
info("Removed configuration for: %s", self))
| pipe.map(self._parse_config)
| pipe.switchmap(lambda conf: stream.empty(
) if conf is None or not conf["enabled"] else
(self._configure_and_stream(conf))))
return data_stream
async def pool_for_each_par_map(loop, pool, f, p, iterator):
zx = stream.iterate(iterator)
zx = zx | pipe.map(lambda x: loop.run_in_executor(pool, p, x))
async with zx.stream() as streamer:
limit = pool._max_workers
unfinished = []
while True:
tasks = []
for i in range(limit):
try:
task = await streamer.__anext__()
except StopAsyncIteration:
limit = 0
else:
tasks.append(task)
tasks = tasks + list(unfinished)
assert len(tasks) <= pool._max_workers
if not tasks:
break
finished, unfinished = await asyncio.wait(
tasks, return_when=asyncio.FIRST_COMPLETED)
for finish in finished:
out = finish.result()
f(out)
limit = len(finished)
async def main():
# Create a counting stream with a 0.2 seconds interval
xs = stream.count(interval=0.2)
# Operators can be piped using '|'
ys = xs | pipe.map(lambda x: x**2)
# Streams can be sliced
zs = ys[1:10:2]
# Use a stream context for proper resource management
async with zs.stream() as streamer:
# Asynchronous iteration
async for z in streamer:
# Print 1, 9, 25, 49 and 81
print("->", z)
# Streams can be awaited and return the last value
print("9² = ", await zs)
# Streams can run several times
print("9² = ", await zs)
# Streams can be concatenated
one_two_three = stream.just(1) + stream.range(2, 4)
# Print [1, 2, 3]
print(await stream.list(one_two_three))
def _configure_and_stream(
self,
config: dict[str, Any] | None) -> AsyncGenerator[DataEvent, None]:
if config is None:
return stream.empty()
try:
# Run all config steps in order (concat) and one at a time (task_limit=1). Drop the output. There is
# nothing to compare them to (filter => false), then read all sensors of the bricklet and process them in
# parallel (flatten).
config_stream = stream.chain(
stream.iterate(config["on_connect"])
| pipe.starmap(lambda func, timeout: stream.just(func()) | pipe
.timeout(timeout))
| pipe.concat(task_limit=1)
| pipe.filter(lambda result: False),
stream.iterate(config["config"].items())
| pipe.starmap(self._parse_callback_configuration)
| pipe.starmap(self._set_callback_configuration)
| pipe.flatten()
|
pipe.map(lambda args: self._read_sensor(config["uuid"], *args))
| pipe.flatten(),
)
return config_stream
except NotConnectedError:
# Do not log it
raise
except Exception:
self._logger.exception("This should not happen.")
raise
async def download_files(urls):
async with aiohttp.ClientSession() as session:
await (stream.repeat(session)
| pipe.zip(stream.iterate(fields.itertuples()))
| pipe.flatmap(generate_urls)
| pipe.filter(ignore_existing_file)
| pipe.take(1000)
| pipe.map(retrieve_file, ordered=False, task_limit=10))
async def _event(self, event, users, **kwargs):
print('assigning event!')
handlers = (stream.iterate(users)
| pipe.filter(lambda u: u.uid in self.handlers)
| pipe.map(lambda u: self.handlers.get(u.uid)))
async with handlers.stream() as handlers:
async for handler in handlers:
print(f'assigning event to: {handler.user.username}#{handler.user.discriminator}')
await handler.events.put((event, kwargs))
async def main():
# This stream computes 11² + 13² in 1.5 second
xs = (
stream.count(interval=0.1) # Count from zero every 0.1 s
| pipe.skip(10) # Skip the first 10 numbers
| pipe.take(5) # Take the following 5
| pipe.filter(lambda x: x % 2) # Keep odd numbers
| pipe.map(lambda x: x ** 2) # Square the results
| pipe.accumulate() # Add the numbers together
)
print('11² + 13² = ', await xs)
def _read_sensor( # pylint: disable=too-many-arguments
self, sid: int, interval: float, unit: str, topic: str, timeout: float
) -> AsyncGenerator[DataEvent, None]:
if self.__uuid is None:
raise SensorNotReady("You must enumerate the sensor before reading.")
return (
stream.repeat(stream.call(self._device.get_by_function_id, sid), interval=interval)
| pipe.concat(task_limit=1)
| pipe.timeout(timeout)
| pipe.map(lambda value: DataEvent(sender=self.__uuid, topic=topic, value=value, sid=sid, unit=str(unit)))
)
async def main():
# This stream computes 11² + 13² in 1.5 second
xs = (
stream.count(interval=0.1) # Count from zero every 0.1 s
| pipe.skip(10) # Skip the first 10 numbers
| pipe.take(5) # Take the following 5
| pipe.filter(lambda x: x % 2) # Keep odd numbers
| pipe.map(lambda x: x**2) # Square the results
| pipe.accumulate() # Add the numbers together
)
print('11² + 13² = ', await xs)
async def run(self):
self.session = aiohttp.ClientSession(auth=self.auth)
self.limiter = FrequencyLock(self.interval)
asyncio.create_task(self.limiter.run())
pl = stream.iterate(self.find_images())
if self.limit:
pl = pl | pipe.take(self.limit)
pl = pl | pipe.map(self.tag, task_limit=self.tasks) | pipe.map(
self.store)
try:
await pl
except aiohttp.ClientResponseError as err:
raise click.ClickException(
"Imagga API failure ({.status}): probably exceeded "
"concurrency limits".format(err))
finally:
await self.limiter.stop()
await self.session.close()
def stream_data(self) -> AsyncGenerator[DataEvent, None]:
"""
Generate the initial configuration of the sensor, configure it, and finally stream the data from the sensor.
If there is a configuration update, reconfigure the sensor and start streaming again.
Returns
-------
AsyncGenerator of DataEvent
The data from the device
"""
# Generates the first configuration
# Query the database and if it does not have a config for the sensor, wait until there is one
data_stream = stream.chain(
stream.just(self),
stream.iterate(
event_bus.subscribe(
f"nodes/tinkerforge/{self.device.uid}/remove"))[:1]
| pipe.map(lambda x: None),
) | pipe.switchmap(
lambda sensor: stream.empty() if sensor is None else
(self._stream_config_updates(sensor)
| pipe.switchmap(lambda config: stream.chain(
stream.just(config),
stream.iterate(
event_bus.subscribe(
f"nodes/by_uuid/{config['uuid']}/remove"))[:1]
| pipe.map(lambda x: None),
))
| pipe.action(lambda config: logging.getLogger(__name__).info(
"Got new configuration for: %s",
sensor.device,
))
| pipe.map(self._create_config)
| pipe.switchmap(lambda config: stream.empty()
if config is None or not config["enabled"] else
(self._configure_and_stream(config)))))
return data_stream
async def render_coroutine(self, states: list, save_config: dict = {}):
"""Draw emulation
Should support following:
- render as images
- render as animations
"""
# map current renderable states into a pipeline buffer
frames = await stream.iterate(states) | aiopipe.map(
self._render_single)
self.plot(frames, self.config)
if save_config:
self.save_plot(save_config)
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()
def _stream_transport(transport: TinkerforgeTransport):
sensor_stream = stream.chain(
stream.call(transport.enumerate) | pipe.filter(lambda x: False),
stream.iterate(transport.read_enumeration())
| pipe.action(lambda enumeration: event_bus.publish(f"nodes/tinkerforge/{enumeration[1].uid}/remove", None))
| pipe.filter(lambda enumeration: enumeration[0] is not EnumerationType.DISCONNECTED)
| pipe.starmap(lambda enumeration_type, sensor: TinkerforgeSensor(sensor))
| pipe.map(lambda sensor: sensor.stream_data())
| pipe.flatten(),
) | context.pipe(
transport,
on_enter=lambda: logging.getLogger(__name__).info(
"Connected to Tinkerforge host at %s (%s).", transport.uri, transport.label
),
on_exit=lambda: logging.getLogger(__name__).info(
"Disconnected from Tinkerforge host at %s (%s).", transport.uri, transport.label
),
)
return sensor_stream
def _stream_data(self, transport):
config_stream = (
with_context(
transport,
on_exit=lambda: logging.getLogger(__name__).info(
"Disconnected from APQ Labnode at %s (%s).", transport.uri, transport.label
),
)
| pipe.action(
lambda _: logging.getLogger(__name__).info(
"Connected to APQ Labnode at %s (%s).", transport.uri, transport.label
)
)
| pipe.map(LabnodeSensor)
| pipe.action(async_(lambda sensor: sensor.enumerate()))
| pipe.switchmap(lambda sensor: sensor.stream_data())
)
return config_stream
def _configure_and_stream(
self, config: dict[str, Any]) -> AsyncGenerator[DataEvent, None]:
if config is None:
return stream.empty()
# Run all config steps in order (concat) and one at a time (task_limit=1). Drop the output. There is nothing to
# compare them to (filter => false), then read the device.
config_stream = stream.chain(
stream.iterate(config["on_connect"])
| pipe.starmap(lambda func, timeout: stream.just(func()) | pipe.
timeout(timeout))
| pipe.concat(task_limit=1)
| pipe.filter(lambda result: False),
self._read_device(config)
| pipe.map(lambda item: DataEvent(sender=config["uuid"],
topic=config["topic"],
value=item,
sid=0,
unit=config["unit"]))
| finally_action.pipe(
stream.call(self._clean_up, config["on_disconnect"])),
) | catch.pipe(TypeError, on_exc=self.on_error)
return config_stream
async def data_stream_async(url,
files,
columns=None,
map_func=None,
reduce_func=None,
initializer=None,
producer_num=2,
data_handler_num=2,
executor_type='process'):
data_type = files[0].split('.')[-1]
columns_mapping = get_parquet_mapping() if data_type == 'parquet' else None
if columns:
c = [k for k, v in columns_mapping.items() if v in columns]
c = c + list(
set(columns).difference(set(list(columns_mapping.values()))))
else:
c = columns
global pbar_handler
pbar_handler = PbarHandler(len(files))
global executor
if executor_type == 'process':
executor = ProcessPoolExecutor(max_workers=data_handler_num)
elif executor_type == 'thread':
executor = ThreadPoolExecutor(max_workers=data_handler_num)
if map_func:
map_task = partial(data_load, func=map_func)
else:
map_task = data_load
pbar_handler.register(mapper, len(files))
if reduce_func:
reduce_task = partial(reducer, func=reduce_func)
else:
reduce_task = partial(reducer, func=concat)
initializer = pd.DataFrame()
file_streams = [
stream.preserve(file_stream(files[i::producer_num], url))
for i in range(producer_num)
]
file_list = []
aws = (stream.merge(*file_streams)
| pipe.map(
async_(lambda x: mapper(x, map_task, c, columns_mapping)),
task_limit=data_handler_num)
| pipe.map(async_(lambda x: file_list.append(x[0]) or x[1]),
task_limit=data_handler_num))
if reduce_func:
pbar_handler.register(reducer, len(files) - 1)
rs = stream.reduce(aws, async_(reduce_task), initializer)
reduced = await stream.takelast(rs, 1)
return reduced
else:
data_list = await asyncio.gather(stream.list(aws))
data_list = data_list[0]
tmp_list = list(zip(file_list, data_list))
tmp_list = sorted(tmp_list, key=lambda pair: files.index(pair[0]))
if map_func:
return tmp_list
else:
return pd.concat(list(map(lambda pair: pair[1], tmp_list)), axis=0)
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_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]
