BUG: Allow a categorical series to accept dict or Series in fillna (pandas-dev#17033)

cleaning up comments and tests

Author: reidy-p

pandas/core/categorical.py

@@ -1660,16 +1660,52 @@ def fillna(self, value=None, method=None, limit=None):
 
         else:
 
-            if not isna(value) and value not in self.categories:
-                raise ValueError("fill value must be in categories")
-
-            mask = values == -1
-            if mask.any():
-                values = values.copy()
-                if isna(value):
-                    values[mask] = -1
+            if isinstance(value, ABCSeries):
+                if not value[~value.isin(self.categories)].isna().all():
+                    raise ValueError("fill value must be in categories")
+
+                # Check if single scalar in the value Series
+                # (e.g., s.fillna(pd.Series('a')))
+                if (len(value[value.notna()]) == 1 and
+                    value[value.notna()].index == 0):
+                    values_codes = _get_codes_for_values(value[value.notna()],
+                                                         self.categories)
+                    mask = values == -1
+                    values[mask] = values_codes
                 else:
-                    values[mask] = self.categories.get_loc(value)
+                    values_codes = _get_codes_for_values(value,
+                                                         self.categories)
+                    index = np.where(values_codes != -1)
+                    values[index] = values_codes[values_codes != -1]
+                    # from pandas import Series
+                    # values_codes = Series(_get_codes_for_values(value[value.notna()],
+                    #                       self.categories), index=value[value.notna()].index)
+                    # values[values_codes.index] = values_codes
+
+            elif isinstance(value, dict):
+                from pandas import Series
+                value = Series(value)
+
+                if not value[~value.isin(self.categories)].isna().all():
+                    raise ValueError("fill value must be in categories")
+
+                # Convert to Series to allow use of index values
+                values_codes = Series(_get_codes_for_values(value,
+                                      self.categories), index=value.index)
+                values[values_codes.index] = values_codes
+
+            # Scalar value
+            else:
+                if not isna(value) and value not in self.categories:
+                    raise ValueError("fill value must be in categories")
+
+                mask = values == -1
+                if mask.any():
+                    values = values.copy()
+                    if isna(value):
+                        values[mask] = -1
+                    else:
+                        values[mask] = self.categories.get_loc(value)
 
         return self._constructor(values, categories=self.categories,
                                  ordered=self.ordered, fastpath=True)

pandas/core/generic.py

@@ -32,7 +32,7 @@
 from pandas.core.common import (_count_not_none,
                                 _maybe_box_datetimelike, _values_from_object,
                                 AbstractMethodError, SettingWithCopyError,
-                                SettingWithCopyWarning)
+                                SettingWithCopyWarning, is_categorical_dtype)
 
 from pandas.core.base import PandasObject, SelectionMixin
 from pandas.core.index import (Index, MultiIndex, _ensure_index,
@@ -4298,7 +4298,9 @@ def fillna(self, value=None, method=None, axis=None, inplace=False,
                 return self
 
             if self.ndim == 1:
-                if isinstance(value, (dict, ABCSeries)):
+                if isinstance(value, dict) and is_categorical_dtype(self):
+                    pass
+                elif isinstance(value, (ABCSeries, dict)):
                     from pandas import Series
                     value = Series(value)
                 elif not is_list_like(value):

pandas/tests/test_categorical.py

@@ -4601,6 +4601,38 @@ def f():
         df = pd.DataFrame({'a': pd.Categorical(idx)})
         tm.assert_frame_equal(df.fillna(value=pd.NaT), df)
 
+    @pytest.mark.parametrize('fill_value expected_output', [
+        ('a', ['a', 'a', 'b', 'a', 'a']),
+        ({1: 'a', 3: 'b', 4: 'b'}, ['a', 'a', 'b', 'b', 'b']),
+        ({1: 'a'}, ['a', 'a', 'b', np.nan, np.nan]),
+        ({1: 'a', 3: 'b'}, ['a', 'a', 'b', 'b', np.nan]),
+        (pd.Series('a'), ['a', 'a', 'b', 'a', 'a']),
+        (pd.Series({1: 'a', 3: 'b'}), ['a', 'a', 'b', 'b', np.nan])
+    ])
+    def fillna_series_categorical(self, fill_value, expected_output):
+        # GH 17033
+        # Test fillna for a Categorical series
+        data = ['a', np.nan, 'b', np.nan, np.nan]
+        s = pd.Series(pd.Categorical(data, categories=['a', 'b']))
+        exp = pd.Series(pd.Categorical(expected_output, categories=['a', 'b']))
+        tm.assert_series_equal(s.fillna(fill_value), exp)
+        
+        s1 = pd.Series(data, categories=['a', 'b'])
+        with tm.assert_raises_regex(ValueError,
+                                    "fill value must be in categories"):
+            s1.fillna('cat')
+
+        s2 = pd.Series(data, categories=['a', 'b'])
+        with tm.assert_raises_regex(ValueError,
+                                    "fill value must be in categories"):
+            s2.fillna(pd.Series('cat'))
+
+        s3 = pd.Series(data, categories=['a', 'b'])
+        with tm.assert_raises_regex(TypeError,
+                                    '"value" parameter must be a scalar or '
+                                    'dict but you passed a "list"'):
+            s3.fillna(['a', 'b'])
+
     def test_astype_to_other(self):
 
         s = self.cat['value_group']