From 14a6b78f15a09bf553b883490639f14fde616b81 Mon Sep 17 00:00:00 2001
From: Taylor Salo <tsalo006@fiu.edu>
Date: Mon, 13 Nov 2023 12:26:27 -0500
Subject: [PATCH] Use context manager for fMRIPrep workflow.

---
 aslprep/workflows/asl/base.py      | 29 ++++++++++++++++++++++-------
 aslprep/workflows/asl/confounds.py |  1 +
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/aslprep/workflows/asl/base.py b/aslprep/workflows/asl/base.py
index adf5a37f7..3431679d2 100644
--- a/aslprep/workflows/asl/base.py
+++ b/aslprep/workflows/asl/base.py
@@ -3,6 +3,7 @@
 """Preprocessing workflows for ASL data."""
 import nibabel as nb
 import numpy as np
+from fmriprep.workflows.bold import confounds
 from fmriprep.workflows.bold.base import get_estimator
 from fmriprep.workflows.bold.registration import init_bold_reg_wf, init_bold_t1_trans_wf
 from fmriprep.workflows.bold.resampling import (
@@ -25,7 +26,7 @@
 from aslprep.utils.bids import collect_run_data
 from aslprep.utils.misc import _create_mem_gb, _get_wf_name
 from aslprep.workflows.asl.cbf import init_compute_cbf_wf, init_parcellate_cbf_wf
-from aslprep.workflows.asl.confounds import init_asl_confounds_wf, init_carpetplot_wf
+from aslprep.workflows.asl.confounds import init_asl_confounds_wf
 from aslprep.workflows.asl.hmc import init_asl_hmc_wf
 from aslprep.workflows.asl.outputs import init_asl_derivatives_wf
 from aslprep.workflows.asl.plotting import init_plot_cbf_wf
@@ -34,6 +35,19 @@
 from aslprep.workflows.asl.util import init_asl_reference_wf, init_validate_asl_wf
 
 
+class OverrideConfoundsDerivativesDataSink:
+    def __enter__(self):
+        # Save the original class
+        self.original_class = confounds.DerivativesDataSink
+        # Replace SomeClass with YourOwnClass
+        confounds.DerivativesDataSink = DerivativesDataSink
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Restore the original class
+        confounds.DerivativesDataSink = self.original_class
+
+
 def init_asl_preproc_wf(asl_file, has_fieldmap=False):
     """Perform the functional preprocessing stages of ASLPrep.
 
@@ -1010,12 +1024,13 @@ def init_asl_preproc_wf(asl_file, has_fieldmap=False):
     # Standard-space outputs requested.
     # Since ASLPrep automatically includes MNI152NLin2009cAsym, this should always be reached.
     if spaces.get_spaces(nonstandard=False, dim=(3,)):
-        carpetplot_wf = init_carpetplot_wf(
-            mem_gb=mem_gb["resampled"],
-            metadata=metadata,
-            cifti_output=config.workflow.cifti_output,
-            name="carpetplot_wf",
-        )
+        with OverrideConfoundsDerivativesDataSink:
+            carpetplot_wf = confounds.init_carpetplot_wf(
+                mem_gb=mem_gb["resampled"],
+                metadata=metadata,
+                cifti_output=config.workflow.cifti_output,
+                name="carpetplot_wf",
+            )
 
         # Xform to "MNI152NLin2009cAsym" is always computed.
         carpetplot_select_std = pe.Node(
diff --git a/aslprep/workflows/asl/confounds.py b/aslprep/workflows/asl/confounds.py
index 446165e9e..cea36acc6 100644
--- a/aslprep/workflows/asl/confounds.py
+++ b/aslprep/workflows/asl/confounds.py
@@ -305,6 +305,7 @@ def init_carpetplot_wf(
     """Build a workflow to generate *carpet* plots.
 
     XXX: Copied from fMRIPrep. Needs to be replaced with some version that works for ASLPrep.
+    TODO: Find a solution that directly uses fMRIPrep's.
 
     Resamples the MNI parcellation (ad-hoc parcellation derived from the
     Harvard-Oxford template and others).