From e1d69138eae050fd1a9f38ae710af13335862b6c Mon Sep 17 00:00:00 2001
From: Sebastian Urchs <surchs@users.noreply.github.com>
Date: Sat, 10 Jun 2023 18:21:46 -0400
Subject: [PATCH] [FIX] Detect categorical column based on annotations (#154)

* Detect categorical column based on Annotations

* Ensure BIDS levels alone doesn't pass as categorical
---
 bagel/pheno_utils.py        |  2 +-
 bagel/tests/test_utility.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/bagel/pheno_utils.py b/bagel/pheno_utils.py
index 23f336f..fa6c7b1 100644
--- 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
index 77bbe2c..2c0e95d 100644
--- 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",
     [