Put all checks in _ensure_listlike_indexer

Author: howsiwei

pandas/core/frame.py

@@ -2825,7 +2825,7 @@ def _setitem_array(self, key, value):
                 for k1, k2 in zip(key, value.columns):
                     self[k1] = value[k2]
             else:
-                self.loc._ensure_listlike_indexer(key, axis=1)
+                self.loc._ensure_listlike_indexer(key, is_indexer_key=False)
                 indexer = self.loc._get_listlike_indexer(
                     key, axis=1, raise_missing=False
                 )[1]

pandas/core/indexing.py

@@ -1728,19 +1728,32 @@ def _getitem_axis(self, key, axis: int):
         self._validate_key(key, axis)
         return self._get_label(key, axis=axis)
 
-    def _ensure_listlike_indexer(self, key: Iterable, axis: int):
+    def _ensure_listlike_indexer(self, key, is_indexer_key: bool):
         """
-        Ensure that a list-like of keys are all present by adding them if they
-        do not already exist.
+        Ensure that a list-like of column labels are all present by adding them if
+        they do not already exist.
 
         Parameters
         ----------
-        key : list-like
-            Target labels
-        axis : int
-            Dimension on which the indexing is being made
+        key : _LocIndexer key or list-like of column labels
+            Target labels.
+        is_indexer_key : bool
+            Whether key is a _LocIndexer key
         """
-        if not isinstance(self.obj._get_axis(axis), ABCMultiIndex) and all(
+        column_axis = 1
+        if is_indexer_key:
+            if not (
+                isinstance(key, tuple)
+                and len(key) >= 2  # key is at least 2-dimensional
+                and is_list_like_indexer(
+                    key[column_axis]
+                )  # key indexes multiple columns
+                and not com.is_bool_indexer(key[column_axis])
+            ):
+                return
+            key = key[column_axis]
+
+        if not isinstance(self.obj._get_axis(column_axis), ABCMultiIndex) and all(
             is_hashable(k) for k in key
         ):
             for k in key:
@@ -1750,13 +1763,7 @@ def _ensure_listlike_indexer(self, key: Iterable, axis: int):
                     self.obj[k] = np.nan
 
     def _get_setitem_indexer(self, key):
-        if (
-            isinstance(key, tuple)
-            and len(key) >= 2  # key is at least 2-dimensional
-            and is_list_like_indexer(key[1])  # key indexes multiple columns
-            and not com.is_bool_indexer(key[1])
-        ):
-            self._ensure_listlike_indexer(key[1], axis=1)
+        self._ensure_listlike_indexer(key, is_indexer_key=True)
         return super()._get_setitem_indexer(key)