BUG: Ensure Index.astype('category') returns a CategoricalIndex

Author: jschendel

doc/source/whatsnew/v0.22.0.txt

@@ -249,6 +249,7 @@ Conversion
 - Fixed a bug where creating a Series from an array that contains both tz-naive and tz-aware values will result in a Series whose dtype is tz-aware instead of object (:issue:`16406`)
 - Adding a ``Period`` object to a ``datetime`` or ``Timestamp`` object will now correctly raise a ``TypeError`` (:issue:`17983`)
 - Fixed a bug where ``FY5253`` date offsets could incorrectly raise an ``AssertionError`` in arithmetic operatons (:issue:`14774`)
+- Bug in :meth:`Index.astype` with a categorical dtype where the resultant index would not be converted to a :class:`CategoricalIndex` for all types of index (:issue:`18630`)
 
 
 Indexing

pandas/core/dtypes/common.py

@@ -1934,7 +1934,7 @@ def pandas_dtype(dtype):
             except TypeError:
                 pass
 
-        elif dtype.startswith('interval[') or dtype.startswith('Interval['):
+        elif dtype.startswith('interval') or dtype.startswith('Interval'):
             try:
                 return IntervalDtype.construct_from_string(dtype)
             except TypeError:

pandas/core/indexes/base.py

@@ -1053,6 +1053,10 @@ def _to_embed(self, keep_tz=False, dtype=None):
 
     @Appender(_index_shared_docs['astype'])
     def astype(self, dtype, copy=True):
+        if is_categorical_dtype(dtype):
+            from pandas.core.indexes.category import CategoricalIndex
+            return CategoricalIndex(self.values, name=self.name, dtype=dtype,
+                                    copy=copy)
         return Index(self.values.astype(dtype, copy=copy), name=self.name,
                      dtype=dtype)
 

pandas/core/indexes/category.py

@@ -9,7 +9,8 @@
     _ensure_platform_int,
     is_list_like,
     is_interval_dtype,
-    is_scalar)
+    is_scalar,
+    pandas_dtype)
 from pandas.core.common import (_asarray_tuplesafe,
                                 _values_from_object)
 from pandas.core.dtypes.missing import array_equivalent, isna
@@ -341,9 +342,13 @@ def __array__(self, dtype=None):
 
     @Appender(_index_shared_docs['astype'])
     def astype(self, dtype, copy=True):
+        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) and (dtype == self.dtype):
+            # fastpath if dtype is the same current
+            return self.copy() if copy else self
         return super(CategoricalIndex, self).astype(dtype=dtype, copy=copy)
 
     @cache_readonly

pandas/core/indexes/datetimes.py

@@ -20,6 +20,7 @@
     is_period_dtype,
     is_bool_dtype,
     is_string_dtype,
+    is_categorical_dtype,
     is_string_like,
     is_list_like,
     is_scalar,
@@ -916,6 +917,10 @@ def astype(self, dtype, copy=True):
             elif copy is True:
                 return self.copy()
             return self
+        elif is_categorical_dtype(dtype):
+            from pandas.core.indexes.category import CategoricalIndex
+            return CategoricalIndex(self.values, name=self.name, dtype=dtype,
+                                    copy=copy)
         elif is_string_dtype(dtype):
             return Index(self.format(), name=self.name, dtype=object)
         elif is_period_dtype(dtype):

pandas/core/indexes/multi.py

@@ -17,6 +17,7 @@
     is_object_dtype,
     is_iterator,
     is_list_like,
+    pandas_dtype,
     is_scalar)
 from pandas.core.dtypes.missing import isna, array_equivalent
 from pandas.errors import PerformanceWarning, UnsortedIndexError
@@ -2715,7 +2716,7 @@ def difference(self, other):
 
     @Appender(_index_shared_docs['astype'])
     def astype(self, dtype, copy=True):
-        if not is_object_dtype(np.dtype(dtype)):
+        if not is_object_dtype(pandas_dtype(dtype)):
             raise TypeError('Setting %s dtype to anything other than object '
                             'is not supported' % self.__class__)
         elif copy is True:

pandas/core/indexes/numeric.py

@@ -7,6 +7,7 @@
     is_float_dtype,
     is_object_dtype,
     is_integer_dtype,
+    is_categorical_dtype,
     is_bool,
     is_bool_dtype,
     is_scalar)
@@ -313,10 +314,14 @@ def astype(self, dtype, copy=True):
             values = self._values.astype(dtype, copy=copy)
         elif is_object_dtype(dtype):
             values = self._values.astype('object', copy=copy)
+        elif is_categorical_dtype(dtype):
+            from pandas.core.indexes.category import CategoricalIndex
+            return CategoricalIndex(self, name=self.name, dtype=dtype,
+                                    copy=copy)
         else:
-            raise TypeError('Setting %s dtype to anything other than '
-                            'float64 or object is not supported' %
-                            self.__class__)
+            raise TypeError('Setting {cls} dtype to anything other than '
+                            'float64, object, or category is not supported'
+                            .format(cls=self.__class__))
         return Index(values, name=self.name, dtype=dtype)
 
     @Appender(_index_shared_docs['_convert_scalar_indexer'])

pandas/core/indexes/period.py

@@ -16,6 +16,7 @@
     is_timedelta64_dtype,
     is_period_dtype,
     is_bool_dtype,
+    is_categorical_dtype,
     pandas_dtype,
     _ensure_object)
 from pandas.core.dtypes.dtypes import PeriodDtype
@@ -517,6 +518,10 @@ def astype(self, dtype, copy=True, how='start'):
             return self.to_timestamp(how=how).tz_localize(dtype.tz)
         elif is_period_dtype(dtype):
             return self.asfreq(freq=dtype.freq)
+        elif is_categorical_dtype(dtype):
+            from pandas.core.indexes.category import CategoricalIndex
+            return CategoricalIndex(self.values, name=self.name, dtype=dtype,
+                                    copy=copy)
         raise TypeError('Cannot cast PeriodIndex to dtype %s' % dtype)
 
     @Substitution(klass='PeriodIndex')

pandas/core/indexes/timedeltas.py

@@ -12,6 +12,8 @@
     is_object_dtype,
     is_timedelta64_dtype,
     is_timedelta64_ns_dtype,
+    is_categorical_dtype,
+    pandas_dtype,
     _ensure_int64)
 from pandas.core.dtypes.missing import isna
 from pandas.core.dtypes.generic import ABCSeries
@@ -479,7 +481,7 @@ def to_pytimedelta(self):
 
     @Appender(_index_shared_docs['astype'])
     def astype(self, dtype, copy=True):
-        dtype = np.dtype(dtype)
+        dtype = pandas_dtype(dtype)
 
         if is_object_dtype(dtype):
             return self._box_values_as_index()
@@ -498,6 +500,10 @@ def astype(self, dtype, copy=True):
         elif is_integer_dtype(dtype):
             return Index(self.values.astype('i8', copy=copy), dtype='i8',
                          name=self.name)
+        elif is_categorical_dtype(dtype):
+            from pandas.core.indexes.category import CategoricalIndex
+            return CategoricalIndex(self.values, name=self.name, dtype=dtype,
+                                    copy=copy)
         raise TypeError('Cannot cast TimedeltaIndex to dtype %s' % dtype)
 
     def union(self, other):

pandas/tests/indexes/common.py

@@ -13,6 +13,7 @@
 from pandas.core.indexes.base import InvalidIndexError
 from pandas.core.indexes.datetimelike import DatetimeIndexOpsMixin
 from pandas.core.dtypes.common import needs_i8_conversion
+from pandas.core.dtypes.dtypes import CategoricalDtype
 from pandas._libs.tslib import iNaT
 
 import pandas.util.testing as tm
@@ -1058,3 +1059,20 @@ def test_putmask_with_wrong_mask(self):
 
         with pytest.raises(ValueError):
             index.putmask('foo', 1)
+
+    def test_astype_category(self):
+        # GH 18630
+        index = self.create_index()
+
+        expected = CategoricalIndex(index.values)
+        result = index.astype('category', copy=True)
+        tm.assert_index_equal(result, expected)
+
+        expected = CategoricalIndex(index.values, name='foo')
+        result = index.rename('foo').astype('category', copy=False)
+        tm.assert_index_equal(result, expected)
+
+        dtype = CategoricalDtype(index.unique()[:-1], ordered=True)
+        expected = CategoricalIndex(index.values, dtype=dtype)
+        result = index.astype(dtype)
+        tm.assert_index_equal(result, expected)

pandas/tests/indexes/test_interval.py

@@ -6,6 +6,7 @@
 from pandas import (
     Interval, IntervalIndex, Index, isna, notna, interval_range, Timestamp,
     Timedelta, compat, date_range, timedelta_range, DateOffset)
+from pandas.core.dtypes.dtypes import CategoricalDtype
 from pandas.compat import lzip
 from pandas.tseries.offsets import Day
 from pandas._libs.interval import IntervalTree
@@ -362,8 +363,15 @@ def test_astype(self, closed):
         tm.assert_index_equal(result, idx)
         assert result.equals(idx)
 
-        result = idx.astype('category')
+    def test_astype_category(self, closed):
+        # GH 18630
+        idx = self.create_index(closed=closed)
         expected = pd.Categorical(idx, ordered=True)
+
+        result = idx.astype('category')
+        tm.assert_categorical_equal(result, expected)
+
+        result = idx.astype(CategoricalDtype())
         tm.assert_categorical_equal(result, expected)
 
     @pytest.mark.parametrize('klass', [list, tuple, np.array, pd.Series])

pandas/tests/indexes/test_multi.py

@@ -16,6 +16,7 @@
                     compat, date_range, period_range)
 from pandas.compat import PY3, long, lrange, lzip, range, u, PYPY
 from pandas.errors import PerformanceWarning, UnsortedIndexError
+from pandas.core.dtypes.dtypes import CategoricalDtype
 from pandas.core.indexes.base import InvalidIndexError
 from pandas._libs import lib
 from pandas._libs.lib import Timestamp
@@ -554,6 +555,15 @@ def test_astype(self):
         with tm.assert_raises_regex(TypeError, "^Setting.*dtype.*object"):
             self.index.astype(np.dtype(int))
 
+    def test_astype_category(self):
+        # GH 18630
+        msg = 'Setting .* dtype to anything other than object is not supported'
+        with tm.assert_raises_regex(TypeError, msg):
+            self.index.astype('category')
+
+        with tm.assert_raises_regex(TypeError, msg):
+            self.index.astype(CategoricalDtype())
+
     def test_constructor_single_level(self):
         result = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux']],
                             labels=[[0, 1, 2, 3]], names=['first'])