From 70c5daf57567127dd6d927bc4352f0fb1bea3f0c Mon Sep 17 00:00:00 2001
From: 36000 <jk6.28@outlook.com>
Date: Wed, 31 Jan 2024 15:28:33 -0800
Subject: [PATCH 1/2] [ENH] Make the cleaning easier to interact with

---
 AFQ/segmentation.py       | 12 +++++++++++-
 AFQ/tasks/segmentation.py | 17 ++++++++++++-----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/AFQ/segmentation.py b/AFQ/segmentation.py
index 010aa07ac..aeef70548 100644
--- a/AFQ/segmentation.py
+++ b/AFQ/segmentation.py
@@ -139,7 +139,8 @@ def __init__(self,
                  presegment_kwargs={},
                  filter_by_endpoints=True,
                  dist_to_atlas=4,
-                 save_intermediates=None):
+                 save_intermediates=None,
+                 cleaning_params={}):
         """
         Segment streamlines into bundles.
 
@@ -260,6 +261,11 @@ def __init__(self,
         save_intermediates : str, optional
             The full path to a folder into which intermediate products
             are saved. Default: None, means no saving of intermediates.
+        cleaning_params : dict, optional
+            Cleaning params to pass to seg.clean_bundle. This will
+            override the default parameters of that method. However, this
+            can be overriden by setting the cleaning parameters in the
+            bundle_dict. Default: {}.
 
         References
         ----------
@@ -297,6 +303,7 @@ def __init__(self,
         self.filter_by_endpoints = filter_by_endpoints
         self.dist_to_atlas = dist_to_atlas
         self.parallel_segmentation = parallel_segmentation
+        self.cleaning_params = cleaning_params
 
         if (save_intermediates is not None) and \
                 (not op.exists(save_intermediates)):
@@ -774,6 +781,9 @@ def segment_afq(self, tg=None):
             if b_sls:
                 accept_idx = b_sls.initiate_selection("Mahalanobis")
                 clean_params = bundle_def.get("mahal", {})
+                clean_params = {
+                    **self.cleaning_params,
+                    **clean_params}
                 clean_params["return_idx"] = True
                 cut = self.clip_edges or ("bundlesection" in bundle_def)
                 _, cleaned_idx = clean_bundle(
diff --git a/AFQ/tasks/segmentation.py b/AFQ/tasks/segmentation.py
index 6eca46ab3..044ff92a6 100644
--- a/AFQ/tasks/segmentation.py
+++ b/AFQ/tasks/segmentation.py
@@ -89,13 +89,20 @@ def segment(data_imap, mapping_imap,
     else:
         tgram, meta = seg_sft.get_sft_and_sidecar()
 
-    segmentation_params_out = {
-        arg_name: value if isinstance(value, (int, float, bool, str)) or (
-            value is None) else str(value)
-        for arg_name, value in segmentation_params.items()}
+    seg_params_out = {}
+    for arg_name, value in segmentation_params.items():
+        if isinstance(value, (int, float, bool, str)):
+            seg_params_out[arg_name] = value
+        elif isinstance(value, (list, tuple)):
+            seg_params_out[arg_name] = [str(v) for v in value]
+        elif isinstance(value, dict):
+            for k, v in value.items():
+                seg_params_out[k] = str(v)
+        else:
+            seg_params_out[arg_name] = str(value)
 
     meta["source"] = streamlines
-    meta["Recognition Parameters"] = segmentation_params_out
+    meta["Recognition Parameters"] = seg_params_out
     meta["Timing"] = time() - start_time
     return tgram, meta
 

From 05f1ab67defc4a393145ad77cb9ad76994463cf5 Mon Sep 17 00:00:00 2001
From: 36000 <jk6.28@outlook.com>
Date: Wed, 31 Jan 2024 15:46:27 -0800
Subject: [PATCH 2/2] moto version pin

---
 setup.cfg | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.cfg b/setup.cfg
index 028ab7643..955e24dd9 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -71,7 +71,7 @@ dev =
     sphinx-autoapi
     rapidfuzz
     xvfbwrapper==0.2.9
-    moto>=3.0.0
+    moto>=3.0.0,<5.0.0
     pip-conflict-checker>=0.6.0
     pydata-sphinx-theme
     sphinx-design