From 4a3791cf111dcb246aba0b39d302aea2fc14c743 Mon Sep 17 00:00:00 2001
From: Michael Milton <michael.r.milton@gmail.com>
Date: Wed, 14 Aug 2024 13:28:42 +1000
Subject: [PATCH] Add channel and time to workflow data frames

---
 core/lls_core/models/results.py | 13 ++++++++++++-
 core/tests/test_workflows.py    |  5 ++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/core/lls_core/models/results.py b/core/lls_core/models/results.py
index 2b6d007d..99d57a0c 100644
--- a/core/lls_core/models/results.py
+++ b/core/lls_core/models/results.py
@@ -97,11 +97,12 @@ def process(self) -> Iterable[Tuple[RoiIndex, ProcessedWorkflowOutput]]:
 
         # Handle each ROI separately
         for roi, roi_results in groupby(self.slices, key=lambda it: it.roi_index):
-            values: list[Writer, dict, tuple, list] = []
+            values: list[Writer | list] = []
             for result in roi_results:
                 # Ensure the data is in a tuple
                 data = (result.data,) if is_arraylike(result.data) else result.data
                 for i, element in enumerate(data):
+                    # If the element is array like, we assume it's an image to write to disk
                     if is_arraylike(element):
                         # Make the writer the first time only
                         if len(values) <= i:
@@ -114,10 +115,20 @@ def process(self) -> Iterable[Tuple[RoiIndex, ProcessedWorkflowOutput]]:
                             )
                         )
                     else:
+                        # Otherwise, we assume it's one row to be added to a data frame
                         if len(values) <= i:
                             values.append([])
 
                         rows = cast(list, values[i])
+
+                        if isinstance(element, list):
+                            # If the row is a list, it has no column names
+                            # We add the channel and time 
+                            element = [f"T{result.time_index}", f"C{result.channel_index}"] + element
+                        elif isinstance(element, dict):
+                            # If the row is a dict, it has column names
+                            element = {"time": f"T{result.time_index}", "channel": f"C{result.channel_index}", **element}
+
                         rows.append(element)
 
             for element in values:
diff --git a/core/tests/test_workflows.py b/core/tests/test_workflows.py
index 113c5ca7..045fa8c5 100644
--- a/core/tests/test_workflows.py
+++ b/core/tests/test_workflows.py
@@ -75,12 +75,15 @@ def test_table_workflow(minimal_image_path: Path, table_workflow: Workflow):
             workflow = table_workflow,
             save_dir = tmpdir
         )
-        for roi, output in params.process_workflow().process():
+        for _roi, output in params.process_workflow().process():
             assert isinstance(output, (DataFrame, Path))
             if isinstance(output, DataFrame):
                 nrow, ncol = output.shape
                 assert nrow == params.nslices
                 assert ncol > 0
+                # Check that time and channel are included
+                assert output.iloc[0, 0] == "T0"
+                assert output.iloc[0, 1] == "C0"
             else:
                 assert valid_image_path(output)