1import pytest
2
3from pandas import DataFrame, IndexSlice, MultiIndex, date_range
4import pandas._testing as tm
5
6
7@pytest.fixture
8def df():
9    #                        c1
10    # 2016-01-01 00:00:00 a   0
11    #                     b   1
12    #                     c   2
13    # 2016-01-01 12:00:00 a   3
14    #                     b   4
15    #                     c   5
16    # 2016-01-02 00:00:00 a   6
17    #                     b   7
18    #                     c   8
19    # 2016-01-02 12:00:00 a   9
20    #                     b  10
21    #                     c  11
22    # 2016-01-03 00:00:00 a  12
23    #                     b  13
24    #                     c  14
25    dr = date_range("2016-01-01", "2016-01-03", freq="12H")
26    abc = ["a", "b", "c"]
27    mi = MultiIndex.from_product([dr, abc])
28    frame = DataFrame({"c1": range(0, 15)}, index=mi)
29    return frame
30
31
32def test_partial_string_matching_single_index(df):
33    # partial string matching on a single index
34    for df_swap in [df.swaplevel(), df.swaplevel(0), df.swaplevel(0, 1)]:
35        df_swap = df_swap.sort_index()
36        just_a = df_swap.loc["a"]
37        result = just_a.loc["2016-01-01"]
38        expected = df.loc[IndexSlice[:, "a"], :].iloc[0:2]
39        expected.index = expected.index.droplevel(1)
40        tm.assert_frame_equal(result, expected)
41
42
43def test_partial_string_timestamp_multiindex(df):
44    # GH10331
45    df_swap = df.swaplevel(0, 1).sort_index()
46    SLC = IndexSlice
47
48    # indexing with IndexSlice
49    result = df.loc[SLC["2016-01-01":"2016-02-01", :], :]
50    expected = df
51    tm.assert_frame_equal(result, expected)
52
53    # match on secondary index
54    result = df_swap.loc[SLC[:, "2016-01-01":"2016-01-01"], :]
55    expected = df_swap.iloc[[0, 1, 5, 6, 10, 11]]
56    tm.assert_frame_equal(result, expected)
57
58    # partial string match on year only
59    result = df.loc["2016"]
60    expected = df
61    tm.assert_frame_equal(result, expected)
62
63    # partial string match on date
64    result = df.loc["2016-01-01"]
65    expected = df.iloc[0:6]
66    tm.assert_frame_equal(result, expected)
67
68    # partial string match on date and hour, from middle
69    result = df.loc["2016-01-02 12"]
70    expected = df.iloc[9:12]
71    tm.assert_frame_equal(result, expected)
72
73    # partial string match on secondary index
74    result = df_swap.loc[SLC[:, "2016-01-02"], :]
75    expected = df_swap.iloc[[2, 3, 7, 8, 12, 13]]
76    tm.assert_frame_equal(result, expected)
77
78    # tuple selector with partial string match on date
79    result = df.loc[("2016-01-01", "a"), :]
80    expected = df.iloc[[0, 3]]
81    tm.assert_frame_equal(result, expected)
82
83    # Slicing date on first level should break (of course)
84    with pytest.raises(KeyError, match="'2016-01-01'"):
85        df_swap.loc["2016-01-01"]
86
87
88def test_partial_string_timestamp_multiindex_str_key_raises(df):
89    # Even though this syntax works on a single index, this is somewhat
90    # ambiguous and we don't want to extend this behavior forward to work
91    # in multi-indexes. This would amount to selecting a scalar from a
92    # column.
93    with pytest.raises(KeyError, match="'2016-01-01'"):
94        df["2016-01-01"]
95
96
97def test_partial_string_timestamp_multiindex_daily_resolution(df):
98    # GH12685 (partial string with daily resolution or below)
99    result = df.loc[IndexSlice["2013-03":"2013-03", :], :]
100    expected = df.iloc[118:180]
101    tm.assert_frame_equal(result, expected)
102