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