Repository: pandas-dev/pandas-stubs (language: python)

Issue:
`DataFrame.agg()` can return `Series` or `DataFrame` with `dict` argument
**Describe the bug**

There is no overload for `agg` returning a scalar, which it can per [the docs](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.agg.html) and:

**To Reproduce**

```py
df = pd.DataFrame([[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9],
                   [np.nan, np.nan, np.nan]],
                  columns=['A', 'B', 'C'])



# Expression of type "DataFrame" cannot be assigned to declared type "float64"
#  "DataFrame" is incompatible with "float64"Pylance[reportGeneralTypeIssues](https://github.com/microsoft/pyright/blob/main/docs/configuration.md#reportGeneralTypeIssues)
A_sum: np.float64 = df.agg({'A': 'sum'})


# But:
print(A_sum)

# A    12.0
# dtype: float64
```


**Please complete the following information:**
 - Linux
 - Ubuntu 20.04.6 LTS
 - python 3.11.4
 - pyright
 - `pandas-stubs == 2.1.4.231227


**Additional context**
Add any other context about the problem here.

`DataFrame.agg()` can return `Series` or `DataFrame` with `dict` argument
**Describe the bug**

There is no overload for `agg` returning a scalar, which it can per [the docs](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.agg.html) and:

**To Reproduce**

```py
df = pd.DataFrame([[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9],
                   [np.nan, np.nan, np.nan]],
                  columns=['A', 'B', 'C'])



# Expression of type "DataFrame" cannot be assigned to declared type "float64"
#  "DataFrame" is incompatible with "float64"Pylance[reportGeneralTypeIssues](https://github.com/microsoft/pyright/blob/main/docs/configuration.md#reportGeneralTypeIssues)
A_sum: np.float64 = df.agg({'A': 'sum'})


# But:
print(A_sum)

# A    12.0
# dtype: float64
```


**Please complete the following information:**
 - Linux
 - Ubuntu 20.04.6 LTS
 - python 3.11.4
 - pyright
 - `pandas-stubs == 2.1.4.231227


**Additional context**
Add any other context about the problem here.



Discussion / hints:
... while writing this I realized this is better anyway 😁 

```py
A_sum = df['A'].sum()
```
You're using a 7 month old version of `pandas-stubs`.  Please upgrade to the latest version (2.1.4.231227) and retest.  I tested with that release and pyright did not report an error in your code.

@Dr-Irv whoops, updated and seeing same thing,
I amended the snippet to actually generate an error (by adding a type annotation):

```py
A_sum: np.float64 = df.agg({'A': 'sum'})
```

```
Expression of type "DataFrame" cannot be assigned to declared type "float64"
  "DataFrame" is incompatible with "float64"Pylance[reportGeneralTypeIssues](https://github.com/microsoft/pyright/blob/main/docs/configuration.md#reportGeneralTypeIssues)
```
I'm not sure that we can fix this in terms of always returning the correct type from a static typing perspective.

```python
import pandas as pd
import numpy as np
df = pd.DataFrame([[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9],
                   [np.nan, np.nan, np.nan]],
                  columns=['A', 'B', 'C'])


A_sum = df.agg({'A': 'sum'})
print(type(A_sum))
print(A_sum)

S_sum = df.agg({'A': 'sum', "B": 'mean'})
print(type(S_sum))
print(S_sum)

T_sum = df.agg({"A": ['sum', 'mean'], "B": ["min", "max"]})
print(type(T_sum))
print(T_sum)
```
produces
```text
<class 'pandas.core.series.Series'>
A    12.0
dtype: float64
<class 'pandas.core.series.Series'>
A    12.0
B     5.0
dtype: float64
<class 'pandas.core.frame.DataFrame'>
         A    B
sum   12.0  NaN
mean   4.0  NaN
min    NaN  2.0
max    NaN  8.0
```
So `A_sum` and `S_sum` are `Series`, while `T_sum` is `DataFrame`, and the result is solely dependent on how many items are in the dictionary or in the dictionary items.  That can't be detected via static typing.

Having said that, we can change the type stub for `DataFrame.agg()` to return `Series | DataFrame` if a `dict` is passed.

I don't see how `agg` on a `DataFrame` can return a `Scalar`, so will change the title of this issue to reflect the issue with the `dict` argument.

> So `A_sum` and `S_sum` are `Series`, while `T_sum` is `DataFrame`, and the result is solely dependent on how many items are in the dictionary or in the dictionary items. That can't be detected via static typing.
> 
> Having said that, we can change the type stub for `DataFrame.agg()` to return `Series | DataFrame` if a `dict` is passed.

I think what we should look at in case of a dictionary is the *type of the dictionary values*, not the number of its items. In your example,  `A_sum` and `S_sum` are `Series` because they are the result of a single aggregation (indicated by the scalar values in the dictionary), while `T_sum` is `DataFrame` because it results from multiple aggregations (2 in this case).

If we look at the following example:
```pycon
>>> U_sum = df.agg({"A": ['sum', 'mean']})
>>> print(type(U_sum))
<class 'pandas.core.frame.DataFrame'>
>>> print(U_sum)
         A
sum   12.0
mean   4.0
```
We can see that even if the dictionary contains only one item, it returns a dataframe because it is performing two aggregations as indicated by the list `['sum', 'mean']`.
I can take a stab at this and add some tests after I get my groupby PR sorted out (unless someone else beats me to it).
> We can see that even if the dictionary contains only one item, it returns a dataframe because it is performing two aggregations as indicated by the list `['sum', 'mean']`.

Good point.

> I can take a stab at this and add some tests after I get my groupby PR sorted out (unless someone else beats me to it).

You can work on the 2 PR's independently.  I'll be reviewing the groupby PR shortly
> You can work on the 2 PR's independently. I'll be reviewing the groupby PR shortly

Of course. I wasn’t suggesting otherwise, it is just that I need to find some free time to work on it.

... while writing this I realized this is better anyway 😁 

```py
A_sum = df['A'].sum()
```
You're using a 7 month old version of `pandas-stubs`.  Please upgrade to the latest version (2.1.4.231227) and retest.  I tested with that release and pyright did not report an error in your code.

@Dr-Irv whoops, updated and seeing same thing,
I amended the snippet to actually generate an error (by adding a type annotation):

```py
A_sum: np.float64 = df.agg({'A': 'sum'})
```

```
Expression of type "DataFrame" cannot be assigned to declared type "float64"
  "DataFrame" is incompatible with "float64"Pylance[reportGeneralTypeIssues](https://github.com/microsoft/pyright/blob/main/docs/configuration.md#reportGeneralTypeIssues)
```
I'm not sure that we can fix this in terms of always returning the correct type from a static typing perspective.

```python
import pandas as pd
import numpy as np
df = pd.DataFrame([[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9],
                   [np.nan, np.nan, np.nan]],
                  columns=['A', 'B', 'C'])


A_sum = df.agg({'A': 'sum'})
print(type(A_sum))
print(A_sum)

S_sum = df.agg({'A': 'sum', "B": 'mean'})
print(type(S_sum))
print(S_sum)

T_sum = df.agg({"A": ['sum', 'mean'], "B": ["min", "max"]})
print(type(T_sum))
print(T_sum)
```
produces
```text
<class 'pandas.core.series.Series'>
A    12.0
dtype: float64
<class 'pandas.core.series.Series'>
A    12.0
B     5.0
dtype: float64
<class 'pandas.core.frame.DataFrame'>
         A    B
sum   12.0  NaN
mean   4.0  NaN
min    NaN  2.0
max    NaN  8.0
```
So `A_sum` and `S_sum` are `Series`, while `T_sum` is `DataFrame`, and the result is solely dependent on how many items are in the dictionary or in the dictionary items.  That can't be detected via static typing.

Having said that, we can change the type stub for `DataFrame.agg()` to return `Series | DataFrame` if a `dict` is passed.

I don't see how `agg` on a `DataFrame` can return a `Scalar`, so will change the title of this issue to reflect the issue with the `dict` argument.

> So `A_sum` and `S_sum` are `Series`, while `T_sum` is `DataFrame`, and the result is solely dependent on how many items are in the dictionary or in the dictionary items. That can't be detected via static typing.
> 
> Having said that, we can change the type stub for `DataFrame.agg()` to return `Series | DataFrame` if a `dict` is passed.

I think what we should look at in case of a dictionary is the *type of the dictionary values*, not the number of its items. In your example,  `A_sum` and `S_sum` are `Series` because they are the result of a single aggregation (indicated by the scalar values in the dictionary), while `T_sum` is `DataFrame` because it results from multiple aggregations (2 in this case).

If we look at the following example:
```pycon
>>> U_sum = df.agg({"A": ['sum', 'mean']})
>>> print(type(U_sum))
<class 'pandas.core.frame.DataFrame'>
>>> print(U_sum)
         A
sum   12.0
mean   4.0
```
We can see that even if the dictionary contains only one item, it returns a dataframe because it is performing two aggregations as indicated by the list `['sum', 'mean']`.
I can take a stab at this and add some tests after I get my groupby PR sorted out (unless someone else beats me to it).
> We can see that even if the dictionary contains only one item, it returns a dataframe because it is performing two aggregations as indicated by the list `['sum', 'mean']`.

Good point.

> I can take a stab at this and add some tests after I get my groupby PR sorted out (unless someone else beats me to it).

You can work on the 2 PR's independently.  I'll be reviewing the groupby PR shortly
> You can work on the 2 PR's independently. I'll be reviewing the groupby PR shortly

Of course. I wasn’t suggesting otherwise, it is just that I need to find some free time to work on it.


Fix patch:
diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
index 45f02ca2..ce9e5ba8 100644
--- a/pandas-stubs/core/frame.pyi
+++ b/pandas-stubs/core/frame.pyi
@@ -61,6 +61,7 @@ from pandas._typing import (
     S1,
     AggFuncTypeBase,
     AggFuncTypeDictFrame,
+    AggFuncTypeDictSeries,
     AggFuncTypeFrame,
     AnyArrayLike,
     ArrayLike,
@@ -1139,7 +1140,9 @@ class DataFrame(NDFrame, OpsMixin):
     ) -> DataFrame: ...
     def diff(self, periods: int = ..., axis: Axis = ...) -> DataFrame: ...
     @overload
-    def agg(self, func: AggFuncTypeBase, axis: Axis = ..., **kwargs) -> Series: ...
+    def agg(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+        self, func: AggFuncTypeBase | AggFuncTypeDictSeries, axis: Axis = ..., **kwargs
+    ) -> Series: ...
     @overload
     def agg(
         self,
@@ -1148,8 +1151,8 @@ class DataFrame(NDFrame, OpsMixin):
         **kwargs,
     ) -> DataFrame: ...
     @overload
-    def aggregate(
-        self, func: AggFuncTypeBase, axis: Axis = ..., **kwargs
+    def aggregate(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+        self, func: AggFuncTypeBase | AggFuncTypeDictSeries, axis: Axis = ..., **kwargs
     ) -> Series: ...
     @overload
     def aggregate(


Test patch:
diff --git a/tests/test_frame.py b/tests/test_frame.py
index cd4ace1f..d53d1a0d 100644
--- a/tests/test_frame.py
+++ b/tests/test_frame.py
@@ -1267,7 +1267,10 @@ def test_types_agg() -> None:
         assert_type(df.agg({"A": ["min", "max"], "B": "min"}), pd.DataFrame),
         pd.DataFrame,
     )
+    check(assert_type(df.agg({"A": ["mean"]}), pd.DataFrame), pd.DataFrame)
     check(assert_type(df.agg("mean", axis=1), pd.Series), pd.Series)
+    check(assert_type(df.agg({"A": "mean"}), pd.Series), pd.Series)
+    check(assert_type(df.agg({"A": "mean", "B": "sum"}), pd.Series), pd.Series)
 
 
 def test_types_aggregate() -> None:
@@ -1290,6 +1293,10 @@ def test_types_aggregate() -> None:
         assert_type(df.aggregate({"A": ["min", "max"], "B": "min"}), pd.DataFrame),
         pd.DataFrame,
     )
+    check(assert_type(df.aggregate({"A": ["mean"]}), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.aggregate("mean", axis=1), pd.Series), pd.Series)
+    check(assert_type(df.aggregate({"A": "mean"}), pd.Series), pd.Series)
+    check(assert_type(df.aggregate({"A": "mean", "B": "sum"}), pd.Series), pd.Series)
 
 
 def test_types_transform() -> None: