[FIX] Detect categorical column based on annotations (#154)

* Detect categorical column based on Annotations * Ensure BIDS levels alone doesn't pass as categorical
neurobagel · Jun 10, 2023 · e1d6913 · e1d6913
1 parent dd54f14
commit e1d6913
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 1 deletion.
diff --git a/bagel/pheno_utils.py b/bagel/pheno_utils.py
@@ -106,7 +106,7 @@ def is_missing_value(
 
 def is_column_categorical(column: str, data_dict: dict) -> bool:
     """Determine whether a column in a Neurobagel data dictionary is categorical"""
-    if "Levels" in data_dict[column]:
+    if "Levels" in data_dict[column]["Annotations"]:
         return True
     return False
 

diff --git a/bagel/tests/test_utility.py b/bagel/tests/test_utility.py
@@ -69,6 +69,41 @@ def test_get_transformed_categorical_value(test_data, load_test_json):
     )
 
 
+@pytest.mark.parametrize(
+    "example,expected_result",
+    [
+        (
+            {
+                "column": {
+                    "Annotations": {
+                        "IsAbout": {"TermURL": "something", "Labels": "other"},
+                        "Levels": {
+                            "val1": {"TermURL": "something", "Label": "other"}
+                        },
+                    }
+                }
+            },
+            True,
+        ),
+        (
+            {
+                "column": {
+                    "Levels": {"val1": "some description"},
+                    "Annotations": {
+                        "IsAbout": {"TermURL": "something", "Labels": "other"}
+                    },
+                }
+            },
+            False,
+        ),
+    ],
+)
+def test_detect_categorical_column(example, expected_result):
+    result = putil.is_column_categorical(column="column", data_dict=example)
+
+    assert result is expected_result
+
+
 @pytest.mark.parametrize(
     "value,column,expected",
     [