Fix boolean types to be supported as categorical (#260)

sodascience · Feb 28, 2024 · 65bd520 · 65bd520
1 parent f7a46b9
commit 65bd520
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/metasyn/var.py b/metasyn/var.py
@@ -96,6 +96,7 @@ def get_var_type(series: pl.Series) -> str:
             "time": "time",
             "str": "string",
             "categorical": "categorical",
+            "bool": "categorical",
             "NoneType": "continuous",
         }
         try:

diff --git a/tests/test_var.py b/tests/test_var.py
@@ -97,11 +97,9 @@ def check_similar(series_a, series_b):
     ]
 )
 def test_categorical(tmp_path, series):
-    # series = pd.Series(np.random.choice(["a", "b", "c", None], size=100), dtype="category")
     new_series = check_var(series, "categorical", tmp_path)
     assert set(_series_drop_nans(series)) == set(np.unique(_series_drop_nans(new_series)))
 
-
 @mark.parametrize("dtype", ["int8", "int16", "int32", "int64", "int"])
 def test_integer(dtype, tmp_path):
     series = pd.Series([np.random.randint(0, 10) for _ in range(300)], dtype=dtype)
@@ -150,8 +148,10 @@ def test_string(tmp_path, series_type):
 )
 def test_bool(tmp_path, series_type):
     series = series_type(np.random.choice([True, False], size=100))
-    with raises(ValueError):
-        check_var(series, "categorical", tmp_path)
+    check_var(series, "categorical", tmp_path)
+    var = MetaVar.fit(series)
+    new_series = var.draw_series(10)
+    assert new_series.dtype == pl.Boolean
 
 
 @mark.parametrize(