Skip to content

Commit b812e6a

Browse files
committed
CLN: Consolidate Index.astype
1 parent cdebcf3 commit b812e6a

File tree

11 files changed

+109
-136
lines changed

11 files changed

+109
-136
lines changed

pandas/core/indexes/base.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -1069,8 +1069,12 @@ def astype(self, dtype, copy=True):
10691069
from .category import CategoricalIndex
10701070
return CategoricalIndex(self.values, name=self.name, dtype=dtype,
10711071
copy=copy)
1072-
return Index(self.values.astype(dtype, copy=copy), name=self.name,
1073-
dtype=dtype)
1072+
try:
1073+
return Index(self.values.astype(dtype, copy=copy), name=self.name,
1074+
dtype=dtype)
1075+
except (TypeError, ValueError):
1076+
msg = 'Cannot cast {name} to dtype {dtype}'
1077+
raise TypeError(msg.format(name=type(self).__name__, dtype=dtype))
10741078

10751079
def _to_safe_for_reshape(self):
10761080
""" convert to object if we are a categorical """

pandas/core/indexes/datetimelike.py

+28-6
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,22 @@
1111

1212
import numpy as np
1313
from pandas.core.dtypes.common import (
14-
is_integer, is_float,
15-
is_bool_dtype, _ensure_int64,
16-
is_scalar, is_dtype_equal,
17-
is_list_like, is_timedelta64_dtype)
14+
_ensure_int64,
15+
is_dtype_equal,
16+
is_float,
17+
is_integer,
18+
is_list_like,
19+
is_scalar,
20+
is_bool_dtype,
21+
is_categorical_dtype,
22+
is_datetime_or_timedelta_dtype,
23+
is_float_dtype,
24+
is_integer_dtype,
25+
is_object_dtype,
26+
is_string_dtype,
27+
is_timedelta64_dtype)
1828
from pandas.core.dtypes.generic import (
19-
ABCIndex, ABCSeries,
20-
ABCPeriodIndex, ABCIndexClass)
29+
ABCIndex, ABCSeries, ABCPeriodIndex, ABCIndexClass)
2130
from pandas.core.dtypes.missing import isna
2231
from pandas.core import common as com, algorithms
2332
from pandas.core.algorithms import checked_add_with_arr
@@ -859,6 +868,19 @@ def _concat_same_dtype(self, to_concat, name):
859868
new_data = np.concatenate([c.asi8 for c in to_concat])
860869
return self._simple_new(new_data, **attribs)
861870

871+
def astype(self, dtype, copy=True):
872+
if is_object_dtype(dtype):
873+
return self._box_values_as_index()
874+
elif is_string_dtype(dtype) and not is_categorical_dtype(dtype):
875+
return Index(self.format(), name=self.name, dtype=object)
876+
elif is_integer_dtype(dtype):
877+
return Index(self.values.astype('i8', copy=copy), name=self.name,
878+
dtype='i8')
879+
elif is_float_dtype(dtype) or is_datetime_or_timedelta_dtype(dtype):
880+
msg = 'Cannot cast {name} to dtype {dtype}'
881+
raise TypeError(msg.format(name=type(self).__name__, dtype=dtype))
882+
return super(DatetimeIndexOpsMixin, self).astype(dtype, copy=copy)
883+
862884

863885
def _ensure_datetimelike_to_i8(other):
864886
""" helper for coercing an input scalar or array to i8 """

pandas/core/indexes/datetimes.py

+14-27
Original file line numberDiff line numberDiff line change
@@ -10,17 +10,19 @@
1010
from pandas.core.base import _shared_docs
1111

1212
from pandas.core.dtypes.common import (
13-
_NS_DTYPE, _INT64_DTYPE,
14-
is_object_dtype, is_datetime64_dtype,
15-
is_datetimetz, is_dtype_equal,
13+
_INT64_DTYPE,
14+
_NS_DTYPE,
15+
is_object_dtype,
16+
is_datetime64_dtype,
17+
is_datetimetz,
18+
is_dtype_equal,
1619
is_timedelta64_dtype,
17-
is_integer, is_float,
20+
is_integer,
21+
is_float,
1822
is_integer_dtype,
1923
is_datetime64_ns_dtype,
2024
is_period_dtype,
2125
is_bool_dtype,
22-
is_string_dtype,
23-
is_categorical_dtype,
2426
is_string_like,
2527
is_list_like,
2628
is_scalar,
@@ -36,20 +38,17 @@
3638
from pandas.core.algorithms import checked_add_with_arr
3739

3840
from pandas.core.indexes.base import Index, _index_shared_docs
39-
from pandas.core.indexes.category import CategoricalIndex
4041
from pandas.core.indexes.numeric import Int64Index, Float64Index
4142
import pandas.compat as compat
42-
from pandas.tseries.frequencies import (
43-
to_offset, get_period_alias,
44-
Resolution)
43+
from pandas.tseries.frequencies import to_offset, get_period_alias, Resolution
4544
from pandas.core.indexes.datetimelike import (
4645
DatelikeOps, TimelikeOps, DatetimeIndexOpsMixin)
4746
from pandas.tseries.offsets import (
4847
DateOffset, generate_range, Tick, CDay, prefix_mapping)
4948

5049
from pandas.core.tools.timedeltas import to_timedelta
51-
from pandas.util._decorators import (Appender, cache_readonly,
52-
deprecate_kwarg, Substitution)
50+
from pandas.util._decorators import (
51+
Appender, cache_readonly, deprecate_kwarg, Substitution)
5352
import pandas.core.common as com
5453
import pandas.tseries.offsets as offsets
5554
import pandas.core.tools.datetimes as tools
@@ -906,25 +905,13 @@ def _format_native_types(self, na_rep='NaT', date_format=None, **kwargs):
906905
@Appender(_index_shared_docs['astype'])
907906
def astype(self, dtype, copy=True):
908907
dtype = pandas_dtype(dtype)
909-
if is_object_dtype(dtype):
910-
return self._box_values_as_index()
911-
elif is_integer_dtype(dtype):
912-
return Index(self.values.astype('i8', copy=copy), name=self.name,
913-
dtype='i8')
914-
elif is_datetime64_ns_dtype(dtype):
908+
if is_datetime64_ns_dtype(dtype):
915909
if self.tz is not None:
916910
return self.tz_convert('UTC').tz_localize(None)
917-
elif copy is True:
918-
return self.copy()
919-
return self
920-
elif is_categorical_dtype(dtype):
921-
return CategoricalIndex(self.values, name=self.name, dtype=dtype,
922-
copy=copy)
923-
elif is_string_dtype(dtype):
924-
return Index(self.format(), name=self.name, dtype=object)
911+
return self.copy() if copy else self
925912
elif is_period_dtype(dtype):
926913
return self.to_period(freq=dtype.freq)
927-
raise TypeError('Cannot cast DatetimeIndex to dtype %s' % dtype)
914+
return super(DatetimeIndex, self).astype(dtype, copy=copy)
928915

929916
def _get_time_micros(self):
930917
values = self.asi8

pandas/core/indexes/interval.py

+2-13
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,6 @@
1212
is_datetime_or_timedelta_dtype,
1313
is_datetime64tz_dtype,
1414
is_integer_dtype,
15-
is_object_dtype,
16-
is_categorical_dtype,
1715
is_float_dtype,
1816
is_interval_dtype,
1917
is_scalar,
@@ -29,7 +27,6 @@
2927
Interval, IntervalMixin, IntervalTree,
3028
intervals_to_interval_bounds)
3129

32-
from pandas.core.indexes.category import CategoricalIndex
3330
from pandas.core.indexes.datetimes import date_range
3431
from pandas.core.indexes.timedeltas import timedelta_range
3532
from pandas.core.indexes.multi import MultiIndex
@@ -671,16 +668,8 @@ def copy(self, deep=False, name=None):
671668
@Appender(_index_shared_docs['astype'])
672669
def astype(self, dtype, copy=True):
673670
if is_interval_dtype(dtype):
674-
if copy:
675-
self = self.copy()
676-
return self
677-
elif is_object_dtype(dtype):
678-
return Index(self.values, dtype=object)
679-
elif is_categorical_dtype(dtype):
680-
return CategoricalIndex(self.values, name=self.name, dtype=dtype,
681-
copy=copy)
682-
raise ValueError('Cannot cast IntervalIndex to dtype {dtype}'
683-
.format(dtype=dtype))
671+
return self.copy() if copy else self
672+
return super(IntervalIndex, self).astype(dtype, copy=copy)
684673

685674
@cache_readonly
686675
def dtype(self):

pandas/core/indexes/numeric.py

+9-20
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,8 @@
44
from pandas.core.dtypes.common import (
55
is_dtype_equal,
66
pandas_dtype,
7-
is_float_dtype,
8-
is_object_dtype,
7+
needs_i8_conversion,
98
is_integer_dtype,
10-
is_categorical_dtype,
119
is_bool,
1210
is_bool_dtype,
1311
is_scalar)
@@ -17,7 +15,6 @@
1715
from pandas.core import algorithms
1816
from pandas.core.indexes.base import (
1917
Index, InvalidIndexError, _index_shared_docs)
20-
from pandas.core.indexes.category import CategoricalIndex
2118
from pandas.util._decorators import Appender, cache_readonly
2219
import pandas.core.dtypes.concat as _concat
2320
import pandas.core.indexes.base as ibase
@@ -315,22 +312,14 @@ def inferred_type(self):
315312
@Appender(_index_shared_docs['astype'])
316313
def astype(self, dtype, copy=True):
317314
dtype = pandas_dtype(dtype)
318-
if is_float_dtype(dtype):
319-
values = self._values.astype(dtype, copy=copy)
320-
elif is_integer_dtype(dtype):
321-
if self.hasnans:
322-
raise ValueError('cannot convert float NaN to integer')
323-
values = self._values.astype(dtype, copy=copy)
324-
elif is_object_dtype(dtype):
325-
values = self._values.astype('object', copy=copy)
326-
elif is_categorical_dtype(dtype):
327-
return CategoricalIndex(self, name=self.name, dtype=dtype,
328-
copy=copy)
329-
else:
330-
raise TypeError('Setting {cls} dtype to anything other than '
331-
'float64, object, or category is not supported'
332-
.format(cls=self.__class__))
333-
return Index(values, name=self.name, dtype=dtype)
315+
if needs_i8_conversion(dtype):
316+
msg = ('Cannot convert Float64Index to dtype {dtype}; integer '
317+
'values are required for conversion').format(dtype=dtype)
318+
raise TypeError(msg)
319+
elif is_integer_dtype(dtype) and self.hasnans:
320+
# GH 13149
321+
raise ValueError('Cannot convert NA to integer')
322+
return super(Float64Index, self).astype(dtype, copy=copy)
334323

335324
@Appender(_index_shared_docs['_convert_scalar_indexer'])
336325
def _convert_scalar_indexer(self, key, kind=None):

pandas/core/indexes/period.py

+7-19
Original file line numberDiff line numberDiff line change
@@ -7,24 +7,21 @@
77
from pandas.core.dtypes.common import (
88
is_integer,
99
is_float,
10-
is_object_dtype,
1110
is_integer_dtype,
1211
is_float_dtype,
1312
is_scalar,
1413
is_datetime64_dtype,
15-
is_datetime64tz_dtype,
14+
is_datetime64_any_dtype,
1615
is_timedelta64_dtype,
1716
is_period_dtype,
1817
is_bool_dtype,
19-
is_categorical_dtype,
2018
pandas_dtype,
2119
_ensure_object)
2220
from pandas.core.dtypes.dtypes import PeriodDtype
2321
from pandas.core.dtypes.generic import ABCSeries
2422

2523
import pandas.tseries.frequencies as frequencies
2624
from pandas.tseries.frequencies import get_freq_code as _gfc
27-
from pandas.core.indexes.category import CategoricalIndex
2825
from pandas.core.indexes.datetimes import DatetimeIndex, Int64Index, Index
2926
from pandas.core.indexes.timedeltas import TimedeltaIndex
3027
from pandas.core.indexes.datetimelike import DatelikeOps, DatetimeIndexOpsMixin
@@ -506,23 +503,14 @@ def asof_locs(self, where, mask):
506503
@Appender(_index_shared_docs['astype'])
507504
def astype(self, dtype, copy=True, how='start'):
508505
dtype = pandas_dtype(dtype)
509-
if is_object_dtype(dtype):
510-
return self._box_values_as_index()
511-
elif is_integer_dtype(dtype):
512-
if copy:
513-
return self._int64index.copy()
514-
else:
515-
return self._int64index
516-
elif is_datetime64_dtype(dtype):
517-
return self.to_timestamp(how=how)
518-
elif is_datetime64tz_dtype(dtype):
519-
return self.to_timestamp(how=how).tz_localize(dtype.tz)
506+
if is_integer_dtype(dtype):
507+
return self._int64index.copy() if copy else self._int64index
508+
elif is_datetime64_any_dtype(dtype):
509+
tz = getattr(dtype, 'tz', None)
510+
return self.to_timestamp(how=how).tz_localize(tz)
520511
elif is_period_dtype(dtype):
521512
return self.asfreq(freq=dtype.freq)
522-
elif is_categorical_dtype(dtype):
523-
return CategoricalIndex(self.values, name=self.name, dtype=dtype,
524-
copy=copy)
525-
raise TypeError('Cannot cast PeriodIndex to dtype %s' % dtype)
513+
return super(PeriodIndex, self).astype(dtype, copy=copy)
526514

527515
@Substitution(klass='PeriodIndex')
528516
@Appender(_shared_docs['searchsorted'])

pandas/core/indexes/timedeltas.py

+7-22
Original file line numberDiff line numberDiff line change
@@ -4,23 +4,20 @@
44
import numpy as np
55
from pandas.core.dtypes.common import (
66
_TD_DTYPE,
7-
is_integer, is_float,
7+
is_integer,
8+
is_float,
89
is_bool_dtype,
910
is_list_like,
1011
is_scalar,
11-
is_integer_dtype,
12-
is_object_dtype,
1312
is_timedelta64_dtype,
1413
is_timedelta64_ns_dtype,
15-
is_categorical_dtype,
1614
pandas_dtype,
1715
_ensure_int64)
1816
from pandas.core.dtypes.missing import isna
1917
from pandas.core.dtypes.generic import ABCSeries
2018
from pandas.core.common import _maybe_box, _values_from_object
2119

2220
from pandas.core.indexes.base import Index
23-
from pandas.core.indexes.category import CategoricalIndex
2421
from pandas.core.indexes.numeric import Int64Index
2522
import pandas.compat as compat
2623
from pandas.compat import u
@@ -483,28 +480,16 @@ def to_pytimedelta(self):
483480
@Appender(_index_shared_docs['astype'])
484481
def astype(self, dtype, copy=True):
485482
dtype = pandas_dtype(dtype)
486-
487-
if is_object_dtype(dtype):
488-
return self._box_values_as_index()
489-
elif is_timedelta64_ns_dtype(dtype):
490-
if copy is True:
491-
return self.copy()
492-
return self
483+
if is_timedelta64_ns_dtype(dtype):
484+
return self.copy() if copy else self
493485
elif is_timedelta64_dtype(dtype):
494486
# return an index (essentially this is division)
495487
result = self.values.astype(dtype, copy=copy)
496488
if self.hasnans:
497-
return Index(self._maybe_mask_results(result,
498-
convert='float64'),
499-
name=self.name)
489+
values = self._maybe_mask_results(result, convert='float64')
490+
return Index(values, name=self.name)
500491
return Index(result.astype('i8'), name=self.name)
501-
elif is_integer_dtype(dtype):
502-
return Index(self.values.astype('i8', copy=copy), dtype='i8',
503-
name=self.name)
504-
elif is_categorical_dtype(dtype):
505-
return CategoricalIndex(self.values, name=self.name, dtype=dtype,
506-
copy=copy)
507-
raise TypeError('Cannot cast TimedeltaIndex to dtype %s' % dtype)
492+
return super(TimedeltaIndex, self).astype(dtype, copy=copy)
508493

509494
def union(self, other):
510495
"""

pandas/tests/indexes/datetimes/test_astype.py

+7-7
Original file line numberDiff line numberDiff line change
@@ -126,15 +126,15 @@ def test_astype_object(self):
126126
tm.assert_index_equal(casted, Index(exp_values, dtype=np.object_))
127127
assert casted.tolist() == exp_values
128128

129-
def test_astype_raises(self):
129+
@pytest.mark.parametrize('dtype', [
130+
float, 'timedelta64', 'timedelta64[ns]', 'datetime64',
131+
'datetime64[D]'])
132+
def test_astype_raises(self, dtype):
130133
# GH 13149, GH 13209
131134
idx = DatetimeIndex(['2016-05-16', 'NaT', NaT, np.NaN])
132-
133-
pytest.raises(TypeError, idx.astype, float)
134-
pytest.raises(TypeError, idx.astype, 'timedelta64')
135-
pytest.raises(TypeError, idx.astype, 'timedelta64[ns]')
136-
pytest.raises(TypeError, idx.astype, 'datetime64')
137-
pytest.raises(TypeError, idx.astype, 'datetime64[D]')
135+
msg = 'Cannot cast DatetimeIndex to dtype'
136+
with tm.assert_raises_regex(TypeError, msg):
137+
idx.astype(dtype)
138138

139139
def test_index_convert_to_datetime_array(self):
140140
def _check_rng(rng):

pandas/tests/indexes/period/test_period.py

+10-6
Original file line numberDiff line numberDiff line change
@@ -39,19 +39,23 @@ def test_astype_conversion(self):
3939
dtype=np.int64)
4040
tm.assert_index_equal(result, expected)
4141

42+
result = idx.astype(str)
43+
expected = Index(str(x) for x in idx)
44+
tm.assert_index_equal(result, expected)
45+
4246
idx = period_range('1990', '2009', freq='A')
4347
result = idx.astype('i8')
4448
tm.assert_index_equal(result, Index(idx.asi8))
4549
tm.assert_numpy_array_equal(result.values, idx.asi8)
4650

47-
def test_astype_raises(self):
51+
@pytest.mark.parametrize('dtype', [
52+
float, 'timedelta64', 'timedelta64[ns]'])
53+
def test_astype_raises(self, dtype):
4854
# GH 13149, GH 13209
4955
idx = PeriodIndex(['2016-05-16', 'NaT', NaT, np.NaN], freq='D')
50-
51-
pytest.raises(TypeError, idx.astype, str)
52-
pytest.raises(TypeError, idx.astype, float)
53-
pytest.raises(TypeError, idx.astype, 'timedelta64')
54-
pytest.raises(TypeError, idx.astype, 'timedelta64[ns]')
56+
msg = 'Cannot cast PeriodIndex to dtype'
57+
with tm.assert_raises_regex(TypeError, msg):
58+
idx.astype(dtype)
5559

5660
def test_pickle_compat_construction(self):
5761
pass

0 commit comments

Comments
 (0)