TST: tests for nullable issues (#45167)

Author: jbrockmendel

pandas/tests/arrays/categorical/test_astype.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+from pandas import (
+    Categorical,
+    CategoricalDtype,
+    array,
+)
+import pandas._testing as tm
+
+
+class TestAstype:
+    def test_astype_str_int_categories_to_nullable_int(self):
+        # GH#39616
+        dtype = CategoricalDtype([str(i) for i in range(5)])
+        arr = Categorical.from_codes(np.random.randint(5, size=20), dtype=dtype)
+
+        res = arr.astype("Int64")
+        expected = array(arr.astype("int64"), dtype="Int64")
+        tm.assert_extension_array_equal(res, expected)

pandas/tests/arrays/categorical/test_constructors.py

@@ -511,6 +511,15 @@ def test_construction_with_null(self, klass, nulls_fixture):
 
         tm.assert_categorical_equal(result, expected)
 
+    def test_from_codes_nullable_int_categories(self, any_numeric_ea_dtype):
+        # GH#39649
+        cats = pd.array(range(5), dtype=any_numeric_ea_dtype)
+        codes = np.random.randint(5, size=3)
+        dtype = CategoricalDtype(cats)
+        arr = Categorical.from_codes(codes, dtype=dtype)
+        assert arr.categories.dtype == cats.dtype
+        tm.assert_index_equal(arr.categories, Index(cats))
+
     def test_from_codes_empty(self):
         cat = ["a", "b", "c"]
         result = Categorical.from_codes([], categories=cat)

pandas/tests/arrays/integer/test_arithmetic.py

@@ -55,7 +55,7 @@ def test_div(dtype):
 
 @pytest.mark.parametrize("zero, negative", [(0, False), (0.0, False), (-0.0, True)])
 def test_divide_by_zero(zero, negative):
-    # https://github.com/pandas-dev/pandas/issues/27398
+    # https://github.com/pandas-dev/pandas/issues/27398, GH#22793
     a = pd.array([0, 1, -1, None], dtype="Int64")
     result = a / zero
     expected = FloatingArray(

pandas/tests/indexing/test_loc.py

@@ -2714,6 +2714,26 @@ def test_loc_getitem_multiindex_tuple_level():
     assert result2 == 6
 
 
+def test_loc_getitem_nullable_index_with_duplicates():
+    # GH#34497
+    df = DataFrame(
+        data=np.array([[1, 2, 3, 4], [5, 6, 7, 8], [1, 2, np.nan, np.nan]]).T,
+        columns=["a", "b", "c"],
+        dtype="Int64",
+    )
+    df2 = df.set_index("c")
+    assert df2.index.dtype == "Int64"
+
+    res = df2.loc[1]
+    expected = Series([1, 5], index=df2.columns, dtype="Int64", name=1)
+    tm.assert_series_equal(res, expected)
+
+    # pd.NA and duplicates in an object-dtype Index
+    df2.index = df2.index.astype(object)
+    res = df2.loc[1]
+    tm.assert_series_equal(res, expected)
+
+
 class TestLocSeries:
     @pytest.mark.parametrize("val,expected", [(2 ** 63 - 1, 3), (2 ** 63, 4)])
     def test_loc_uint64(self, val, expected):