1from datetime import timedelta 2 3import numpy as np 4import pytest 5 6import pandas as pd 7from pandas import ( 8 Float64Index, 9 Index, 10 Int64Index, 11 NaT, 12 Timedelta, 13 TimedeltaIndex, 14 timedelta_range, 15) 16import pandas._testing as tm 17 18 19class TestTimedeltaIndex: 20 def test_astype_object(self): 21 idx = timedelta_range(start="1 days", periods=4, freq="D", name="idx") 22 expected_list = [ 23 Timedelta("1 days"), 24 Timedelta("2 days"), 25 Timedelta("3 days"), 26 Timedelta("4 days"), 27 ] 28 result = idx.astype(object) 29 expected = Index(expected_list, dtype=object, name="idx") 30 tm.assert_index_equal(result, expected) 31 assert idx.tolist() == expected_list 32 33 def test_astype_object_with_nat(self): 34 idx = TimedeltaIndex( 35 [timedelta(days=1), timedelta(days=2), NaT, timedelta(days=4)], name="idx" 36 ) 37 expected_list = [ 38 Timedelta("1 days"), 39 Timedelta("2 days"), 40 NaT, 41 Timedelta("4 days"), 42 ] 43 result = idx.astype(object) 44 expected = Index(expected_list, dtype=object, name="idx") 45 tm.assert_index_equal(result, expected) 46 assert idx.tolist() == expected_list 47 48 def test_astype(self): 49 # GH 13149, GH 13209 50 idx = TimedeltaIndex([1e14, "NaT", NaT, np.NaN], name="idx") 51 52 result = idx.astype(object) 53 expected = Index( 54 [Timedelta("1 days 03:46:40")] + [NaT] * 3, dtype=object, name="idx" 55 ) 56 tm.assert_index_equal(result, expected) 57 58 result = idx.astype(int) 59 expected = Int64Index( 60 [100000000000000] + [-9223372036854775808] * 3, dtype=np.int64, name="idx" 61 ) 62 tm.assert_index_equal(result, expected) 63 64 result = idx.astype(str) 65 expected = Index([str(x) for x in idx], name="idx") 66 tm.assert_index_equal(result, expected) 67 68 rng = timedelta_range("1 days", periods=10) 69 result = rng.astype("i8") 70 tm.assert_index_equal(result, Index(rng.asi8)) 71 tm.assert_numpy_array_equal(rng.asi8, result.values) 72 73 def test_astype_uint(self): 74 arr = timedelta_range("1H", periods=2) 75 expected = pd.UInt64Index( 76 np.array([3600000000000, 90000000000000], dtype="uint64") 77 ) 78 79 tm.assert_index_equal(arr.astype("uint64"), expected) 80 tm.assert_index_equal(arr.astype("uint32"), expected) 81 82 def test_astype_timedelta64(self): 83 # GH 13149, GH 13209 84 idx = TimedeltaIndex([1e14, "NaT", NaT, np.NaN]) 85 86 result = idx.astype("timedelta64") 87 expected = Float64Index([1e14] + [np.NaN] * 3, dtype="float64") 88 tm.assert_index_equal(result, expected) 89 90 result = idx.astype("timedelta64[ns]") 91 tm.assert_index_equal(result, idx) 92 assert result is not idx 93 94 result = idx.astype("timedelta64[ns]", copy=False) 95 tm.assert_index_equal(result, idx) 96 assert result is idx 97 98 @pytest.mark.parametrize("dtype", [float, "datetime64", "datetime64[ns]"]) 99 def test_astype_raises(self, dtype): 100 # GH 13149, GH 13209 101 idx = TimedeltaIndex([1e14, "NaT", NaT, np.NaN]) 102 msg = "Cannot cast TimedeltaArray to dtype" 103 with pytest.raises(TypeError, match=msg): 104 idx.astype(dtype) 105 106 def test_astype_category(self): 107 obj = timedelta_range("1H", periods=2, freq="H") 108 109 result = obj.astype("category") 110 expected = pd.CategoricalIndex([Timedelta("1H"), Timedelta("2H")]) 111 tm.assert_index_equal(result, expected) 112 113 result = obj._data.astype("category") 114 expected = expected.values 115 tm.assert_categorical_equal(result, expected) 116 117 def test_astype_array_fallback(self): 118 obj = timedelta_range("1H", periods=2) 119 result = obj.astype(bool) 120 expected = Index(np.array([True, True])) 121 tm.assert_index_equal(result, expected) 122 123 result = obj._data.astype(bool) 124 expected = np.array([True, True]) 125 tm.assert_numpy_array_equal(result, expected) 126