fix Dataframe.__setitem__ with Hashable (#1535)

* StrLike -> Hashable

* swap overload order

* write test

* reinstate test

* list -> Iterable, remove non-intended test case

* sync

Author: MarcoGorelli

pandas-stubs/core/frame.pyi

@@ -283,9 +283,7 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
     @overload
     def __setitem__(
         self,
-        key: (
-            MaskType | StrLike | _IndexSliceTuple | list[ScalarT] | IndexingInt | slice
-        ),
+        key: tuple[_IndexSliceTuple, Hashable],
         value: (
             Scalar
             | NAType
@@ -294,15 +292,21 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
             | IndexOpsMixin
             | Sequence[Scalar]
             | Sequence[Sequence[Scalar]]
-            | DataFrame
             | Mapping[Hashable, Scalar | NAType | NaTType]
             | None
         ),
     ) -> None: ...
     @overload
     def __setitem__(
         self,
-        key: tuple[_IndexSliceTuple, Hashable],
+        key: (
+            MaskType
+            | Hashable
+            | _IndexSliceTuple
+            | Iterable[Scalar]
+            | IndexingInt
+            | slice
+        ),
         value: (
             Scalar
             | NAType
@@ -311,6 +315,7 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
             | IndexOpsMixin
             | Sequence[Scalar]
             | Sequence[Sequence[Scalar]]
+            | DataFrame
             | Mapping[Hashable, Scalar | NAType | NaTType]
             | None
         ),
@@ -848,9 +853,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     @overload
     def __setitem__(
         self,
-        idx: (
-            MaskType | StrLike | _IndexSliceTuple | list[ScalarT] | IndexingInt | slice
-        ),
+        idx: tuple[_IndexSliceTuple, Hashable],
         value: (
             Scalar
             | NAType
@@ -859,15 +862,21 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
             | IndexOpsMixin
             | Sequence[Scalar]
             | Sequence[Sequence[Scalar]]
-            | DataFrame
             | Mapping[Hashable, Scalar | NAType | NaTType]
             | None
         ),
     ) -> None: ...
     @overload
     def __setitem__(
         self,
-        idx: tuple[_IndexSliceTuple, Hashable],
+        idx: (
+            MaskType
+            | Hashable
+            | _IndexSliceTuple
+            | Iterable[Scalar]
+            | IndexingInt
+            | slice
+        ),
         value: (
             Scalar
             | NAType
@@ -876,6 +885,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
             | IndexOpsMixin
             | Sequence[Scalar]
             | Sequence[Sequence[Scalar]]
+            | DataFrame
             | Mapping[Hashable, Scalar | NAType | NaTType]
             | None
         ),

tests/test_frame.py

@@ -290,11 +290,16 @@ def test_slice_setitem() -> None:
 
 def test_types_setitem() -> None:
     df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4], 5: [6, 7]})
+    h = cast(Hashable, "col1")
     i = pd.Index(["col1", "col2"])
     s = pd.Series(["col1", "col2"])
     a = np.array(["col1", "col2"])
     df["col1"] = [1, 2]
     df[5] = [5, 6]
+    df[h] = [5, 6]
+    df.loc[:, h] = [5, 6]
+    df.loc[:, UserList([h])] = [[5], [6]]
+    df.loc[:, iter([h])] = [[5], [6]]
     df[["col1", "col2"]] = [[1, 2], [3, 4]]
     df[s] = [5, 6]
     df.loc[:, s] = [5, 6]