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