def test_task_map_with_all_inputs_unmapped(x, y, out):
@prefect.task
def add(x, y):
return x + y
with Flow(name="test") as f:
res = add.map(unmapped(x), unmapped(y))
flow_state = f.run()
assert flow_state.is_successful()
assert flow_state.result[res].is_successful()
assert flow_state.result[res].result == out
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_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_unmapped_applied(self):
with Flow("test") as flow:
z = add.map(x=[1, 2, 3], y=edges.unmapped(100))
state = flow.run()
assert state.result[z].result == [101, 102, 103]
def test_tasks_have_all_non_unmapped_constant_args_as_transitive_upstream_deps(
self,
):
def func(a, b, c, d):
m = inc.copy(name="m").bind(1)
n = inc.copy(name="n").bind(a)
o = inc.copy(name="o").bind(b)
p = add.copy(name="p").bind(n, o)
q = add.copy(name="q").bind(c, d)
r = add.copy(name="r").bind(q, m)
return m, n, o, p, q, r
with Flow("test") as flow:
a = ranged.copy(name="a").bind(3)
b = inc.copy(name="b").bind(1)
c = Constant(1, name="c")
d = Constant(range(3), name="d")
m, n, o, p, q, r = apply_map(
func, a, edges.unmapped(b), c=edges.unmapped(c), d=d
)
def edge_info(task):
"""Returns a map of {upstream: (is_data_dep, is_mapped)}"""
return {
e.upstream_task: (e.key is not None, e.mapped)
for e in flow.edges_to(task)
}
assert edge_info(m) == {a: (False, True), b: (False, False), d: (False, True)}
assert edge_info(n) == {a: (True, True), b: (False, False), d: (False, True)}
assert edge_info(o) == {a: (False, True), b: (True, False), d: (False, True)}
assert edge_info(p) == {n: (True, True), o: (True, True)}
assert edge_info(q) == {a: (False, True), b: (False, False), d: (True, True)}
assert edge_info(r) == {q: (True, True), m: (True, True)}
state = flow.run()
res = {t: state.result[t].result for t in [m, n, o, p, q, r]}
sol = {
m: [2, 2, 2],
n: [1, 2, 3],
o: [3, 3, 3],
p: [4, 5, 6],
q: [1, 2, 3],
r: [3, 4, 5],
}
assert res == sol
def test_map_can_handle_nonkeyed_nonmapped_upstreams_and_mapped_args(executor):
ii = IdTask()
ll = ListTask()
with Flow(name="test") as f:
res = ll.map(start=ll(), upstream_tasks=[unmapped(ii(5))])
s = f.run(executor=executor)
m = s.result[res]
assert s.is_successful()
assert isinstance(m.map_states, list)
assert len(m.result) == 3
assert m.result == [[1 + i, 2 + i, 3 + i] for i in range(3)]
def test_map_can_handle_fixed_kwargs(executor):
ll = ListTask()
a = AddTask()
with Flow(name="test") as f:
res = a.map(ll, y=unmapped(5))
s = f.run(executor=executor)
m = s.result[res]
assert s.is_successful()
assert isinstance(m.map_states, list)
assert len(m.result) == 3
assert m.result == [6, 7, 8]
def test_unmapped_on_mapped(executor):
@prefect.task
def add_one(x):
if isinstance(x, list):
return x + x
return x + 1
with Flow("wild") as flow:
res = add_one.map(unmapped(add_one.map([1, 2, 3])))
flow_state = flow.run(executor=executor)
assert flow_state.is_successful()
assert flow_state.result[res].result == [2, 3, 4, 2, 3, 4]
def main(component_spec_file, src_dir, dest_dir, title):
"""
Runs a parameterized pipeline that processes all the SSP's in the
SRC_DIR and places output files in the DEST_DIR. The COMPONENT_SPEC_FILE
is supplied to the *match* phase.
"""
with Flow("SSP Pipeline") as flow:
src = Parameter("src")
dest = Parameter("dest")
component_spec = Parameter("component_spec")
ssps = discover_ssps(src)
components = match_ssp.map(ssps, unmapped(component_spec))
combined = combine(components)
oscal = oscalize(combined, title)
write_file(oscal, Path(dest_dir) / "oscal.json")
markdown(combined, dest)
flow.run(component_spec=component_spec_file, src=src_dir, dest=dest_dir)
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_map_tracks_non_mapped_upstream_tasks(executor):
div = DivTask()
@task
def zeros():
return [0, 0, 0]
@task(trigger=prefect.triggers.all_failed)
def register(x):
return True
with Flow(name="test") as f:
res = register.map(div.map(zeros()), upstream_tasks=[unmapped(div(1))])
s = f.run(executor=executor)
assert s.is_failed()
assert all(sub.is_failed() for sub in s.result[res].map_states)
assert all([
isinstance(sub, prefect.engine.state.TriggerFailed)
for sub in s.result[res].map_states
])
def test_apply_map_inside_resource_manager_works(self):
@resource_manager
class MyResource:
def setup(self):
return 1
def cleanup(self, _):
pass
def func(x, a):
return add(x, a), add(x, 2)
with Flow("test") as flow:
with MyResource() as a:
b, c = apply_map(func, range(4), edges.unmapped(a))
d = add.map(b, c)
state = flow.run()
assert state.result[a].result == 1
assert state.result[b].result == [1, 2, 3, 4]
assert state.result[c].result == [2, 3, 4, 5]
assert state.result[d].result == [3, 5, 7, 9]
def test_apply_map_inside_case_statement_works(self):
def func(x, a):
return add(x, 1), add(x, a)
with Flow("test") as flow:
branch = Parameter("branch")
with case(branch, True):
a = inc(1)
b, c = apply_map(func, range(4), edges.unmapped(a))
d = add.map(b, c)
state = flow.run(branch=True)
assert state.result[a].result == 2
assert state.result[b].result == [1, 2, 3, 4]
assert state.result[c].result == [2, 3, 4, 5]
assert state.result[d].result == [3, 5, 7, 9]
state = flow.run(branch=False)
assert state.result[a].is_skipped()
assert state.result[b].is_skipped()
assert state.result[c].is_skipped()
assert state.result[d].is_skipped()
def test_unmapped_annotation_takes_precedence(self):
e = Edge(edges.unmapped(Task()), Task(), mapped=True)
assert e.mapped is False
def test_unmapped(self):
e = Edge(edges.unmapped(Task()), Task())
assert e.mapped is False
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_multiple_annotations(self):
ea = edges.mapped(edges.flatten(edges.unmapped(edges.mapped(Task()))))
assert ea.annotations == {"flattened": True, "mapped": True}
Behavior is configured in the config.toml
"""
with Flow('Loading Socrata data to Postgres',
result=LocalResult(dir=os.path.join(os.path.dirname(__file__),
"results"),
validate_dir=True),
state_handlers=[postgres.log_to_database]) as flow:
datasets = Parameter("datasets")
# get last updated from database
since = postgres.get_last_updated()
# download dataset from Socrata
downloads = socrata.download_dataset.map(dataset=datasets,
since=unmapped(since))
# get the temp tables ready for load
prep = postgres.prep_load()
# load each downloaded file
load = postgres.load_datafile.map(datafile=downloads)
# commit new data to database and clean up
complete = postgres.complete_load()
# make sure prep runs before load
flow.add_edge(upstream_task=prep, downstream_task=load)
# make sure load runs before complete
flow.add_edge(upstream_task=load, downstream_task=complete)
if __name__ == "__main__":
logger = prefect.context.get("logger")
def test_unmapped(self):
ea = edges.unmapped(Task())
assert ea.annotations == {"mapped": False}
with Flow('sample') as flow:
extract_results = fake_extract.map(periods=[5, 6, 8], )
upstream_repos = Parameter('upstream_repos')
with semantic_operation(
flow=flow,
upstream_repos=upstream_repos,
versions_to_retain=Parameter('versions_to_retain'),
) as op:
# load_done = df_to_table_task(fake_extract(1), upstream_tasks=[base_ref])
workspaces = op.workspaces
foo1 = workspaces['unittest']
foo2 = workspaces['foo']
load_done = df_to_table_task.map(request=extract_results,
repo_uri=unmapped(foo1['repo_uri']))
load_done2 = df_to_table_task.map(request=extract_results,
repo_uri=unmapped(foo2['repo_uri']))
say_something('after_load', upstream_tasks=[load_done, load_done2])
say_something('after with', upstream_tasks=[op.push])
class RepoTasksTest(unittest.TestCase):
repo: Repository
def setUp(self):
self.repo = build_repo()
def tearDown(self):
self.repo.delete(unregister=True, uncheckout=True)