refactor dtype update

Author: jschendel

pandas/core/dtypes/dtypes.py

@@ -340,6 +340,33 @@ def _validate_categories(categories, fastpath=False):
 
         return categories
 
+    def _update_dtype(self, dtype):
+        """
+        Returns a CategoricalDtype with categories and ordered taken from dtype
+        if specified, otherwise falling back to self if unspecified
+
+        Parameters
+        ----------
+        dtype : CategoricalDtype
+
+        Returns
+        -------
+        new_dtype : CategoricalDtype
+        """
+        if isinstance(dtype, compat.string_types) and dtype == 'category':
+            # dtype='category' should not change anything
+            return self
+        elif not self.is_dtype(dtype):
+            msg = ('a CategoricalDtype must be passed to perform an update, '
+                   'got {dtype!r}').format(dtype=dtype)
+            raise ValueError(msg)
+
+        # dtype is CDT: keep current categories if None (ordered can't be None)
+        new_categories = dtype.categories
+        if new_categories is None:
+            new_categories = self.categories
+        return CategoricalDtype(new_categories, dtype.ordered)
+
     @property
     def categories(self):
         """

pandas/core/indexes/category.py

@@ -10,8 +10,7 @@
     _ensure_platform_int,
     is_list_like,
     is_interval_dtype,
-    is_scalar,
-    pandas_dtype)
+    is_scalar)
 from pandas.core.common import (_asarray_tuplesafe,
                                 _values_from_object)
 from pandas.core.dtypes.missing import array_equivalent, isna
@@ -341,23 +340,12 @@ def __array__(self, dtype=None):
 
     @Appender(_index_shared_docs['astype'])
     def astype(self, dtype, copy=True):
-        if isinstance(dtype, compat.string_types) and dtype == 'category':
-            # GH 18630: CI.astype('category') should not change anything
-            return self.copy() if copy else self
-
-        dtype = pandas_dtype(dtype)
         if is_interval_dtype(dtype):
             from pandas import IntervalIndex
             return IntervalIndex.from_intervals(np.array(self))
         elif is_categorical_dtype(dtype):
-            # GH 18630: keep current categories if None (ordered can't be None)
-            if dtype.categories is None:
-                new_categories = self.categories
-            else:
-                new_categories = dtype.categories
-            dtype = CategoricalDtype(new_categories, dtype.ordered)
-
-            # fastpath if dtypes are equal
+            # GH 18630
+            dtype = self.dtype._update_dtype(dtype)
             if dtype == self.dtype:
                 return self.copy() if copy else self
 

pandas/tests/dtypes/test_dtypes.py

@@ -9,6 +9,7 @@
 from pandas import (
     Series, Categorical, CategoricalIndex, IntervalIndex, date_range)
 
+from pandas.compat import string_types
 from pandas.core.dtypes.dtypes import (
     DatetimeTZDtype, PeriodDtype,
     IntervalDtype, CategoricalDtype)
@@ -123,6 +124,41 @@ def test_tuple_categories(self):
         result = CategoricalDtype(categories)
         assert all(result.categories == categories)
 
+    @pytest.mark.parametrize('dtype', [
+        CategoricalDtype(list('abc'), False),
+        CategoricalDtype(list('abc'), True)])
+    @pytest.mark.parametrize('new_dtype', [
+        'category',
+        CategoricalDtype(None, False),
+        CategoricalDtype(None, True),
+        CategoricalDtype(list('abc'), False),
+        CategoricalDtype(list('abc'), True),
+        CategoricalDtype(list('cba'), False),
+        CategoricalDtype(list('cba'), True),
+        CategoricalDtype(list('wxyz'), False),
+        CategoricalDtype(list('wxyz'), True)])
+    def test_update_dtype(self, dtype, new_dtype):
+        if isinstance(new_dtype, string_types) and new_dtype == 'category':
+            expected_categories = dtype.categories
+            expected_ordered = dtype.ordered
+        else:
+            expected_categories = new_dtype.categories
+            if expected_categories is None:
+                expected_categories = dtype.categories
+            expected_ordered = new_dtype.ordered
+
+        result = dtype._update_dtype(new_dtype)
+        tm.assert_index_equal(result.categories, expected_categories)
+        assert result.ordered is expected_ordered
+
+    @pytest.mark.parametrize('bad_dtype', [
+        'foo', object, np.int64, PeriodDtype('Q'), IntervalDtype(object)])
+    def test_update_dtype_errors(self, bad_dtype):
+        dtype = CategoricalDtype(list('abc'), False)
+        msg = 'a CategoricalDtype must be passed to perform an update, '
+        with tm.assert_raises_regex(ValueError, msg):
+            dtype._update_dtype(bad_dtype)
+
 
 class TestDatetimeTZDtype(Base):