def test_as_task_toggles_constants(self):
with Flow("test") as f:
t = tasks.as_task(4)
s = tasks.as_task(5, convert_constants=False)
assert isinstance(t, Task)
assert t.name == "4"
assert not isinstance(s, Task)
assert s == 5
def test_list_of_tasks():
"""
Test that a list of tasks can be set as a switch condition
"""
with Flow(name="test") as flow:
condition = Condition()
true_branch = [SuccessTask(), SuccessTask()]
list_task = as_task(true_branch)
false_branch = SuccessTask()
ifelse(condition, list_task, false_branch)
with prefect.context(CONDITION=True):
state = flow.run()
for t in true_branch:
assert isinstance(state.result[t], Success)
assert isinstance(state.result[false_branch], Skipped)
with prefect.context(CONDITION=False):
state = flow.run()
for t in true_branch:
# the tasks in the list ran becuase they have no upstream dependencies.
assert isinstance(state.result[t], Success)
# but the list itself skipped
assert isinstance(state.result[list_task], Skipped)
assert isinstance(state.result[false_branch], Success)
def test_ordered_collections(self):
"""
Tests that ordered collections maintain order
"""
val = [[list(range(100))]]
with Flow("test") as f:
task = tasks.as_task(val)
f.add_task(task)
assert f.run().result[task].result == val
def replace(self, old: Task, new: Task, validate: bool = True) -> None:
"""
Performs an inplace replacement of the old task with the provided new task.
Args:
- old (Task): the old task to replace
- new (Task): the new task to replace the old with; if not a Prefect
Task, Prefect will attempt to convert it to one
- validate (boolean, optional): whether to validate the Flow after
the replace has been completed; defaults to `True`
Raises:
- ValueError: if the `old` task is not a part of this flow
"""
if old not in self.tasks:
raise ValueError("Task {t} was not found in Flow {f}".format(
t=old, f=self))
new = as_task(new)
# update tasks
self.tasks.remove(old)
self.add_task(new)
self._cache.clear()
affected_edges = {e for e in self.edges if old in e.tasks}
# remove old edges
for edge in affected_edges:
self.edges.remove(edge)
# replace with new edges
for edge in affected_edges:
upstream = new if edge.upstream_task == old else edge.upstream_task
downstream = new if edge.downstream_task == old else edge.downstream_task
self.add_edge(
upstream_task=upstream,
downstream_task=downstream,
key=edge.key,
mapped=edge.mapped,
validate=False,
)
# update auxiliary task collections
ref_tasks = self.reference_tasks()
new_refs = [t for t in ref_tasks if t != old
] + ([new] if old in ref_tasks else [])
self.set_reference_tasks(new_refs)
if validate:
self.validate()
def test_apply_map_simple(self, api):
if api == "functional":
def func(x, y, z):
a = add(x, y)
a.name = "add-a"
b = add(a, z)
b.name = "add-b"
c = add(b, 1)
c.name = "add-c"
return inc(c)
with Flow("test") as flow:
y = ranged(3)
z = edges.unmapped(1)
res = apply_map(func, range(3, 6), y=y, z=z)
else:
def func(x, y, z, flow):
a = add.copy(name="add-a").bind(x, y, flow=flow)
b = add.copy(name="add-b").bind(a, z, flow=flow)
c = add.copy(name="add-c").bind(b, 1, flow=flow)
return inc.copy().bind(c, flow=flow)
flow = Flow("test")
y = ranged.copy().bind(3, flow=flow)
z = edges.unmapped(tasks.as_task(1, flow=flow))
res = apply_map(func, range(3, 6), y=y, z=z, flow=flow)
consts = {t.name: c for t, c in flow.constants.items()}
assert consts == {
"ranged": {
"n": 3
},
"add-b": {
"y": 1
},
"add-c": {
"y": 1
}
}
for task in flow.tasks:
if task.name != "ranged":
for e in flow.edges_to(task):
assert e.mapped
state = flow.run()
assert state.result[res].result == [6, 8, 10]
def test_as_task_with_basic_python_objs(self, obj):
@tasks.task
def return_val(x):
"Necessary because constant tasks aren't tracked inside the flow"
return x
with Flow("test") as f:
t = tasks.as_task(obj)
val = return_val(t)
assert isinstance(t, Task)
res = FlowRunner(f).run(return_tasks=[val])
assert res.is_successful()
assert res.result[val].result == obj
def test_nested_collections(self, val):
with Flow("test") as f:
task = tasks.as_task(val)
f.add_task(task)
assert f.run().result[task].result == val
def test_nested_collections_of_mixed_constants_are_not_constants(
self, val):
with Flow("test"):
task = tasks.as_task(val)
assert not isinstance(task, Constant)
def test_nested_collections_of_constants_are_constants(self, val):
task = tasks.as_task(val)
assert isinstance(task, Constant)
assert task.value == val
def test_as_task_doesnt_label_tasks_as_auto_generated(self):
t = Task()
assert t.auto_generated is False
assert tasks.as_task(t).auto_generated is False
def test_as_task_toggles_constants(self):
with Flow("test"):
t = tasks.as_task(4)
assert isinstance(t, Task)
assert t.name == "4"
def set_dependencies(
self,
task: object,
upstream_tasks: Iterable[object] = None,
downstream_tasks: Iterable[object] = None,
keyword_tasks: Mapping[str, object] = None,
mapped: bool = False,
validate: bool = None,
) -> None:
"""
Convenience function for adding task dependencies.
Args:
- task (object): a Task that will become part of the Flow. If the task is not a
Task subclass, Prefect will attempt to convert it to one.
- upstream_tasks ([object], optional): Tasks that will run before the task runs. If any task
is not a Task subclass, Prefect will attempt to convert it to one.
- downstream_tasks ([object], optional): Tasks that will run after the task runs. If any task
is not a Task subclass, Prefect will attempt to convert it to one.
- keyword_tasks ({key: object}, optional): The results of these tasks
will be provided to the task under the specified keyword
arguments. If any task is not a Task subclass, Prefect will attempt to
convert it to one.
- mapped (bool, optional): Whether the upstream tasks (both keyed
and non-keyed) should be mapped over; defaults to `False`. If `True`, any
tasks wrapped in the `prefect.utilities.tasks.unmapped` container will
_not_ be mapped over.
- validate (bool, optional): Whether or not to check the validity of
the flow (e.g., presence of cycles). Defaults to the value of `eager_edge_validation`
in your Prefect configuration file.
Returns:
- None
"""
task = as_task(task)
assert isinstance(task, Task) # mypy assert
# add the main task (in case it was called with no arguments)
self.add_task(task)
# add upstream tasks
for t in upstream_tasks or []:
is_mapped = mapped & (not isinstance(t, unmapped))
t = as_task(t)
assert isinstance(t, Task) # mypy assert
self.add_edge(
upstream_task=t,
downstream_task=task,
validate=validate,
mapped=is_mapped,
)
# add downstream tasks
for t in downstream_tasks or []:
t = as_task(t)
assert isinstance(t, Task) # mypy assert
self.add_edge(upstream_task=task,
downstream_task=t,
validate=validate)
# add data edges to upstream tasks
for key, t in (keyword_tasks or {}).items():
is_mapped = mapped & (not isinstance(t, unmapped))
t = as_task(t)
assert isinstance(t, Task) # mypy assert
self.add_edge(
upstream_task=t,
downstream_task=task,
key=key,
validate=validate,
mapped=is_mapped,
)
def test_as_task_unpacks_unmapped_objects(self):
t1 = Task()
unmapped_t1 = tasks.unmapped(t1)
assert tasks.as_task(t1) is t1
assert tasks.as_task(t1).auto_generated is False
def test_as_task_unpacks_unmapped_objects(self, edge_annotation):
t1 = Task()
annotated = edge_annotation(t1)
assert tasks.as_task(annotated) is t1
assert tasks.as_task(annotated).auto_generated is False
def test_as_task_unpacks_unmapped_objects(self):
t1 = Task()
unmapped_t1 = tasks.unmapped(t1)
assert tasks.as_task(t1) is t1