From 5f107e0ec22ac5e1dca422d53d1377f5595c58d6 Mon Sep 17 00:00:00 2001
From: Michelle Wang <tomichellewang@gmail.com>
Date: Thu, 13 Jun 2024 17:09:22 -0400
Subject: [PATCH] [FIX] Write updated doughnut file to disk at the end of
 `DicomReorgWorkflow` and `BidsConversionWorkflow` (#260)

* write doughnut in DicomReorgWorkflow's cleanup step

* write doughnut in BidsConversionWorkflow's cleanup step

* address code review comment about docstrings
---
 .../nipoppy/workflows/bids_conversion.py      | 10 ++++++
 nipoppy_cli/nipoppy/workflows/dicom_reorg.py  | 10 ++++++
 .../tests/test_workflow_bids_conversion.py    | 33 +++++++++++++++++++
 .../tests/test_workflow_dicom_reorg.py        | 29 ++++++++++++++++
 4 files changed, 82 insertions(+)

diff --git a/nipoppy_cli/nipoppy/workflows/bids_conversion.py b/nipoppy_cli/nipoppy/workflows/bids_conversion.py
index 78d9eebf..ef00818c 100644
--- a/nipoppy_cli/nipoppy/workflows/bids_conversion.py
+++ b/nipoppy_cli/nipoppy/workflows/bids_conversion.py
@@ -97,3 +97,13 @@ def run_single(self, participant_id: str, session_id: str):
         )
 
         return invocation_and_descriptor
+
+    def run_cleanup(self, **kwargs):
+        """
+        Clean up after main BIDS conversion part is run.
+
+        Specifically:
+        - Write updated doughnut file
+        """
+        self.save_tabular_file(self.doughnut, self.layout.fpath_doughnut)
+        return super().run_cleanup(**kwargs)
diff --git a/nipoppy_cli/nipoppy/workflows/dicom_reorg.py b/nipoppy_cli/nipoppy/workflows/dicom_reorg.py
index c0b671f8..b4a426fa 100644
--- a/nipoppy_cli/nipoppy/workflows/dicom_reorg.py
+++ b/nipoppy_cli/nipoppy/workflows/dicom_reorg.py
@@ -175,3 +175,13 @@ def run_main(self):
                     "Error reorganizing DICOM files for participant "
                     f"{participant_id} session {session_id}: {exception}"
                 )
+
+    def run_cleanup(self):
+        """
+        Clean up after main DICOM reorg part is run.
+
+        Specifically:
+        - Write updated doughnut file
+        """
+        self.save_tabular_file(self.doughnut, self.layout.fpath_doughnut)
+        return super().run_cleanup()
diff --git a/nipoppy_cli/tests/test_workflow_bids_conversion.py b/nipoppy_cli/tests/test_workflow_bids_conversion.py
index 0501ef0d..80938a63 100644
--- a/nipoppy_cli/tests/test_workflow_bids_conversion.py
+++ b/nipoppy_cli/tests/test_workflow_bids_conversion.py
@@ -42,6 +42,39 @@ def test_setup(config: Config, tmp_path: Path):
     assert not workflow.dpath_pipeline.exists()
 
 
+@pytest.mark.parametrize(
+    "doughnut",
+    [
+        Doughnut(),
+        Doughnut(
+            data={
+                Doughnut.col_participant_id: ["01"],
+                Doughnut.col_visit_id: ["1"],
+                Doughnut.col_session_id: ["1"],
+                Doughnut.col_datatype: "['anat']",
+                Doughnut.col_participant_dicom_dir: ["01"],
+                Doughnut.col_in_raw_imaging: [True],
+                Doughnut.col_in_sourcedata: [True],
+                Doughnut.col_in_bids: [True],
+            }
+        ).validate(),
+    ],
+)
+def test_cleanup(doughnut: Doughnut, tmp_path: Path):
+    workflow = BidsConversionRunner(
+        dpath_root=tmp_path / "my_dataset",
+        pipeline_name="",
+        pipeline_version="",
+        pipeline_step="",
+    )
+    workflow.doughnut = doughnut
+
+    workflow.run_cleanup()
+
+    assert workflow.layout.fpath_doughnut.exists()
+    assert Doughnut.load(workflow.layout.fpath_doughnut).equals(doughnut)
+
+
 @pytest.mark.parametrize(
     "doughnut_data,participant_id,session_id,expected",
     [
diff --git a/nipoppy_cli/tests/test_workflow_dicom_reorg.py b/nipoppy_cli/tests/test_workflow_dicom_reorg.py
index 084a92c1..39df01e7 100644
--- a/nipoppy_cli/tests/test_workflow_dicom_reorg.py
+++ b/nipoppy_cli/tests/test_workflow_dicom_reorg.py
@@ -7,6 +7,7 @@
 import pytest
 
 from nipoppy.tabular.dicom_dir_map import DicomDirMap
+from nipoppy.tabular.doughnut import Doughnut
 from nipoppy.tabular.manifest import Manifest
 from nipoppy.utils import participant_id_to_bids_participant, session_id_to_bids_session
 from nipoppy.workflows.dicom_reorg import DicomReorgWorkflow, is_derived_dicom
@@ -430,3 +431,31 @@ def test_run_main(
 
             else:
                 assert not dpath_to_check.exists()
+
+
+@pytest.mark.parametrize(
+    "doughnut",
+    [
+        Doughnut(),
+        Doughnut(
+            data={
+                Doughnut.col_participant_id: ["01"],
+                Doughnut.col_visit_id: ["1"],
+                Doughnut.col_session_id: ["1"],
+                Doughnut.col_datatype: "['anat']",
+                Doughnut.col_participant_dicom_dir: ["01"],
+                Doughnut.col_in_raw_imaging: [True],
+                Doughnut.col_in_sourcedata: [True],
+                Doughnut.col_in_bids: [True],
+            }
+        ).validate(),
+    ],
+)
+def test_cleanup(doughnut: Doughnut, tmp_path: Path):
+    workflow = DicomReorgWorkflow(dpath_root=tmp_path / "my_dataset")
+    workflow.doughnut = doughnut
+
+    workflow.run_cleanup()
+
+    assert workflow.layout.fpath_doughnut.exists()
+    assert Doughnut.load(workflow.layout.fpath_doughnut).equals(doughnut)