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