- update tests for categorical values

Author: adrien pacifico

pandas/tests/reshape/concat/test_categorical.py

@@ -7,6 +7,7 @@
 import pandas as pd
 from pandas import (
     Categorical,
+    CategoricalIndex,
     DataFrame,
     Series,
 )
@@ -75,13 +76,13 @@ def test_concat_categoricalindex(self):
         # GH 16111, categories that aren't lexsorted
         categories = [9, 0, 1, 2, 3]
 
-        a = Series(1, index=pd.CategoricalIndex([9, 0], categories=categories))
-        b = Series(2, index=pd.CategoricalIndex([0, 1], categories=categories))
-        c = Series(3, index=pd.CategoricalIndex([1, 2], categories=categories))
+        a = Series(1, index=CategoricalIndex([9, 0], categories=categories))
+        b = Series(2, index=CategoricalIndex([0, 1], categories=categories))
+        c = Series(3, index=CategoricalIndex([1, 2], categories=categories))
 
         result = pd.concat([a, b, c], axis=1)
 
-        exp_idx = pd.CategoricalIndex([9, 0, 1, 2], categories=categories)
+        exp_idx = CategoricalIndex([9, 0, 1, 2], categories=categories)
         exp = DataFrame(
             {
                 0: [1, 1, np.nan, np.nan],
@@ -99,7 +100,7 @@ def test_categorical_concat_preserve(self):
         s = Series(list("abc"), dtype="category")
         s2 = Series(list("abd"), dtype="category")
 
-        exp = Series(list("abcabd"))
+        exp = Series(list("abcabd"), dtype="category")
         res = pd.concat([s, s2], ignore_index=True)
         tm.assert_series_equal(res, exp)
 
@@ -147,8 +148,8 @@ def test_categorical_index_preserver(self):
         result = pd.concat([df2, df3])
         expected = pd.concat(
             [
-                df2.set_axis(df2.index.astype(object), axis=0),
-                df3.set_axis(df3.index.astype(object), axis=0),
+                df2.set_axis(df2.index.astype("category"), axis=0),
+                df3.set_axis(df3.index.astype("category"), axis=0),
             ]
         )
         tm.assert_frame_equal(result, expected)
@@ -179,7 +180,8 @@ def test_concat_categorical_datetime(self):
 
         result = pd.concat([df1, df2])
         expected = DataFrame(
-            {"x": Series([datetime(2021, 1, 1), datetime(2021, 1, 2)])}
+            {"x": Series([datetime(2021, 1, 1), datetime(2021, 1, 2)])},
+            dtype="category",
         )
 
         tm.assert_equal(result, expected)
@@ -227,15 +229,17 @@ def test_categorical_index_upcast(self):
         b = DataFrame({"foo": [4, 3]}, index=Categorical(["baz", "bar"]))
 
         res = pd.concat([a, b])
-        exp = DataFrame({"foo": [1, 2, 4, 3]}, index=["foo", "bar", "baz", "bar"])
+        exp = DataFrame(
+            {"foo": [1, 2, 4, 3]}, index=Categorical(["foo", "bar", "baz", "bar"])
+        )
 
         tm.assert_equal(res, exp)
 
         a = Series([1, 2], index=Categorical(["foo", "bar"]))
         b = Series([4, 3], index=Categorical(["baz", "bar"]))
 
         res = pd.concat([a, b])
-        exp = Series([1, 2, 4, 3], index=["foo", "bar", "baz", "bar"])
+        exp = Series([1, 2, 4, 3], index=Categorical(["foo", "bar", "baz", "bar"]))
 
         tm.assert_equal(res, exp)
 
@@ -257,9 +261,9 @@ def test_categorical_missing_from_one_frame(self):
     def test_concat_categorical_same_categories_different_order(self):
         # https://github.com/pandas-dev/pandas/issues/24845
 
-        c1 = pd.CategoricalIndex(["a", "a"], categories=["a", "b"], ordered=False)
-        c2 = pd.CategoricalIndex(["b", "b"], categories=["b", "a"], ordered=False)
-        c3 = pd.CategoricalIndex(
+        c1 = CategoricalIndex(["a", "a"], categories=["a", "b"], ordered=False)
+        c2 = CategoricalIndex(["b", "b"], categories=["b", "a"], ordered=False)
+        c3 = CategoricalIndex(
             ["a", "a", "b", "b"], categories=["a", "b"], ordered=False
         )