From 40a99dd0e11e935f59b9d4322b6d691a56fe8234 Mon Sep 17 00:00:00 2001 From: Fabian Gebhart Date: Sat, 21 Nov 2020 18:20:17 +0000 Subject: [PATCH 1/4] add test to verify column does not lose categorical type when using loc --- pandas/tests/dtypes/test_dtypes.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pandas/tests/dtypes/test_dtypes.py b/pandas/tests/dtypes/test_dtypes.py index a419cb0dded79..b4e1358288ef0 100644 --- a/pandas/tests/dtypes/test_dtypes.py +++ b/pandas/tests/dtypes/test_dtypes.py @@ -998,3 +998,10 @@ def test_period_dtype_compare_to_string(): dtype = PeriodDtype(freq="M") assert (dtype == "period[M]") is True assert (dtype != "period[M]") is False + + +def test_column_not_loses_categorical_when_using_loc(): + # GH16360 + df = pd.DataFrame({"A": [1]}) + df.loc[:, "B"] = Categorical(["a"]) + assert is_categorical_dtype(df["B"]) From 0452a41134e5302cae9f456e01ffa24001034715 Mon Sep 17 00:00:00 2001 From: Fabian Gebhart Date: Sun, 22 Nov 2020 09:07:35 +0000 Subject: [PATCH 2/4] move test and compare with expected df --- pandas/tests/dtypes/test_dtypes.py | 7 ------- pandas/tests/frame/indexing/test_categorical.py | 9 +++++++++ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/pandas/tests/dtypes/test_dtypes.py b/pandas/tests/dtypes/test_dtypes.py index b4e1358288ef0..a419cb0dded79 100644 --- a/pandas/tests/dtypes/test_dtypes.py +++ b/pandas/tests/dtypes/test_dtypes.py @@ -998,10 +998,3 @@ def test_period_dtype_compare_to_string(): dtype = PeriodDtype(freq="M") assert (dtype == "period[M]") is True assert (dtype != "period[M]") is False - - -def test_column_not_loses_categorical_when_using_loc(): - # GH16360 - df = pd.DataFrame({"A": [1]}) - df.loc[:, "B"] = Categorical(["a"]) - assert is_categorical_dtype(df["B"]) diff --git a/pandas/tests/frame/indexing/test_categorical.py b/pandas/tests/frame/indexing/test_categorical.py index 6137cadc93125..71ceb5cd8698c 100644 --- a/pandas/tests/frame/indexing/test_categorical.py +++ b/pandas/tests/frame/indexing/test_categorical.py @@ -386,3 +386,12 @@ def test_loc_indexing_preserves_index_category_dtype(self): result = df.loc[["a"]].index.levels[0] tm.assert_index_equal(result, expected) + + +def test_column_not_loses_categorical_when_using_loc(): + # GH16360 + result = DataFrame({"A": [1]}) + result.loc[:, "B"] = Categorical(["b"]) + expected = DataFrame({"A": [1], "B": ["b"]}) + expected["B"] = expected["B"].astype("category") + tm.assert_frame_equal(result, expected) From 353bce18a8bbc276a67fef1e9a9b37396bed75bc Mon Sep 17 00:00:00 2001 From: Fabian Gebhart Date: Sun, 22 Nov 2020 09:15:11 +0000 Subject: [PATCH 3/4] move again and rename function --- pandas/tests/frame/indexing/test_categorical.py | 9 --------- pandas/tests/indexing/test_loc.py | 9 +++++++++ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pandas/tests/frame/indexing/test_categorical.py b/pandas/tests/frame/indexing/test_categorical.py index 71ceb5cd8698c..6137cadc93125 100644 --- a/pandas/tests/frame/indexing/test_categorical.py +++ b/pandas/tests/frame/indexing/test_categorical.py @@ -386,12 +386,3 @@ def test_loc_indexing_preserves_index_category_dtype(self): result = df.loc[["a"]].index.levels[0] tm.assert_index_equal(result, expected) - - -def test_column_not_loses_categorical_when_using_loc(): - # GH16360 - result = DataFrame({"A": [1]}) - result.loc[:, "B"] = Categorical(["b"]) - expected = DataFrame({"A": [1], "B": ["b"]}) - expected["B"] = expected["B"].astype("category") - tm.assert_frame_equal(result, expected) diff --git a/pandas/tests/indexing/test_loc.py b/pandas/tests/indexing/test_loc.py index 28846bcf2f14d..de11165c3e7aa 100644 --- a/pandas/tests/indexing/test_loc.py +++ b/pandas/tests/indexing/test_loc.py @@ -12,6 +12,7 @@ import pandas as pd from pandas import ( + Categorical, CategoricalIndex, DataFrame, Index, @@ -1285,6 +1286,14 @@ def test_loc_setitem_datetime_keys_cast(self): expected = DataFrame({"one": [100.0, 200.0]}, index=[dt1, dt2]) tm.assert_frame_equal(df, expected) + def test_loc_setitem_categorical_column_retains_dtype(self): + # GH16360 + result = DataFrame({"A": [1]}) + result.loc[:, "B"] = Categorical(["b"]) + expected = DataFrame({"A": [1], "B": ["b"]}) + expected["B"] = expected["B"].astype("category") + tm.assert_frame_equal(result, expected) + class TestLocCallable: def test_frame_loc_getitem_callable(self): From ef9daf4ed8f8bc31b36ae0271f8243d5b2600a25 Mon Sep 17 00:00:00 2001 From: Fabian Gebhart Date: Mon, 23 Nov 2020 16:40:34 +0000 Subject: [PATCH 4/4] use ordered fixture for categorical --- pandas/tests/indexing/test_loc.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pandas/tests/indexing/test_loc.py b/pandas/tests/indexing/test_loc.py index de11165c3e7aa..9dbae874c1a93 100644 --- a/pandas/tests/indexing/test_loc.py +++ b/pandas/tests/indexing/test_loc.py @@ -1286,12 +1286,11 @@ def test_loc_setitem_datetime_keys_cast(self): expected = DataFrame({"one": [100.0, 200.0]}, index=[dt1, dt2]) tm.assert_frame_equal(df, expected) - def test_loc_setitem_categorical_column_retains_dtype(self): + def test_loc_setitem_categorical_column_retains_dtype(self, ordered): # GH16360 result = DataFrame({"A": [1]}) - result.loc[:, "B"] = Categorical(["b"]) - expected = DataFrame({"A": [1], "B": ["b"]}) - expected["B"] = expected["B"].astype("category") + result.loc[:, "B"] = Categorical(["b"], ordered=ordered) + expected = DataFrame({"A": [1], "B": Categorical(["b"], ordered=ordered)}) tm.assert_frame_equal(result, expected)