getitem for Job delegates to its output

Author: vineetbansal

src/jobflow/core/job.py

@@ -389,6 +389,24 @@ def __repr__(self):
         name, uuid = self.name, self.uuid
         return f"Job({name=}, {uuid=})"
 
+    def __getitem__(self, key: Any) -> OutputReference:
+        """
+        Get the corresponding `OutputReference` for the `Job`.
+
+        This is for when it is indexed like a dictionary or list.
+
+        Parameters
+        ----------
+        key
+            The index/key.
+
+        Returns
+        -------
+        OutputReference
+            The equivalent of `Job.output[k]`
+        """
+        return self.output[key]
+
     def __contains__(self, item: Hashable) -> bool:
         """
         Check if the job contains a reference to a given UUID.

tests/core/test_flow.py

@@ -1353,3 +1353,23 @@ def test_flow_repr():
     assert len(lines) == len(flow_repr)
     for expected, line in zip(lines, flow_repr):
         assert line.startswith(expected), f"{line=} doesn't start with {expected=}"
+
+
+def test_get_item():
+    from jobflow import Flow, job, run_locally
+
+    @job
+    def make_str(s):
+        return {"hello": s}
+
+    @job
+    def capitalize(s):
+        return s.upper()
+
+    job1 = make_str("world")
+    job2 = capitalize(job1["hello"])
+
+    flow = Flow([job1, job2])
+
+    responses = run_locally(flow, ensure_success=True)
+    assert responses[job2.uuid][1].output == "WORLD"