update typing expectations for DataFrame.combine_first

Author: angela-tarantula

pandas/tests/frame/methods/test_combine_first.py

@@ -26,7 +26,12 @@ def test_combine_first_mixed(self):
         b = Series(range(2), index=range(5, 7))
         g = DataFrame({"A": a, "B": b})
 
-        exp = DataFrame({"A": list("abab"), "B": [0, 1, 0, 1]}, index=[0, 1, 5, 6])
+        exp = DataFrame(
+            {
+                "A": list("abab"),
+                "B": Series([0, 1, 0, 1], index=[0, 1, 5, 6], dtype="Int64"),
+            }
+        )
         combined = f.combine_first(g)
         tm.assert_frame_equal(combined, exp)
 
@@ -52,7 +57,7 @@ def test_combine_first(self, float_frame):
         combined = fcopy.combine_first(fcopy2)
 
         assert (combined["A"] == 1).all()
-        tm.assert_series_equal(combined["B"], fcopy["B"])
+        tm.assert_series_equal(combined["B"], fcopy["B"].astype("Float64"))
         tm.assert_series_equal(combined["C"], fcopy2["C"])
         tm.assert_series_equal(combined["D"], fcopy["D"])
 
@@ -118,12 +123,15 @@ def test_combine_first_same_as_in_update(self):
 
         other = DataFrame([[45, 45]], index=[0], columns=["A", "B"])
         result = df.combine_first(other)
-        tm.assert_frame_equal(result, df)
+        expected = df.copy()
+        expected["A"] = expected["A"].astype("Float64")
+        expected["B"] = expected["B"].astype("Float64")
+        tm.assert_frame_equal(result, expected)
 
         df.loc[0, "A"] = np.nan
         result = df.combine_first(other)
-        df.loc[0, "A"] = 45
-        tm.assert_frame_equal(result, df)
+        expected.loc[0, "A"] = 45
+        tm.assert_frame_equal(result, expected)
 
     def test_combine_first_doc_example(self):
         # doc example
@@ -202,21 +210,25 @@ def test_combine_first_align_nan(self):
 
         res = dfa.combine_first(dfb)
         exp = DataFrame(
-            {"a": [pd.Timestamp("2011-01-01"), pd.NaT], "b": [2, 5]},
+            {
+                "a": [pd.Timestamp("2011-01-01"), pd.NaT],
+                "b": Series([2, 5], dtype="Int64"),
+            },
             columns=["a", "b"],
         )
         tm.assert_frame_equal(res, exp)
         assert res["a"].dtype == "datetime64[s]"
-        # TODO: this must be int64
-        assert res["b"].dtype == "int64"
+        assert res["b"].dtype == "Int64"
 
         res = dfa.iloc[:0].combine_first(dfb)
-        exp = DataFrame({"a": [np.nan, np.nan], "b": [4, 5]}, columns=["a", "b"])
+        exp = DataFrame(
+            {"a": [np.nan, np.nan], "b": Series([4, 5], dtype="Int64")},
+            columns=["a", "b"],
+        )
         tm.assert_frame_equal(res, exp)
         # TODO: this must be datetime64
         assert res["a"].dtype == "float64"
-        # TODO: this must be int64
-        assert res["b"].dtype == "int64"
+        assert res["b"].dtype == "Int64"
 
     def test_combine_first_timezone(self, unit):
         # see gh-7630
@@ -366,21 +378,21 @@ def test_combine_first_int(self):
         df2 = DataFrame({"a": [1, 4]}, dtype="int64")
 
         result_12 = df1.combine_first(df2)
-        expected_12 = DataFrame({"a": [0, 1, 3, 5]})
+        expected_12 = DataFrame({"a": Series([0, 1, 3, 5], dtype="Int64")})
         tm.assert_frame_equal(result_12, expected_12)
 
         result_21 = df2.combine_first(df1)
-        expected_21 = DataFrame({"a": [1, 4, 3, 5]})
+        expected_21 = DataFrame({"a": Series([1, 4, 3, 5], dtype="Int64")})
         tm.assert_frame_equal(result_21, expected_21)
 
-    @pytest.mark.parametrize("val", [1, 1.0])
-    def test_combine_first_with_asymmetric_other(self, val):
+    @pytest.mark.parametrize("val,dtype", [(1, "Int64"), (1.0, "float64")])
+    def test_combine_first_with_asymmetric_other(self, val, dtype):
         # see gh-20699
         df1 = DataFrame({"isNum": [val]})
         df2 = DataFrame({"isBool": [True]})
 
         res = df1.combine_first(df2)
-        exp = DataFrame({"isNum": [val], "isBool": [True]})
+        exp = DataFrame({"isNum": Series([val], dtype=dtype), "isBool": [True]})
 
         tm.assert_frame_equal(res, exp)
 
@@ -472,9 +484,10 @@ def test_combine_first_with_nan_multiindex():
     expected = DataFrame(
         {
             "c": [np.nan, np.nan, 1, 1, 1, 1, 1, np.nan, 1, np.nan, 1],
-            "d": [1.0, 4.0, np.nan, 2.0, 5.0, np.nan, np.nan, 3.0, np.nan, 6.0, np.nan],
+            "d": [1, 4, np.nan, 2, 5, np.nan, np.nan, 3, np.nan, 6, np.nan],
         },
         index=mi_expected,
+        dtype="Int64",
     )
     tm.assert_frame_equal(res, expected)
 
@@ -492,10 +505,9 @@ def test_combine_preserve_dtypes():
     expected = DataFrame(
         {
             "A": ["a", "b", np.nan, np.nan],
-            "B": [0, 1, -1, 0],
+            "B": Series([0, 1, -1, 0], index=[0, 1, 5, 6], dtype="Int64"),
             "C": [np.nan, np.nan, "a", "b"],
-        },
-        index=[0, 1, 5, 6],
+        }
     )
     combined = df1.combine_first(df2)
     tm.assert_frame_equal(combined, expected)
@@ -515,12 +527,13 @@ def test_combine_first_duplicates_rows_for_nan_index_values():
 
     expected = DataFrame(
         {
-            "x": [9.0, 10.0, 11.0, np.nan],
-            "y": [12.0, 13.0, np.nan, 14.0],
+            "x": [9, 10, 11, np.nan],
+            "y": [12, 13, np.nan, 14],
         },
         index=MultiIndex.from_arrays(
             [[1, 2, 3, 4], [np.nan, 5, 6, 7]], names=["a", "b"]
         ),
+        dtype="Int64",
     )
     combined = df1.combine_first(df2)
     tm.assert_frame_equal(combined, expected)
@@ -531,7 +544,9 @@ def test_combine_first_int64_not_cast_to_float64():
     df_1 = DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
     df_2 = DataFrame({"A": [1, 20, 30], "B": [40, 50, 60], "C": [12, 34, 65]})
     result = df_1.combine_first(df_2)
-    expected = DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [12, 34, 65]})
+    expected = DataFrame(
+        {"A": [1, 2, 3], "B": [4, 5, 6], "C": [12, 34, 65]}, dtype="Int64"
+    )
     tm.assert_frame_equal(result, expected)
 
 
@@ -545,7 +560,7 @@ def test_midx_losing_dtype():
     expected_midx = MultiIndex.from_arrays(
         [[0, 0, 1, 1], [np.nan, np.nan, np.nan, np.nan]]
     )
-    expected = DataFrame({"a": [np.nan, 4, 3, 3]}, index=expected_midx)
+    expected = DataFrame({"a": [pd.NA, 4, 3, 3]}, index=expected_midx, dtype="Float64")
     tm.assert_frame_equal(result, expected)
 
 
@@ -563,5 +578,10 @@ def test_combine_first_preserve_column_order():
     df2 = DataFrame({"A": [5]}, index=[1])
 
     result = df1.combine_first(df2)
-    expected = DataFrame({"B": [1, 2, 3], "A": [4.0, 5.0, 6.0]})
+    expected = DataFrame(
+        {
+            "B": Series([1, 2, 3], dtype="Int64"),
+            "A": Series([4.0, 5.0, 6.0], dtype="Float64"),
+        }
+    )
     tm.assert_frame_equal(result, expected)