Skip to content

TYP: Correct type annotation for to_dict. #55130

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Oct 2, 2023
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions pandas/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
Hashable,
Iterator,
Mapping,
MutableMapping,
Sequence,
)
from datetime import (
Expand Down Expand Up @@ -100,6 +101,7 @@
TypeGuard: Any = None

HashableT = TypeVar("HashableT", bound=Hashable)
MutableMappingT = TypeVar("MutableMappingT", bound=MutableMapping)

# array-like

Expand Down
22 changes: 12 additions & 10 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,7 @@
Level,
MergeHow,
MergeValidate,
MutableMappingT,
NaAction,
NaPosition,
NsmallestNlargestKeep,
Expand Down Expand Up @@ -1925,18 +1926,18 @@ def _create_data_for_split_and_tight_to_dict(
def to_dict(
self,
orient: Literal["dict", "list", "series", "split", "tight", "index"] = ...,
into: type[dict] = ...,
into: type[MutableMappingT] | MutableMappingT = ...,
index: bool = ...,
) -> dict:
) -> MutableMappingT:
...

@overload
def to_dict(
self,
orient: Literal["records"],
into: type[dict] = ...,
into: type[MutableMappingT] | MutableMappingT = ...,
index: bool = ...,
) -> list[dict]:
) -> list[MutableMappingT]:
...

@deprecate_nonkeyword_arguments(
Expand All @@ -1947,9 +1948,9 @@ def to_dict(
orient: Literal[
"dict", "list", "series", "split", "tight", "records", "index"
] = "dict",
into: type[dict] = dict,
into = dict,
index: bool = True,
) -> dict | list[dict]:
):
Copy link
Member

@twoertwein twoertwein Sep 15, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add the annotations for the non-overload definition? If it triggers a mypy error within the function, it is okay to add an ignore comment.

edit: having the ignore comment will function as a TODO enforced by mypy (when a future mypy allows TypeVars+default )

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jsspencer I pushed to your branch so that pd.DataFrame([]).to_dict("dict") works correctly

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should now work for all cases:

import pandas as pd
from typing import MutableMapping

class MyDict(dict): ...

def test(into: MutableMapping):
    # MutableMapping
    reveal_type(pd.DataFrame([]).to_dict("dict", into=into))
    reveal_type(pd.Series([]).to_dict(into=into))

    # dict
    reveal_type(pd.DataFrame([]).to_dict("dict"))
    reveal_type(pd.Series([]).to_dict())

    # MyDict
    reveal_type(pd.DataFrame([]).to_dict("dict", into=MyDict))
    reveal_type(pd.Series([]).to_dict(into=MyDict))

"""
Convert the DataFrame to a dictionary.

Expand Down Expand Up @@ -1977,7 +1978,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.
Expand All @@ -1991,9 +1992,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
--------
Expand Down
31 changes: 26 additions & 5 deletions pandas/core/methods/to_dict.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from typing import (
TYPE_CHECKING,
Literal,
overload,
)
import warnings

Expand All @@ -17,16 +18,36 @@

if TYPE_CHECKING:
from pandas import DataFrame
from pandas._typing import MutableMappingT


@overload
def to_dict(
df: DataFrame,
orient: Literal[
"dict", "list", "series", "split", "tight", "index"] = ...,
into: type[MutableMappingT] | MutableMappingT = ...,
index: bool = True,
) -> MutableMappingT:
...

@overload
def to_dict(
df: DataFrame,
orient: Literal["records"],
into: type[MutableMappingT] | MutableMappingT,
index: bool = True,
) -> list[MutableMappingT]:
...

def to_dict(
df: DataFrame,
orient: Literal[
"dict", "list", "series", "split", "tight", "records", "index"
] = "dict",
into: type[dict] = dict,
into = dict,
index: bool = True,
) -> dict | list[dict]:
):
"""
Convert the DataFrame to a dictionary.

Expand Down Expand Up @@ -54,7 +75,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.
Expand All @@ -69,8 +90,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(
Expand Down
27 changes: 21 additions & 6 deletions pandas/core/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,7 @@
IndexKeyFunc,
IndexLabel,
Level,
MutableMappingT,
NaPosition,
NumpySorter,
NumpyValueArrayLike,
Expand Down Expand Up @@ -1926,21 +1927,35 @@ 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:
...

@overload
def to_dict(
self,
into: MutableMappingT = ...,
) -> MutableMappingT:
...

def to_dict(self, into = dict):
"""
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
Expand Down