def get_flow(self, flow_location: str = None) -> "Flow":
"""
Given a flow_location within this Storage object, returns the underlying Flow (if possible).
Args:
- flow_location (str, optional): the location of a flow within this Storage; in this case,
a file path or python path where a Flow has been serialized to. Will use `path`
if not provided.
Returns:
- Flow: the requested flow
Raises:
- ValueError: if the flow is not contained in this storage
"""
if flow_location:
if flow_location not in self.flows.values():
raise ValueError("Flow is not contained in this Storage")
elif self.path:
flow_location = self.path
else:
raise ValueError("No flow location provided")
# check if the path given is a file path
if os.path.isfile(flow_location):
if self.stored_as_script:
return extract_flow_from_file(file_path=flow_location)
else:
return prefect.core.flow.Flow.load(flow_location)
# otherwise the path is given in the module format
else:
return extract_flow_from_module(module_str=flow_location)
def get_flow(self, flow_name: str) -> "Flow":
"""
Given a flow name within this Storage object, load and return the Flow.
Args:
- flow_name (str): the name of the flow to return.
Returns:
- Flow: the requested flow
"""
if flow_name not in self.flows:
raise ValueError("Flow is not contained in this Storage")
flow_location = self.flows[flow_name]
# check if the path given is a file path
if os.path.isfile(flow_location):
if self.stored_as_script:
return extract_flow_from_file(
file_path=flow_location, flow_name=flow_name
)
else:
with open(flow_location, "rb") as f:
return flow_from_bytes_pickle(f.read())
# otherwise the path is given in the module format
else:
return extract_flow_from_module(
module_str=flow_location, flow_name=flow_name
)
def test_extract_flow_from_module_callable_objects(mymodule):
flow1 = Flow("flow 1")
flow2 = Flow("flow 2")
class Obj:
def build_flow(self):
return flow2
mymodule.build_flow = lambda: flow1
mymodule.multi_level = Obj()
mymodule.bad_type = lambda: 1
assert extract_flow_from_module("mymodule:build_flow") is flow1
assert extract_flow_from_module("mymodule:build_flow", "flow 1") is flow1
assert extract_flow_from_module("mymodule:multi_level.build_flow") is flow2
with pytest.raises(TypeError, match="Object at 'mymodule:bad_type'"):
extract_flow_from_module("mymodule:bad_type")
def get_flow(self, flow_name: str) -> "Flow":
"""
Given a flow name within this Storage object, load and return the Flow.
Args:
- flow_name (str): the name of the flow to return.
Returns:
- Flow: the requested flow
"""
if flow_name not in self.flows:
raise ValueError("Flow is not contained in this Storage")
return extract_flow_from_module(module_str=self.module,
flow_name=flow_name)
def test_extract_flow_from_module(mymodule):
class Obj:
flow = Flow("multi-level flow")
mymodule.flow = Flow("top level flow")
mymodule.multi_level = Obj()
mymodule.bad_type = 1
# module with single top-level flow has flow auto-inferred
assert extract_flow_from_module("mymodule") is mymodule.flow
# Specifying name/attribute still works
assert extract_flow_from_module("mymodule",
"top level flow") is mymodule.flow
assert extract_flow_from_module("mymodule:flow",
"top level flow") is mymodule.flow
# Multi-level attrs work
assert extract_flow_from_module("mymodule:multi_level.flow") is Obj.flow
# Multiple top-level flows
mymodule.flow2 = Flow("a second flow")
assert extract_flow_from_module("mymodule",
"top level flow") is mymodule.flow
assert extract_flow_from_module("mymodule",
"a second flow") is mymodule.flow2
# Multiple flows not auto-inferred
with pytest.raises(ValueError, match="Multiple flows found"):
extract_flow_from_module("mymodule")
# Name not found
with pytest.raises(ValueError, match="Failed to find flow"):
extract_flow_from_module("mymodule", "unknown name")
# Name doesn't match specified object
with pytest.raises(ValueError, match="Flow at 'mymodule:flow' is named"):
extract_flow_from_module("mymodule:flow", "incorrect name")
# Not a flow object
with pytest.raises(TypeError, match="Object at 'mymodule:bad_type'"):
extract_flow_from_module("mymodule:bad_type")
def test_extract_flow_from_module():
module_name = "tests.utilities.test_storage"
multi_level_flow = Flow("test-module-loading-multilevel")
factory_flow = Flow("test-module-loading-callable")
test_extract_flow_from_module.multi_level_flow = multi_level_flow
test_extract_flow_from_module.not_a_flow = None
test_extract_flow_from_module.not_a_flow_factory = lambda: object()
test_extract_flow_from_module.invalid_callable = lambda _a, _b, **_kwargs: None
class FlowFactory:
@classmethod
def default_flow(cls):
return factory_flow
test_extract_flow_from_module.callable_flow = FlowFactory.default_flow
default_flow = extract_flow_from_module(module_name)
attribute_flow = extract_flow_from_module(module_name, "flow")
module_flow = extract_flow_from_module(f"{module_name}:flow")
multi_level_default_flow = extract_flow_from_module(
f"{module_name}:test_extract_flow_from_module.multi_level_flow"
)
multi_level_arg_flow = extract_flow_from_module(
module_name, "test_extract_flow_from_module.multi_level_flow"
)
callable_flow = extract_flow_from_module(
f"{module_name}:test_extract_flow_from_module.callable_flow"
)
assert flow == default_flow == attribute_flow == module_flow
assert multi_level_flow == multi_level_default_flow == multi_level_arg_flow
assert factory_flow == callable_flow
with pytest.raises(AttributeError):
extract_flow_from_module("tests.utilities.test_storage:should_not_exist_flow")
with pytest.raises(AttributeError):
extract_flow_from_module(
"tests.utilities.test_storage", "should_not_exist_flow"
)
with pytest.raises(ValueError, match="without an attribute specifier or remove"):
extract_flow_from_module("tests.utilities.test_storage:flow", "flow")
with pytest.raises(ValueError, match="must return `prefect.Flow`"):
extract_flow_from_module(
f"{module_name}:test_extract_flow_from_module.not_a_flow"
)
with pytest.raises(ValueError, match="must return `prefect.Flow`"):
extract_flow_from_module(
f"{module_name}:test_extract_flow_from_module.not_a_flow_factory"
)
with pytest.raises(TypeError):
extract_flow_from_module(
f"{module_name}:test_extract_flow_from_module.invalid_callable"
)
