def test_flatmap_one_unnested_input(self, executor):
a = AddTask()
with Flow("test") as flow:
z = a.map(x=flatten([1]), y=flatten([[5]]))
state = flow.run()
assert state.result[z].is_mapped()
assert state.result[z].result == [6]
def test_flatmap_one_unmappable_input(self, executor):
a = AddTask()
with Flow("test") as flow:
z = a.map(x=flatten(1), y=flatten([[1]]))
state = flow.run()
assert state.result[z].is_failed()
assert (state.result[z].message ==
"At least one upstream state has an unmappable result.")
def test_flatmap_flatmapped_result(self, executor):
# flatmap over a flattened mapped task
ll = ListTask()
nest = NestTask()
a = AddTask()
with Flow(name="test") as f:
nested = nest.map(ll())
nested2 = nest.map(flatten(nested))
x = a.map(flatten(nested2))
s = f.run(executor=executor)
assert s.result[x].result == [2, 3, 4]
def test_flatmap_constant(self, executor):
# flatmap over a constant
add = AddTask()
with Flow(name="test") as f:
a = add.map(flatten([[1, 2, 3]]))
b = add.map(flatten([[1], [2], [3]]))
c = add.map(flatten([[1], [2, 3]]))
d = add.map(flatten([1, 2, 3]))
s = f.run(executor=executor)
# all results should be the same
for task in [a, b, c, d]:
assert s.result[task].result == [2, 3, 4]
def test_multiple_annotations_applied(self):
with Flow("test") as flow:
z = add(x=edges.mapped(edges.flatten([[1], [2, 3]])),
y=edges.unmapped(100))
state = flow.run()
assert state.result[z].result == [101, 102, 103]
def test_flatmap_nested_noniterable_input(self, executor, caplog):
class NestTaskWithNonIterable(Task):
def run(self, x):
if x == 2:
return x
return [x]
ll = ListTask()
nest = NestTaskWithNonIterable()
a = AddTask()
with Flow("test") as flow:
nested = nest.map(ll())
z = a.map(flatten(nested))
state = flow.run(executor=executor)
if not isinstance(executor, prefect.executors.dask.DaskExecutor):
# When multiprocessing with Dask, caplog fails
assert any(
[
"`flatten` was used on upstream task that did not return an iterable"
in record.message
for record in caplog.records
]
), "Warning is logged"
assert state.result[z].result == [2, 3, 4]
def test_flatmap_reduced_result(self, executor):
# flatmap over a reduced flattened mapped task
ll = ListTask()
nest = NestTask()
a = AddTask()
with Flow(name="test") as f:
nested = nest.map(ll())
nested2 = nest(flatten(nested))
x = a.map(flatten(nested2))
from prefect.utilities.debug import raise_on_exception
with raise_on_exception():
s = f.run(executor=executor)
assert s.result[x].result == [2, 3, 4]
def test_apply_map_flatten_works(self, input_type):
def func(a):
return a * 2
@tasks.task
def nested_list():
return [[1], [1, 2], [1, 2, 3]]
constant_nested = nested_list.run()
with Flow("test") as flow:
nested = nested_list()
if input_type == "constant":
a = apply_map(func, edges.flatten(nested))
else:
a = apply_map(func, edges.flatten(constant_nested))
state = flow.run()
assert state.result[a].result == [2, 2, 4, 2, 4, 6]
def test_flatmap_nested_signal_input(self, executor):
class NestTaskWithSignal(Task):
def run(self, x):
if x == 2:
raise SKIP("Skip this task execution")
return [x]
ll = ListTask()
nest = NestTaskWithSignal()
a = AddTask()
with Flow("test") as flow:
nested = nest.map(ll())
z = a.map(flatten(nested))
state = flow.run(executor=executor)
assert state.result[z].result == [2, None, 4]
def test_apply_map_mixed_edge_types(self):
@tasks.task
def get_mixed_types():
return 3, [1, 2, 3], [[1, 2], [3]]
@tasks.task
def identity(a, b, c):
return a, b, c
def func(u, m, f):
return identity(u, m, f)
with Flow("test") as flow:
m = get_mixed_types()
a = apply_map(func, edges.unmapped(m[0]), m[1],
edges.flatten(m[2]))
state = flow.run()
assert state.result[a].result == [(3, 1, 1), (3, 2, 2), (3, 3, 3)]
def test_flatmap_nested_exception_input(self, executor):
class NestTaskWithException(Task):
def run(self, x):
if x == 2:
raise ValueError("This task failed!")
return [x]
ll = ListTask()
nest = NestTaskWithException()
a = AddTask()
with Flow("test") as flow:
nested = nest.map(ll())
z = a.map(flatten(nested))
state = flow.run(executor=executor)
result = state.result[z].result
assert len(result) == 3
assert (result[0], result[2]) == (2, 4)
assert isinstance(result[1], TRIGGERFAIL)
def test_flat_applied_to_constant(self):
with Flow("test") as flow:
z = add.map(x=[1, 2, 3], y=edges.unmapped(edges.flatten(100)))
state = flow.run()
assert state.result[z].result == [101, 102, 103]
def test_flat_applied(self):
with Flow("test") as flow:
z = add.map(x=edges.flatten([[1], [2, 3]]), y=edges.unmapped(100))
state = flow.run()
assert state.result[z].result == [101, 102, 103]
def test_unmapped_annotation_takes_precedence(self):
e = Edge(edges.flatten(Task()), Task(), flattened=False)
assert e.flattened is True
def test_flat(self):
ea = edges.flatten(Task())
assert ea.annotations == {"flattened": True}
def test_nested_annotation(self):
e = Edge(edges.flatten(edges.mapped(Task())), Task())
assert e.flattened is True
assert e.mapped is True
def test_multiple_annotations(self):
ea = edges.mapped(edges.flatten(edges.unmapped(edges.mapped(Task()))))
assert ea.annotations == {"flattened": True, "mapped": True}
def test_flat(self):
e = Edge(edges.flatten(Task()), Task())
assert e.flattened is True
