Unify _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, is_indexer_key=False)
+                self.loc._ensure_listlike_indexer(key)
                 indexer = self.loc._get_listlike_indexer(
                     key, axis=1, raise_missing=False
                 )[1]

pandas/core/indexing.py

@@ -160,6 +160,9 @@ def _slice(self, obj, axis: int, kind=None):
         return self.obj._slice(obj, axis=axis, kind=kind)
 
     def _get_setitem_indexer(self, key):
+        if self.name == "loc":
+            self._ensure_listlike_indexer(key)
+
         if self.axis is not None:
             return self._convert_tuple(key)
 
@@ -1728,7 +1731,7 @@ 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, is_indexer_key: bool):
+    def _ensure_listlike_indexer(self, key):
         """
         Ensure that a list-like of column labels are all present by adding them if
         they do not already exist.
@@ -1737,36 +1740,38 @@ def _ensure_listlike_indexer(self, key, is_indexer_key: bool):
         ----------
         key : _LocIndexer key or list-like of column labels
             Target labels.
-        is_indexer_key : bool
-            Whether key is a _LocIndexer key
         """
         column_axis = 1
-        if is_indexer_key:
+
+        # check if self.obj is at least 2-dimensional
+        if len(self.obj.shape) <= column_axis:
+            return
+
+        if isinstance(key, tuple):
+            # key is a _LocIndexer 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
+                # key indexes multiple columns
+                is_list_like_indexer(key[column_axis])
                 and not com.is_bool_indexer(key[column_axis])
             ):
                 return
+
+            # set key to the column part of key which is a list-like of column
+            # labels
             key = key[column_axis]
 
-        if not isinstance(self.obj._get_axis(column_axis), ABCMultiIndex) and all(
-            is_hashable(k) for k in key
+        if (
+            not isinstance(self.obj._get_axis(column_axis), ABCMultiIndex)
+            and is_list_like(key)
+            and all(is_hashable(k) for k in key)
         ):
             for k in key:
                 try:
                     self.obj[k]
                 except KeyError:
                     self.obj[k] = np.nan
 
-    def _get_setitem_indexer(self, key):
-        self._ensure_listlike_indexer(key, is_indexer_key=True)
-        return super()._get_setitem_indexer(key)
-
-
 class _iLocIndexer(_LocationIndexer):
     """
     Purely integer-location based indexing for selection by position.