Resolve Job inputs automatically in decorated flows

pyproject.toml

@@ -79,7 +79,7 @@ jobflow = ["py.typed"]
 
 [tool.versioningit.vcs]
 method = "git"
-default-tag = "0.0.1"
+default-tag = "0.3.1"
 
 [tool.flake8]
 max-line-length = 88

src/jobflow/core/flow.py

@@ -13,7 +13,12 @@
 
 import jobflow
 from jobflow.core.reference import find_and_get_references
-from jobflow.utils import ValueEnum, contains_flow_or_job, suid
+from jobflow.utils import (
+    ValueEnum,
+    contains_flow_or_job,
+    replace_job_or_flow_with_output,
+    suid,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
@@ -965,15 +970,14 @@ def __init__(self, fn, *args, **kwargs):
 
         if isinstance(output, (jobflow.Job, jobflow.Flow)):
             warnings.warn(
-                f"@flow decorated function '{name}' contains a Flow or"
-                f"Job as an output. Usually the output should be the output of"
-                f"a Job or another Flow (e.g. job.output). Replacing the"
-                f"output of the @flow with the output of the Flow/Job."
-                f"If this message is unexpected then double check the outputs"
-                f"of your @flow decorated function.",
+                f"@flow decorated function '{name}' contains a Flow or Job as an "
+                "output. Usually the output should be the output of a Job or another "
+                "Flow (e.g. job.output). Replacing the output of the @flow with the "
+                "output of the Flow/Job. If this message is unexpected then double "
+                "check the outputs of your @flow decorated function.",
                 stacklevel=2,
             )
-            output = output.output
+        output = replace_job_or_flow_with_output(output)
 
         super().__init__(name=name, jobs=children_list, output=output)
 

src/jobflow/core/job.py

@@ -13,6 +13,7 @@
 
 from jobflow.core.flow import _current_flow_context
 from jobflow.core.reference import OnMissing, OutputReference
+from jobflow.utils.find import replace_job_or_flow_with_output
 from jobflow.utils.uid import suid
 
 if typing.TYPE_CHECKING:
@@ -213,6 +214,10 @@ def get_job(*args, **kwargs) -> Job:
                         f = met
                         args = args[1:]
 
+            if _current_flow_context.get() is not None:
+                args = replace_job_or_flow_with_output(args)
+                kwargs = replace_job_or_flow_with_output(kwargs)
+
             return Job(
                 function=f, function_args=args, function_kwargs=kwargs, **job_kwargs
             )

src/jobflow/utils/__init__.py

@@ -5,6 +5,7 @@
     contains_flow_or_job,
     find_key,
     find_key_value,
+    replace_job_or_flow_with_output,
     update_in_dictionary,
 )
 from jobflow.utils.log import initialize_logger

src/jobflow/utils/find.py

@@ -11,6 +11,24 @@
     from monty.json import MSONable
 
 
+def replace_job_or_flow_with_output(value):
+    """Replace Jobs and Flows in a value with their outputs."""
+    from jobflow import Flow, Job
+
+    if isinstance(value, (Job, Flow)):
+        return value.output
+    if isinstance(value, list):
+        return [replace_job_or_flow_with_output(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(replace_job_or_flow_with_output(item) for item in value)
+    if isinstance(value, dict):
+        return {
+            replace_job_or_flow_with_output(key): replace_job_or_flow_with_output(item)
+            for key, item in value.items()
+        }
+    return value
+
+
 def find_key(
     d: dict[Hashable, Any] | list[Any],
     key: Hashable | type[MSONable],

tests/core/test_flow_decorator.py

@@ -159,6 +159,67 @@ def my_flow(a, b):
     assert result[flow1.output.uuid][1].output == 7
 
 
+def test_flow_resolves_minimal_job_input_example():
+    """Test the minimal Job-as-input example from issue #884."""
+    from jobflow import flow, job
+    from jobflow.managers.local import run_locally
+
+    @job
+    def mult(a, b):
+        return a * b
+
+    @flow
+    def workflow(a, b):
+        sum_result = add(a, b)
+        return mult(a, sum_result)
+
+    flow1 = workflow(1, 2)
+    result = run_locally(flow1, ensure_success=True)
+
+    assert result[flow1.output.uuid][1].output == 3
+
+
+def test_flow_resolves_job_inputs_to_outputs():
+    """Test top-level container Jobs resolve to outputs."""
+    from jobflow import flow, job
+
+    @job
+    def consume(values):
+        return values
+
+    @flow
+    def my_flow():
+        source = add(1, 2)
+        list_job = consume([source])
+        tuple_job = consume((source,))
+        dict_job = consume({"source": source})
+        return [
+            list_job.output,
+            tuple_job.output,
+            dict_job.output,
+        ]
+
+    flow1 = my_flow()
+    source = flow1.jobs[0]
+
+    assert flow1.jobs[1].function_args == ([source.output],)
+    assert flow1.jobs[2].function_args == ((source.output,),)
+    assert flow1.jobs[3].function_args == ({"source": source.output},)
+
+
+def test_flow_normalizes_top_level_job_outputs():
+    """Test Jobs at the top level of a decorated flow output are normalized."""
+    from jobflow import flow
+
+    @flow
+    def my_flow():
+        return {"job": add(1, 2)}
+
+    flow1 = my_flow()
+
+    assert flow1.output == {"job": flow1.jobs[0].output}
+
+
 def test_flow_returns_list():
     """Test that a flow that returns a list of OutputReferences
     can be created and run."""

tests/core/test_job.py

@@ -83,7 +83,6 @@ def test_job_run(capsys, memory_jobstore, memory_data_jobstore):
     response = test_job.run(memory_jobstore)
     assert capsys.readouterr().out == "I am a job\n"
     assert isinstance(response, Response)
-
     # test run with outputs
     test_job = Job(add, function_args=(1,), function_kwargs={"b": 2})
     response = test_job.run(memory_jobstore)
@@ -1155,7 +1154,7 @@ def make(self):
     test_job = maker.make()
     test_job.update_metadata(
         {"v": 5},
-        callback_filter=lambda job: (job.maker is not None and job.maker.value == 42),
+        callback_filter=lambda job: job.maker is not None and job.maker.value == 42,
     )
     assert test_job.metadata["v"] == 5