Warn about bundle dict instead of KeyError

yeatmanlab · Oct 9, 2020 · d9e04b8 · d9e04b8
1 parent 2f0f5f3
commit d9e04b8
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/AFQ/api.py b/AFQ/api.py
@@ -147,16 +147,21 @@ def make_bundle_dict(bundle_names=DEFAULT_BUNDLES,
                     uid += 1
             else:
                 for hemi in ['_R', '_L']:
-                    afq_bundles[name + hemi] = {
-                        'ROIs': [templates[name + '_roi1' + hemi],
-                                 templates[name + '_roi2' + hemi]],
-                        'rules': [True, True],
-                        'prob_map': templates[name + hemi + '_prob_map'],
-                        'cross_midline': False,
-                        'uid': uid}
+                    if (templates.get(name + '_roi1' + hemi)
+                        and templates.get(name + '_roi2' + hemi)
+                        and templates.get(name + hemi + '_prob_map')):
+                        afq_bundles[name + hemi] = {
+                            'ROIs': [templates[name + '_roi1' + hemi],
+                                     templates[name + '_roi2' + hemi]],
+                            'rules': [True, True],
+                            'prob_map': templates[name + hemi + '_prob_map'],
+                            'cross_midline': False,
+                            'uid': uid}
+                    else:
+                        logger = logging.getLogger('AFQ.api')
+                        logger.warning(f"{name} is not in AFQ templates")
 
                     uid += 1
-
     elif seg_algo.startswith("reco"):
         if seg_algo.endswith("80"):
             bundle_dict = afd.read_hcp_atlas(80)

diff --git a/AFQ/tests/test_api.py b/AFQ/tests/test_api.py
@@ -190,10 +190,9 @@ def test_make_bundle_dict():
 
     # Vertical Occipital Fasciculus
     # not included and does not exist in templates
-    try:
-        api.make_bundle_dict(bundle_names=["VOF"])
-    except KeyError:
-        pass
+    afq_bundles = api.make_bundle_dict(bundle_names=["VOF"])
+
+    assert len(afq_bundles) == 0
 
 
 @pytest.mark.nightly4
@@ -205,7 +204,7 @@ def test_AFQ_custom_tract():
     _, bids_path, sub_path = get_temp_hardi()
     afd.fetch_stanford_hardi_tractography()
 
-    bundle_names = ["SLF", "ARC", "CST", "FP"]
+    bundle_names = ["SLF", "ARC", "CST", "FP", "AntFrontal"]
 
     # move subsampled tractography into bids folder
     os.rename(