TYP: Correct type annotation for to_dict. (#55130)

* Correct type annotation for to_dict.

The `into` argument of DataFrame.to_dict and Series.to_dict can be
either a class or instance of a class of dict; this is covariant -
subclasses of dict can also be used. The argument was annotated as
`type[dict]` though, so type checkers marked passing initialized objects
(required for collections.defaultdict) as an incorrect argument type.

Fix by annotating `into` to take either a subclass of dict or an
initialized instance of a subclass of dict.

* Use generic MutableMapping type for to_dict method.

Unfortunately a generic type annotation with a default triggers an
existing mypy limitation (python/mypy#3737).
The current workaround is to use overloads and then not annotate the
implementation containing the default parameter; this still enables mypy
to deduce correct return types.

Two overloads are added for Series.to_dict, even though they could be
combined using a Union type, as at least two overloads are required for
a single method.

* Fix formatting

* return annotation for non-overload

* no keyword should return dict

* swap overload order to work for dict subclasses that are passed as keywords

* fix tests

---------

Co-authored-by: Torsten Wörtwein <[email protected]>

Author: jsspencer

pandas/_typing.py

@@ -4,6 +4,7 @@
     Hashable,
     Iterator,
     Mapping,
+    MutableMapping,
     Sequence,
 )
 from datetime import (
@@ -103,6 +104,7 @@
     TypeGuard: Any = None
 
 HashableT = TypeVar("HashableT", bound=Hashable)
+MutableMappingT = TypeVar("MutableMappingT", bound=MutableMapping)
 
 # array-like
 

pandas/core/frame.py

@@ -230,6 +230,7 @@
         Level,
         MergeHow,
         MergeValidate,
+        MutableMappingT,
         NaAction,
         NaPosition,
         NsmallestNlargestKeep,
@@ -1927,6 +1928,27 @@ def _create_data_for_split_and_tight_to_dict(
     def to_dict(
         self,
         orient: Literal["dict", "list", "series", "split", "tight", "index"] = ...,
+        *,
+        into: type[MutableMappingT] | MutableMappingT,
+        index: bool = ...,
+    ) -> MutableMappingT:
+        ...
+
+    @overload
+    def to_dict(
+        self,
+        orient: Literal["records"],
+        *,
+        into: type[MutableMappingT] | MutableMappingT,
+        index: bool = ...,
+    ) -> list[MutableMappingT]:
+        ...
+
+    @overload
+    def to_dict(
+        self,
+        orient: Literal["dict", "list", "series", "split", "tight", "index"] = ...,
+        *,
         into: type[dict] = ...,
         index: bool = ...,
     ) -> dict:
@@ -1936,11 +1958,14 @@ def to_dict(
     def to_dict(
         self,
         orient: Literal["records"],
+        *,
         into: type[dict] = ...,
         index: bool = ...,
     ) -> list[dict]:
         ...
 
+    # error: Incompatible default for argument "into" (default has type "type
+    # [dict[Any, Any]]", argument has type "type[MutableMappingT] | MutableMappingT")
     @deprecate_nonkeyword_arguments(
         version="3.0", allowed_args=["self", "orient"], name="to_dict"
     )
@@ -1949,9 +1974,10 @@ def to_dict(
         orient: Literal[
             "dict", "list", "series", "split", "tight", "records", "index"
         ] = "dict",
-        into: type[dict] = dict,
+        into: type[MutableMappingT]
+        | MutableMappingT = dict,  # type: ignore[assignment]
         index: bool = True,
-    ) -> dict | list[dict]:
+    ) -> MutableMappingT | list[MutableMappingT]:
         """
         Convert the DataFrame to a dictionary.
 
@@ -1979,7 +2005,7 @@ def to_dict(
                 'tight' as an allowed value for the ``orient`` argument
 
         into : class, default dict
-            The collections.abc.Mapping subclass used for all Mappings
+            The collections.abc.MutableMapping subclass used for all Mappings
             in the return value.  Can be the actual class or an empty
             instance of the mapping type you want.  If you want a
             collections.defaultdict, you must pass it initialized.
@@ -1993,9 +2019,10 @@ def to_dict(
 
         Returns
         -------
-        dict, list or collections.abc.Mapping
-            Return a collections.abc.Mapping object representing the DataFrame.
-            The resulting transformation depends on the `orient` parameter.
+        dict, list or collections.abc.MutableMapping
+            Return a collections.abc.MutableMapping object representing the
+            DataFrame. The resulting transformation depends on the `orient`
+            parameter.
 
         See Also
         --------
@@ -2054,7 +2081,7 @@ def to_dict(
         """
         from pandas.core.methods.to_dict import to_dict
 
-        return to_dict(self, orient, into, index)
+        return to_dict(self, orient, into=into, index=index)
 
     @deprecate_nonkeyword_arguments(
         version="3.0", allowed_args=["self", "destination_table"], name="to_gbq"

pandas/core/methods/to_dict.py

@@ -3,6 +3,7 @@
 from typing import (
     TYPE_CHECKING,
     Literal,
+    overload,
 )
 import warnings
 
@@ -16,17 +17,66 @@
 from pandas.core import common as com
 
 if TYPE_CHECKING:
+    from pandas._typing import MutableMappingT
+
     from pandas import DataFrame
 
 
+@overload
+def to_dict(
+    df: DataFrame,
+    orient: Literal["dict", "list", "series", "split", "tight", "index"] = ...,
+    *,
+    into: type[MutableMappingT] | MutableMappingT,
+    index: bool = ...,
+) -> MutableMappingT:
+    ...
+
+
+@overload
+def to_dict(
+    df: DataFrame,
+    orient: Literal["records"],
+    *,
+    into: type[MutableMappingT] | MutableMappingT,
+    index: bool = ...,
+) -> list[MutableMappingT]:
+    ...
+
+
+@overload
+def to_dict(
+    df: DataFrame,
+    orient: Literal["dict", "list", "series", "split", "tight", "index"] = ...,
+    *,
+    into: type[dict] = ...,
+    index: bool = ...,
+) -> dict:
+    ...
+
+
+@overload
+def to_dict(
+    df: DataFrame,
+    orient: Literal["records"],
+    *,
+    into: type[dict] = ...,
+    index: bool = ...,
+) -> list[dict]:
+    ...
+
+
+# error: Incompatible default for argument "into" (default has type "type[dict
+# [Any, Any]]", argument has type "type[MutableMappingT] | MutableMappingT")
 def to_dict(
     df: DataFrame,
     orient: Literal[
         "dict", "list", "series", "split", "tight", "records", "index"
     ] = "dict",
-    into: type[dict] = dict,
+    *,
+    into: type[MutableMappingT] | MutableMappingT = dict,  # type: ignore[assignment]
     index: bool = True,
-) -> dict | list[dict]:
+) -> MutableMappingT | list[MutableMappingT]:
     """
     Convert the DataFrame to a dictionary.
 
@@ -54,7 +104,7 @@ def to_dict(
             'tight' as an allowed value for the ``orient`` argument
 
     into : class, default dict
-        The collections.abc.Mapping subclass used for all Mappings
+        The collections.abc.MutableMapping subclass used for all Mappings
         in the return value.  Can be the actual class or an empty
         instance of the mapping type you want.  If you want a
         collections.defaultdict, you must pass it initialized.
@@ -69,8 +119,8 @@ def to_dict(
     Returns
     -------
     dict, list or collections.abc.Mapping
-        Return a collections.abc.Mapping object representing the DataFrame.
-        The resulting transformation depends on the `orient` parameter.
+        Return a collections.abc.MutableMapping object representing the
+        DataFrame. The resulting transformation depends on the `orient` parameter.
     """
     if not df.columns.is_unique:
         warnings.warn(
@@ -103,7 +153,7 @@ def to_dict(
     are_all_object_dtype_cols = len(box_native_indices) == len(df.dtypes)
 
     if orient == "dict":
-        return into_c((k, v.to_dict(into)) for k, v in df.items())
+        return into_c((k, v.to_dict(into=into)) for k, v in df.items())
 
     elif orient == "list":
         object_dtype_indices_as_set: set[int] = set(box_native_indices)

pandas/core/series.py

@@ -48,6 +48,7 @@
 from pandas.util._decorators import (
     Appender,
     Substitution,
+    deprecate_nonkeyword_arguments,
     doc,
 )
 from pandas.util._exceptions import find_stack_level
@@ -167,6 +168,7 @@
         IndexKeyFunc,
         IndexLabel,
         Level,
+        MutableMappingT,
         NaPosition,
         NumpySorter,
         NumpyValueArrayLike,
@@ -1922,21 +1924,40 @@ def keys(self) -> Index:
         """
         return self.index
 
-    def to_dict(self, into: type[dict] = dict) -> dict:
+    @overload
+    def to_dict(
+        self, *, into: type[MutableMappingT] | MutableMappingT
+    ) -> MutableMappingT:
+        ...
+
+    @overload
+    def to_dict(self, *, into: type[dict] = ...) -> dict:
+        ...
+
+    # error: Incompatible default for argument "into" (default has type "type[
+    # dict[Any, Any]]", argument has type "type[MutableMappingT] | MutableMappingT")
+    @deprecate_nonkeyword_arguments(
+        version="3.0", allowed_args=["self"], name="to_dict"
+    )
+    def to_dict(
+        self,
+        into: type[MutableMappingT]
+        | MutableMappingT = dict,  # type: ignore[assignment]
+    ) -> MutableMappingT:
         """
         Convert Series to {label -> value} dict or dict-like object.
 
         Parameters
         ----------
         into : class, default dict
-            The collections.abc.Mapping subclass to use as the return
-            object. Can be the actual class or an empty
-            instance of the mapping type you want.  If you want a
-            collections.defaultdict, you must pass it initialized.
+            The collections.abc.MutableMapping subclass to use as the return
+            object. Can be the actual class or an empty instance of the mapping
+            type you want.  If you want a collections.defaultdict, you must
+            pass it initialized.
 
         Returns
         -------
-        collections.abc.Mapping
+        collections.abc.MutableMapping
             Key-value representation of Series.
 
         Examples
@@ -1945,10 +1966,10 @@ def to_dict(self, into: type[dict] = dict) -> dict:
         >>> s.to_dict()
         {0: 1, 1: 2, 2: 3, 3: 4}
         >>> from collections import OrderedDict, defaultdict
-        >>> s.to_dict(OrderedDict)
+        >>> s.to_dict(into=OrderedDict)
         OrderedDict([(0, 1), (1, 2), (2, 3), (3, 4)])
         >>> dd = defaultdict(list)
-        >>> s.to_dict(dd)
+        >>> s.to_dict(into=dd)
         defaultdict(<class 'list'>, {0: 1, 1: 2, 2: 3, 3: 4})
         """
         # GH16122

pandas/tests/series/methods/test_to_dict.py

@@ -13,12 +13,12 @@ class TestSeriesToDict:
     )
     def test_to_dict(self, mapping, datetime_series):
         # GH#16122
-        result = Series(datetime_series.to_dict(mapping), name="ts")
+        result = Series(datetime_series.to_dict(into=mapping), name="ts")
         expected = datetime_series.copy()
         expected.index = expected.index._with_freq(None)
         tm.assert_series_equal(result, expected)
 
-        from_method = Series(datetime_series.to_dict(collections.Counter))
+        from_method = Series(datetime_series.to_dict(into=collections.Counter))
         from_constructor = Series(collections.Counter(datetime_series.items()))
         tm.assert_series_equal(from_method, from_constructor)