ENH: add dict to return type of func argument of DataFrame#apply

[skatsuta]

• Closes #xxxx (Replace xxxx with the Github issue number)
• Tests added: Please use assert_type() to assert the type of any return value

Problem

It is perfectly valid to pass a function with a return type of dict as the func argument to DataFrame#apply when result_type="expand", but the current type stub does not support it.

For example, the following code

def returns_dict(x: pd.Series) -> dict[str, int]:
    return {'x': 1, 'y': 2}

df.apply(returns_dict, axis=1, result_type="expand")

results in the following errors when type-checked by Pyright:

error: Argument of type "(x: Series[Unknown]) -> dict[str, int]" cannot be assigned to parameter "f" of type "(...) -> Series[Unknown]" in function "apply"
    Type "(x: Series[Unknown]) -> dict[str, int]" cannot be assigned to type "(...) -> Series[Unknown]"
      Function return type "dict[str, int]" is incompatible with type "Series[Unknown]"
        "dict[str, int]" is incompatible with "Series[Unknown]" (reportGeneralTypeIssues)
error: Argument of type "Literal['expand']" cannot be assigned to parameter "result_type" of type "Literal['reduce']" in function "apply"
    "Literal['expand']" cannot be assigned to type "Literal['reduce']" (reportGeneralTypeIssues)

Solution

Add dict to the return type of func argument of DataFrame#apply when result_type="expand".

Also, add tests for passing a function that returns a dict to DataFrame#apply in the cases where axis is 0 and 1, respectively.

[gandhis1]

I recently modified these apply overloads and the number of permutations of overloads is a lot higher now. So any changes here need to be very careful to make sure they apply new data types everywhere they need to be present.

The changes you made look good to me, I only have one question. During the last round of changes to apply() annotations, we decided that even if a particular combination of function return and return_type did not make logical sense, if pandas allowed it without errors, then we would continue to permit it in the annotations.

I cannot test right now, but does dict work for any other result_type value (including the default of None)?

[Dr-Irv]

As suggested by @gandhis1 here: #470 (comment)

you should modify all of the overloads to return dict, and create tests for all of the various combinations of axis and result_type as is already present in the tests.

If some combination of axis and result_type doesn't work with dict as the return type, then you leave dict out as a possibility for the return type of the Callable in the f argument.

[skatsuta]

@gandhis1 You are the one who added those comprehensive overloads! Thanks for the hard work!
Will check other combinations and add support for them if they also work!

@Dr-Irv Thanks for your suggestion! Will add support for func returning dict to other overloads and also additional comprehensive test cases!

[skatsuta]

@Dr-Irv I've added Mapping return type to all possible overloads, and comprehensive test cases for all axis and result_type parameter combinations, which are 2 x 4 = 8 assertions in total (2 axes and 4 result types).

I've also replaced Literal[None] with just None, and reformatted and reordered some existing assertions for readability and clarity, but none of them are removed.

Would appreciate if you could take a look when you have time, thanks!

[skatsuta]

These assertions are moved to the bottom so that test cases for result_type="broadcast" are grouped together.

[skatsuta]

This does not necessarily have to be Literal[None], just None is OK, and so are others.

[Dr-Irv]

thanks @skatsuta