1from datetime import datetime
2
3import numpy as np
4import pytest
5
6import pandas as pd
7from pandas import DataFrame, Index, MultiIndex, RangeIndex, Series, date_range
8import pandas._testing as tm
9
10
11class TestResetIndex:
12    def test_reset_index_dti_round_trip(self):
13        dti = date_range(start="1/1/2001", end="6/1/2001", freq="D")._with_freq(None)
14        d1 = DataFrame({"v": np.random.rand(len(dti))}, index=dti)
15        d2 = d1.reset_index()
16        assert d2.dtypes[0] == np.dtype("M8[ns]")
17        d3 = d2.set_index("index")
18        tm.assert_frame_equal(d1, d3, check_names=False)
19
20        # GH#2329
21        stamp = datetime(2012, 11, 22)
22        df = DataFrame([[stamp, 12.1]], columns=["Date", "Value"])
23        df = df.set_index("Date")
24
25        assert df.index[0] == stamp
26        assert df.reset_index()["Date"][0] == stamp
27
28    def test_reset_index(self):
29        df = tm.makeDataFrame()[:5]
30        ser = df.stack()
31        ser.index.names = ["hash", "category"]
32
33        ser.name = "value"
34        df = ser.reset_index()
35        assert "value" in df
36
37        df = ser.reset_index(name="value2")
38        assert "value2" in df
39
40        # check inplace
41        s = ser.reset_index(drop=True)
42        s2 = ser
43        return_value = s2.reset_index(drop=True, inplace=True)
44        assert return_value is None
45        tm.assert_series_equal(s, s2)
46
47        # level
48        index = MultiIndex(
49            levels=[["bar"], ["one", "two", "three"], [0, 1]],
50            codes=[[0, 0, 0, 0, 0, 0], [0, 1, 2, 0, 1, 2], [0, 1, 0, 1, 0, 1]],
51        )
52        s = Series(np.random.randn(6), index=index)
53        rs = s.reset_index(level=1)
54        assert len(rs.columns) == 2
55
56        rs = s.reset_index(level=[0, 2], drop=True)
57        tm.assert_index_equal(rs.index, Index(index.get_level_values(1)))
58        assert isinstance(rs, Series)
59
60    def test_reset_index_name(self):
61        s = Series([1, 2, 3], index=Index(range(3), name="x"))
62        assert s.reset_index().index.name is None
63        assert s.reset_index(drop=True).index.name is None
64
65    def test_reset_index_level(self):
66        df = DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
67
68        for levels in ["A", "B"], [0, 1]:
69            # With MultiIndex
70            s = df.set_index(["A", "B"])["C"]
71
72            result = s.reset_index(level=levels[0])
73            tm.assert_frame_equal(result, df.set_index("B"))
74
75            result = s.reset_index(level=levels[:1])
76            tm.assert_frame_equal(result, df.set_index("B"))
77
78            result = s.reset_index(level=levels)
79            tm.assert_frame_equal(result, df)
80
81            result = df.set_index(["A", "B"]).reset_index(level=levels, drop=True)
82            tm.assert_frame_equal(result, df[["C"]])
83
84            with pytest.raises(KeyError, match="Level E "):
85                s.reset_index(level=["A", "E"])
86
87            # With single-level Index
88            s = df.set_index("A")["B"]
89
90            result = s.reset_index(level=levels[0])
91            tm.assert_frame_equal(result, df[["A", "B"]])
92
93            result = s.reset_index(level=levels[:1])
94            tm.assert_frame_equal(result, df[["A", "B"]])
95
96            result = s.reset_index(level=levels[0], drop=True)
97            tm.assert_series_equal(result, df["B"])
98
99            with pytest.raises(IndexError, match="Too many levels"):
100                s.reset_index(level=[0, 1, 2])
101
102        # Check that .reset_index([],drop=True) doesn't fail
103        result = Series(range(4)).reset_index([], drop=True)
104        expected = Series(range(4))
105        tm.assert_series_equal(result, expected)
106
107    def test_reset_index_range(self):
108        # GH 12071
109        s = Series(range(2), name="A", dtype="int64")
110        series_result = s.reset_index()
111        assert isinstance(series_result.index, RangeIndex)
112        series_expected = DataFrame(
113            [[0, 0], [1, 1]], columns=["index", "A"], index=RangeIndex(stop=2)
114        )
115        tm.assert_frame_equal(series_result, series_expected)
116
117    def test_reset_index_drop_errors(self):
118        #  GH 20925
119
120        # KeyError raised for series index when passed level name is missing
121        s = Series(range(4))
122        with pytest.raises(KeyError, match="does not match index name"):
123            s.reset_index("wrong", drop=True)
124        with pytest.raises(KeyError, match="does not match index name"):
125            s.reset_index("wrong")
126
127        # KeyError raised for series when level to be dropped is missing
128        s = Series(range(4), index=MultiIndex.from_product([[1, 2]] * 2))
129        with pytest.raises(KeyError, match="not found"):
130            s.reset_index("wrong", drop=True)
131
132    def test_reset_index_with_drop(self, series_with_multilevel_index):
133        ser = series_with_multilevel_index
134
135        deleveled = ser.reset_index()
136        assert isinstance(deleveled, DataFrame)
137        assert len(deleveled.columns) == len(ser.index.levels) + 1
138        assert deleveled.index.name == ser.index.name
139
140        deleveled = ser.reset_index(drop=True)
141        assert isinstance(deleveled, Series)
142        assert deleveled.index.name == ser.index.name
143
144
145@pytest.mark.parametrize(
146    "array, dtype",
147    [
148        (["a", "b"], object),
149        (
150            pd.period_range("12-1-2000", periods=2, freq="Q-DEC"),
151            pd.PeriodDtype(freq="Q-DEC"),
152        ),
153    ],
154)
155def test_reset_index_dtypes_on_empty_series_with_multiindex(array, dtype):
156    # GH 19602 - Preserve dtype on empty Series with MultiIndex
157    idx = MultiIndex.from_product([[0, 1], [0.5, 1.0], array])
158    result = Series(dtype=object, index=idx)[:0].reset_index().dtypes
159    expected = Series(
160        {"level_0": np.int64, "level_1": np.float64, "level_2": dtype, 0: object}
161    )
162    tm.assert_series_equal(result, expected)
163