fix tests

Author: R-Palazzo

rdt/transformers/utils.py

@@ -177,8 +177,11 @@ def strings_from_regex(regex, max_repeat=16):
 
 def _fill_nan_with_none_series(data):
     sentinel = object()
-    if isinstance(data.dtype, pd.CategoricalDtype):
+    dtype = data.dtype
+    if isinstance(dtype, pd.CategoricalDtype):
         data = data.cat.add_categories([sentinel])
+        data = data.fillna(sentinel).replace({sentinel: None})
+        return pd.Series(pd.Categorical(data, categories=dtype.categories), index=data.index)
 
     return data.fillna(sentinel).replace({sentinel: None})
 

tests/unit/transformers/test_utils.py

@@ -181,7 +181,9 @@ def test__fill_nan_with_none_series():
     # Assert
     expected_result = pd.Series([1.0, 2.0, 3.0, None], dtype='object')
     pd.testing.assert_series_equal(result, expected_result)
-    expected_result_categorical = pd.Series(['a', 'b', 'c', 'd', None], dtype='category')
+    expected_result_categorical = pd.Series(
+        pd.Categorical(['a', 'b', 'c', 'd', None], categories=['a', 'b', 'c', 'd'])
+    )
     pd.testing.assert_series_equal(result_categorical, expected_result_categorical)