fix(DataFrame): to_dict("index") and typevar

Author: cmp0xff

pandas-stubs/core/frame.pyi

@@ -19,6 +19,7 @@ from typing import (
     Generic,
     Literal,
     NoReturn,
+    TypeVar,
     final,
     overload,
 )
@@ -165,6 +166,8 @@ from pandas._typing import (
 from pandas.io.formats.style import Styler
 from pandas.plotting import PlotAccessor
 
+_T_MUTABLE_MAPPING = TypeVar("_T_MUTABLE_MAPPING", bound=MutableMapping, covariant=True)
+
 class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
     @overload
     def __getitem__(self, idx: tuple[int, int]) -> Scalar: ...
@@ -396,9 +399,9 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         self,
         orient: Literal["records"],
         *,
-        into: MutableMapping | type[MutableMapping],
+        into: _T_MUTABLE_MAPPING | type[_T_MUTABLE_MAPPING],
         index: Literal[True] = ...,
-    ) -> list[MutableMapping[Hashable, Any]]: ...
+    ) -> list[_T_MUTABLE_MAPPING]: ...
     @overload
     def to_dict(
         self,
@@ -410,39 +413,55 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     @overload
     def to_dict(
         self,
-        orient: Literal["dict", "list", "series", "index"],
+        orient: Literal["index"],
+        *,
+        into: _T_MUTABLE_MAPPING | type[_T_MUTABLE_MAPPING],
+        index: Literal[True] = ...,
+    ) -> MutableMapping[Hashable, _T_MUTABLE_MAPPING]: ...
+    @overload
+    def to_dict(
+        self,
+        orient: Literal["index"],
+        *,
+        into: type[dict] = ...,
+        index: Literal[True] = ...,
+    ) -> dict[Hashable, dict[Hashable, Any]]: ...
+    @overload
+    def to_dict(
+        self,
+        orient: Literal["dict", "list", "series"],
         *,
-        into: MutableMapping | type[MutableMapping],
+        into: _T_MUTABLE_MAPPING | type[_T_MUTABLE_MAPPING],
         index: Literal[True] = ...,
-    ) -> MutableMapping[Hashable, Any]: ...
+    ) -> _T_MUTABLE_MAPPING: ...
     @overload
     def to_dict(
         self,
         orient: Literal["split", "tight"],
         *,
-        into: MutableMapping | type[MutableMapping],
+        into: _T_MUTABLE_MAPPING | type[_T_MUTABLE_MAPPING],
         index: bool = ...,
-    ) -> MutableMapping[Hashable, Any]: ...
+    ) -> _T_MUTABLE_MAPPING: ...
     @overload
     def to_dict(
         self,
-        orient: Literal["dict", "list", "series", "index"] = ...,
+        orient: Literal["dict", "list", "series"] = ...,
         *,
-        into: MutableMapping | type[MutableMapping],
+        into: _T_MUTABLE_MAPPING | type[_T_MUTABLE_MAPPING],
         index: Literal[True] = ...,
-    ) -> MutableMapping[Hashable, Any]: ...
+    ) -> _T_MUTABLE_MAPPING: ...
     @overload
     def to_dict(
         self,
         orient: Literal["split", "tight"] = ...,
         *,
-        into: MutableMapping | type[MutableMapping],
+        into: _T_MUTABLE_MAPPING | type[_T_MUTABLE_MAPPING],
         index: bool = ...,
-    ) -> MutableMapping[Hashable, Any]: ...
+    ) -> _T_MUTABLE_MAPPING: ...
     @overload
     def to_dict(
         self,
-        orient: Literal["dict", "list", "series", "index"] = ...,
+        orient: Literal["dict", "list", "series"] = ...,
         *,
         into: type[dict] = ...,
         index: Literal[True] = ...,

tests/test_frame.py

@@ -3638,33 +3638,83 @@ def test_to_records() -> None:
     )
 
 
-def test_to_dict() -> None:
+def test_to_dict_simple() -> None:
     check(assert_type(DF.to_dict(), dict[Hashable, Any]), dict)
     check(assert_type(DF.to_dict("split"), dict[Hashable, Any]), dict)
+    check(assert_type(DF.to_dict("records"), list[dict[Hashable, Any]]), list)
+
+    if TYPE_CHECKING_INVALID_USAGE:
+
+        def test(mapping: Mapping) -> None:  # pyright: ignore[reportUnusedFunction]
+            DF.to_dict(  # type: ignore[call-overload]
+                into=mapping  # pyright: ignore[reportArgumentType,reportCallIssue]
+            )
+
+
+def test_to_dict_into_defaultdict_any() -> None:
+    """Test DataFrame.to_dict with `into=defaultdict[Any, list]`"""
+
+    data = pd.DataFrame({("str", "rts"): [[1, 2, 4], [2, 3], [3]]})
+    target: defaultdict[Hashable, list[Any]] = defaultdict(list)
 
-    target: MutableMapping = defaultdict(list)
     check(
-        assert_type(DF.to_dict(into=target), MutableMapping[Hashable, Any]), defaultdict
+        assert_type(data.to_dict(into=target), defaultdict[Hashable, list[Any]]),
+        defaultdict,
     )
-    target = defaultdict(list)
     check(
-        assert_type(DF.to_dict("tight", into=target), MutableMapping[Hashable, Any]),
+        assert_type(
+            data.to_dict("index", into=target),
+            MutableMapping[Hashable, defaultdict[Hashable, list[Any]]],
+        ),
+        defaultdict,
+    )
+    check(
+        assert_type(
+            data.to_dict("tight", into=target), defaultdict[Hashable, list[Any]]
+        ),
         defaultdict,
     )
-    target = defaultdict(list)
-    check(assert_type(DF.to_dict("records"), list[dict[Hashable, Any]]), list)
     check(
         assert_type(
-            DF.to_dict("records", into=target), list[MutableMapping[Hashable, Any]]
+            data.to_dict("records", into=target), list[defaultdict[Hashable, list[Any]]]
         ),
         list,
     )
-    if TYPE_CHECKING_INVALID_USAGE:
 
-        def test(mapping: Mapping) -> None:  # pyright: ignore[reportUnusedFunction]
-            DF.to_dict(  # type: ignore[call-overload]
-                into=mapping  # pyright: ignore[reportArgumentType,reportCallIssue]
-            )
+
+def test_to_dict_into_defaultdict_typed() -> None:
+    """Test DataFrame.to_dict with `into=defaultdict[tuple[str, str], list[int]]`"""
+
+    data = pd.DataFrame({("str", "rts"): [[1, 2, 4], [2, 3], [3]]})
+    target: defaultdict[tuple[str, str], list[int]] = defaultdict(list)
+    target[("str", "rts")].append(1)
+
+    check(
+        assert_type(data.to_dict(into=target), defaultdict[tuple[str, str], list[int]]),
+        defaultdict,
+        tuple,
+    )
+    check(
+        assert_type(
+            data.to_dict("index", into=target),
+            MutableMapping[Hashable, defaultdict[tuple[str, str], list[int]]],
+        ),
+        defaultdict,
+    )
+    check(
+        assert_type(
+            data.to_dict("tight", into=target), defaultdict[tuple[str, str], list[int]]
+        ),
+        defaultdict,
+    )
+    check(
+        assert_type(
+            data.to_dict("records", into=target),
+            list[defaultdict[tuple[str, str], list[int]]],
+        ),
+        list,
+        defaultdict,
+    )
 
 
 def test_neg() -> None:
@@ -4111,7 +4161,10 @@ def test_to_dict_index() -> None:
         assert_type(df.to_dict(orient="series", index=True), dict[Hashable, Any]), dict
     )
     check(
-        assert_type(df.to_dict(orient="index", index=True), dict[Hashable, Any]), dict
+        assert_type(
+            df.to_dict(orient="index", index=True), dict[Hashable, dict[Hashable, Any]]
+        ),
+        dict,
     )
     check(
         assert_type(df.to_dict(orient="split", index=True), dict[Hashable, Any]), dict