Parametrize some tests

Author: howsiwei

pandas/tests/frame/test_indexing.py

@@ -204,47 +204,29 @@ def test_setitem_list_of_tuples(self, float_frame):
         expected = Series(tuples, index=float_frame.index, name="tuples")
         tm.assert_series_equal(result, expected)
 
-    def test_setitem_list_all_missing_columns_scalar(self, float_frame):
-        # GH 26534
-        result = float_frame.copy()
-        result[["E", "F"]] = 1
-        expected = float_frame.copy()
-        # force the dtypes to be float as currently multcolumn assignment does not
-        # change column dtype from float to int even when it's being assigned an int
-        expected["E"] = 1.0
-        expected["F"] = 1.0
-        assert_frame_equal(result, expected)
-
-    def test_setitem_list_some_missing_columns_list(self, float_frame):
-        # GH 26534
-        result = float_frame.copy()
-        result[["A", "E"]] = [1, 2]
-        expected = float_frame.copy()
-        # force the dtypes to be float as currently multcolumn assignment does not
-        # change column dtype from float to int even when it's being assigned an int
-        expected["A"] = 1.0
-        expected["E"] = 2.0
-        assert_frame_equal(result, expected)
-
-    def test_setitem_list_some_missing_columns_dataframe(self, float_frame):
+    @pytest.mark.parametrize("columns", [["A", "E"], ["E", "F"]])
+    @pytest.mark.parametrize(
+        "box",
+        [
+            lambda x: 1,
+            lambda x: [1, 2],
+            lambda x: np.array([1, 2]),
+            lambda x: x[["B", "C"]],
+            lambda x: x[["B", "A"]].values,
+            lambda x: x[["A", "C"]].values.tolist(),
+        ],
+    )
+    def test_setitem_list_missing_columns(self, float_frame, columns, box):
         # GH 26534
         result = float_frame.copy()
-        result[["A", "E"]] = float_frame[["B", "C"]]
+        result[columns] = box(float_frame)
         expected = float_frame.copy()
-        expected["A"] = float_frame["B"]
-        expected["E"] = float_frame["C"]
+        for col in columns:
+            if col not in expected.columns:
+                expected[col] = np.nan
+        expected[columns] = box(float_frame)
         assert_frame_equal(result, expected)
 
-    def test_setitem_list_some_missing_columns_2dlist(self):
-        # GH 26534
-        result = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=["A", "B"])
-        result[["B", "C", "D"]] = [[7, 8, 9], [10, 11, 12], [13, 14, 15]]
-        expected = pd.DataFrame(
-            [[1, 7, 8, 9], [3, 10, 11, 12], [5, 13, 14, 15]],
-            columns=["A", "B", "C", "D"],
-        )
-        tm.assert_frame_equal(result, expected)
-
     def test_setitem_mulit_index(self):
         # GH7655, test that assigning to a sub-frame of a frame
         # with multi-index columns aligns both rows and columns

pandas/tests/indexing/test_loc.py

@@ -856,45 +856,29 @@ def test_loc_setitem_with_scalar_index(self, indexer, value):
 
         assert is_scalar(result) and result == "Z"
 
-    def test_loc_setitem_missing_columns_scalar_index_list_value(self):
-        # GH 26534
-        df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=["A", "B"])
-        df.loc[1, ["C", "D"]] = [7, 8]
-        expected = pd.DataFrame(
-            [[1, 2, np.nan, np.nan], [3, 4, 7, 8], [5, 6, np.nan, np.nan]],
-            columns=["A", "B", "C", "D"],
-        )
-        tm.assert_frame_equal(df, expected)
-
-    def test_loc_setitem_missing_columns_full_index_dataframe_value(self):
-        # GH 26534
-        df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=["A", "B"])
-        df2 = pd.DataFrame([[7, 8], [9, 10], [11, 12]], columns=["A", "C"])
-        df.loc[:, ["A", "C"]] = df2
-        expected = pd.DataFrame(
-            [[7, 2, 8], [9, 4, 10], [11, 6, 12]], columns=["A", "B", "C"]
-        )
-        tm.assert_frame_equal(df, expected)
-
-    def test_loc_setitem_missing_columns_list_index_scalar_value(self):
-        # GH 26534
-        df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=["A", "B"])
-        df.loc[[0, 2], ["B", "C", "D"]] = 9
-        expected = pd.DataFrame(
-            [[1, 9, 9, 9], [3, 4, np.nan, np.nan], [5, 9, 9, 9]],
-            columns=["A", "B", "C", "D"],
-        )
-        tm.assert_frame_equal(df, expected)
-
-    def test_loc_setitem_missing_columns_range_index_2dlist_value(self):
+    @pytest.mark.parametrize(
+        "index,box",
+        [
+            ((1, ["C", "D"]), [7, 8]),
+            (
+                (slice(None, None, None), ["A", "C"]),
+                pd.DataFrame([[7, 8], [9, 10], [11, 12]], columns=["A", "C"]),
+            ),
+            (([0, 2], ["B", "C", "D"]), 9),
+            ((slice(1, 3, None), ["B", "C", "D"]), [[7, 8, 9], [10, 11, 12]]),
+        ],
+    )
+    def test_loc_setitem_missing_columns(self, index, box):
         # GH 26534
         df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=["A", "B"])
-        df.loc[1:3, ["B", "C", "D"]] = [[7, 8, 9], [10, 11, 12]]
-        expected = pd.DataFrame(
-            [[1, 2, np.nan, np.nan], [3, 7, 8, 9], [5, 10, 11, 12]],
-            columns=["A", "B", "C", "D"],
-        )
-        tm.assert_frame_equal(df, expected)
+        result = df.copy()
+        result.loc[index] = box
+        expected = df.copy()
+        for col in index[1]:
+            if col not in expected.columns:
+                expected[col] = np.nan
+        expected.loc[index] = box
+        tm.assert_frame_equal(result, expected)
 
     def test_loc_coercion(self):