|
12 | 12 | Index,
|
13 | 13 | Interval,
|
14 | 14 | Series,
|
| 15 | + Timedelta, |
15 | 16 | Timestamp,
|
16 | 17 | )
|
17 | 18 | from pandas.api.types import CategoricalDtype as CDT
|
@@ -763,27 +764,74 @@ def test_map_with_dict_or_series(self):
|
763 | 764 | tm.assert_index_equal(expected, output)
|
764 | 765 |
|
765 | 766 | @pytest.mark.parametrize(
|
766 |
| - "idx_values", [[1, 2, 3], [-1, -2, -3], [1.5, 2.5, 3.5], [-1.5, -2.5, -3.5]] |
| 767 | + "idx_values", |
| 768 | + [ |
| 769 | + # python types |
| 770 | + [1, 2, 3], |
| 771 | + [-1, -2, -3], |
| 772 | + [1.5, 2.5, 3.5], |
| 773 | + [-1.5, -2.5, -3.5], |
| 774 | + # numpy int/uint |
| 775 | + *[ |
| 776 | + np.array([1, 2, 3], dtype=dtype) |
| 777 | + for dtype in [ |
| 778 | + np.int8, |
| 779 | + np.int16, |
| 780 | + np.int32, |
| 781 | + np.int64, |
| 782 | + np.uint8, |
| 783 | + np.uint16, |
| 784 | + np.uint32, |
| 785 | + np.uint64, |
| 786 | + ] |
| 787 | + ], |
| 788 | + # numpy floats |
| 789 | + *[ |
| 790 | + np.array([1.5, 2.5, 3.5], dtype=dtype) |
| 791 | + for dtype in (np.float16, np.float32, np.float64) |
| 792 | + ], |
| 793 | + # pandas scalars |
| 794 | + [Interval(1, 4), Interval(4, 6), Interval(6, 9)], |
| 795 | + [Timestamp(2019, 1, 1), Timestamp(2019, 2, 1), Timestamp(2019, 3, 1)], |
| 796 | + [Timedelta(1, "d"), Timedelta(2, "d"), Timedelta(3, "D")], |
| 797 | + # pandas Integer arrays |
| 798 | + *[ |
| 799 | + pd.array([1, 2, 3], dtype=dtype) |
| 800 | + for dtype in [ |
| 801 | + "Int8", |
| 802 | + "Int16", |
| 803 | + "Int32", |
| 804 | + "Int64", |
| 805 | + "UInt8", |
| 806 | + "UInt32", |
| 807 | + "UInt64", |
| 808 | + ] |
| 809 | + ], |
| 810 | + # other pandas arrays |
| 811 | + pd.IntervalIndex.from_breaks([1, 4, 6, 9]).array, |
| 812 | + pd.date_range("2019-01-01", periods=3).array, |
| 813 | + pd.timedelta_range(start="1d", periods=3).array, |
| 814 | + ], |
767 | 815 | )
|
768 | 816 | def test_loc_with_non_string_categories(self, idx_values, ordered_fixture):
|
769 | 817 | # GH-17569
|
770 | 818 | cat_idx = CategoricalIndex(idx_values, ordered=ordered_fixture)
|
771 |
| - cat = DataFrame({"A": ["foo", "bar", "baz"]}, index=cat_idx) |
772 |
| - # scalar |
773 |
| - result = cat.loc[idx_values[0]] |
| 819 | + df = DataFrame({"A": ["foo", "bar", "baz"]}, index=cat_idx) |
| 820 | + # scalar selection |
| 821 | + result = df.loc[idx_values[0]] |
774 | 822 | expected = Series(["foo"], index=["A"], name=idx_values[0])
|
775 | 823 | tm.assert_series_equal(result, expected)
|
776 |
| - # list |
777 |
| - result = cat.loc[idx_values[:2]] |
| 824 | + # list selection |
| 825 | + result = df.loc[idx_values[:2]] |
778 | 826 | expected = DataFrame(["foo", "bar"], index=cat_idx[:2], columns=["A"])
|
779 | 827 | tm.assert_frame_equal(result, expected)
|
780 | 828 | # scalar assignment
|
781 |
| - result = cat.copy() |
| 829 | + result = df.copy() |
782 | 830 | result.loc[idx_values[0]] = "qux"
|
783 | 831 | expected = DataFrame({"A": ["qux", "bar", "baz"]}, index=cat_idx)
|
784 | 832 | tm.assert_frame_equal(result, expected)
|
785 | 833 | # list assignment
|
786 |
| - result = cat.copy() |
| 834 | + result = df.copy() |
787 | 835 | result.loc[idx_values[:2], "A"] = ["qux", "qux2"]
|
788 | 836 | expected = DataFrame({"A": ["qux", "qux2", "baz"]}, index=cat_idx)
|
789 | 837 | tm.assert_frame_equal(result, expected)
|
0 commit comments