Repository: pyjanitor-devs/pyjanitor (python)

Issue:
`summarize` 
# Brief Description

<!-- Please provide a brief description of what you'd like to propose. -->

I would like to propose a `summarize` function, similar to dplyr's `summarise` function and pandas' `agg` function, but for grouping operations, and more flexible

# Example API

```python
df.summarize(y='sum',n=lambda df: df.nth(1), by='x')

# summarize on multiple columns
df.summarize((['a','b','c'], 'sum'), by = 'x')

# replicate dplyr's across
# https://stackoverflow.com/q/63200530/7175713
# select_columns syntax can fit in nicely here
mtcars.summarize(("*t", "mean"), ("*p", "sum"), by='cyl')
```


Pull request:
[ENH] summarize function
Finally, as commits come in, don't forget to regularly rebase!
-->

# PR Description

Please describe the changes proposed in the pull request:

- Add `summarize` function


**This PR resolves #1225 **


Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8f4c663..2740ad2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 -  [ENH] Added support for pd.Series.select - Issue #1394 @samukweku
 -  [ENH] Added suport for janitor.mutate - Issue #1226 @samukweku
+-  [ENH] Added support for janitor.summarise - Issue #1225 @samukweku
 
 ## [v0.30.0] - 2024-12-04
 
diff --git a/janitor/functions/__init__.py b/janitor/functions/__init__.py
index dfd755d..03ed1bb 100644
--- a/janitor/functions/__init__.py
+++ b/janitor/functions/__init__.py
@@ -82,6 +82,7 @@ from .select import (
 from .shuffle import shuffle
 from .sort_column_value_order import sort_column_value_order
 from .sort_naturally import sort_naturally
+from .summarise import summarise
 from .take_first import take_first
 from .then import then
 from .to_datetime import to_datetime
@@ -160,6 +161,7 @@ __all__ = [
     "shuffle",
     "sort_column_value_order",
     "sort_naturally",
+    "summarise",
     "take_first",
     "then",
     "to_datetime",
diff --git a/janitor/functions/summarise.py b/janitor/functions/summarise.py
new file mode 100644
index 0000000..267f342
--- /dev/null
+++ b/janitor/functions/summarise.py
@@ -0,0 +1,265 @@
+"""Implementation of summarise."""
+
+from __future__ import annotations
+
+from functools import singledispatch
+from typing import Any
+
+import pandas as pd
+import pandas_flavor as pf
+from pandas.api.types import is_scalar
+from pandas.core.groupby.generic import DataFrameGroupBy
+
+from janitor.functions.select import get_index_labels
+
+
+@pf.register_dataframe_method
+def summarise(
+    df: pd.DataFrame,
+    *args: tuple[dict | tuple],
+    by: Any = None,
+) -> pd.DataFrame | pd.Series:
+    """
+
+    !!! info "New in version 0.31.0"
+
+    !!!note
+
+        Before reaching for `summarise`, try `pd.DataFrame.agg`.
+
+    summarise creates a new dataframe;
+    it returns one row for each combination of grouping columns.
+    If there are no grouping variables,
+    the output will have a single row
+    summarising all observations in the input.
+
+    The argument provided to *args* should be either a dictionary or a tuple.
+
+    - **dictionary argument**:
+    If the argument is a dictionary,
+    the value in the `{key:value}` pairing
+    should be either a string, a callable or a tuple.
+
+        - If the value in the dictionary
+        is a string or a callable,
+        the key of the dictionary
+        should be an existing column name.
+
+        !!!note
+
+            - If the value is a string,
+            the string should be a pandas string function,
+            e.g "sum", "mean", etc.
+
+        - If the value of the dictionary is a tuple,
+        it should be of length 2, and of the form
+        `(column_name, mutation_func)`,
+        where `column_name` should exist in the DataFrame,
+        and `mutation_func` should be either a string or a callable.
+
+        !!!note
+
+            - If `mutation_func` is a string,
+            the string should be a pandas string function,
+            e.g "sum", "mean", etc.
+
+        The key in the dictionary can be a new column name.
+
+    - **tuple argument**:
+    If the argument is a tuple, it should be of length 2,
+    and of the form
+    `(column_name, mutation_func)`,
+    where column_name should exist in the DataFrame,
+    and `mutation_func` should be either a string or a callable.
+
+        !!!note
+
+            - if `mutation_func` is a string,
+            the string should be a pandas string function,
+            e.g "sum", "mean", etc.
+
+        !!!note
+
+            - `column_name` can be anything supported by the
+            [`select`][janitor.functions.select.select] syntax;
+            as such multiple columns can be processed here -
+            they will be processed individually.
+
+    - **callable argument**:
+    If the argument is a callable, the callable is applied
+    on the DataFrame or GroupBy object.
+    The result from the callable should be a pandas Series
+    or DataFrame.
+
+
+    Aggregated columns cannot be reused in `summarise`.
+
+
+    `by` can be a `DataFrameGroupBy` object; it is assumed that
+    `by` was created from `df` - the onus is on the user to
+    ensure that, or the aggregations may yield incorrect results.
+
+    `by` accepts anything supported by `pd.DataFrame.groupby`.
+
+    Arguments supported in `pd.DataFrame.groupby`
+    can also be passed to `by` via a dictionary.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import janitor
+        >>> data = {'avg_jump': [3, 4, 1, 2, 3, 4],
+        ...         'avg_run': [3, 4, 1, 3, 2, 4],
+        ...         'combine_id': [100200, 100200,
+        ...                        101200, 101200,
+        ...                        102201, 103202]}
+        >>> df = pd.DataFrame(data)
+        >>> df
+           avg_jump  avg_run  combine_id
+        0         3        3      100200
+        1         4        4      100200
+        2         1        1      101200
+        3         2        3      101200
+        4         3        2      102201
+        5         4        4      103202
+
+        Aggregation via a callable:
+        >>> df.summarise(lambda df: df.sum(),by='combine_id')
+                    avg_jump  avg_run
+        combine_id
+        100200             7        7
+        101200             3        4
+        102201             3        2
+        103202             4        4
+
+        Aggregation via a tuple:
+        >>> df.summarise(("avg_run","mean"), by='combine_id')
+                    avg_run
+        combine_id
+        100200          3.5
+        101200          2.0
+        102201          2.0
+        103202          4.0
+
+        Aggregation via a dictionary:
+        >>> df.summarise({"avg_run":"mean"}, by='combine_id')
+                    avg_run
+        combine_id
+        100200          3.5
+        101200          2.0
+        102201          2.0
+        103202          4.0
+        >>> df.summarise({"avg_run_2":("avg_run","mean")}, by='combine_id')
+                    avg_run_2
+        combine_id
+        100200            3.5
+        101200            2.0
+        102201            2.0
+        103202            4.0
+
+    Args:
+        df: A pandas DataFrame.
+        args: Either a dictionary or a tuple.
+        by: Column(s) to group by.
+
+    Raises:
+        ValueError: If a tuple is passed and the length is not 2.
+
+    Returns:
+        A pandas DataFrame or Series with aggregated columns.
+
+    """  # noqa: E501
+
+    if by is not None:
+        # it is assumed that by is created from df
+        # onus is on user to ensure that
+        if isinstance(by, DataFrameGroupBy):
+            pass
+        elif isinstance(by, dict):
+            by = df.groupby(**by)
+        else:
+            if is_scalar(by):
+                by = [by]
+            by = df.groupby(by, sort=False, observed=True)
+    dictionary = {}
+    for arg in args:
+        aggregate = _mutator(arg, df=df, by=by)
+        dictionary.update(aggregate)
+    values = map(is_scalar, dictionary.values())
+    if all(values):
+        return pd.Series(dictionary)
+    return pd.concat(dictionary, axis=1, sort=False, copy=False)
+
+
+@singledispatch
+def _mutator(arg, df, by):
+    if not callable(arg):
+        raise NotImplementedError(
+            f"janitor.summarise is not supported for {type(arg)}"
+        )
+    if by is None:
+        val = df
+    else:
+        val = by
+    outcome = _process_maybe_callable(func=arg, obj=val)
+    if isinstance(outcome, pd.Series):
+        if not outcome.name:
+            raise ValueError("Ensure the pandas Series object has a name")
+        return {outcome.name: outcome}
+    # assumption: should return a DataFrame
+    outcome = {key: outcome[key] for key in outcome}
+    return outcome
+
+
+@_mutator.register(dict)
+def _(arg, df, by):
+    """Dispatch function for dictionary"""
+    if by is None:
+        val = df
+    else:
+        val = by
+
+    dictionary = {}
+    for column_name, mutator in arg.items():
+        if isinstance(mutator, tuple):
+            column, func = mutator
+            column = _process_within_dict(mutator=func, obj=val[column])
+        else:
+            column = _process_within_dict(
+                mutator=mutator, obj=val[column_name]
+            )
+        dictionary[column_name] = column
+    return dictionary
+
+
+@_mutator.register(tuple)
+def _(arg, df, by):
+    """Dispatch function for tuple"""
+    if len(arg) != 2:
+        raise ValueError("the tuple has to be a length of 2")
+    column_names, mutator = arg
+    column_names = get_index_labels(arg=[column_names], df=df, axis="columns")
+    mapping = {column_name: mutator for column_name in column_names}
+    return _mutator(mapping, df=df, by=by)
+
+
+def _process_maybe_callable(func: callable, obj):
+    """Function to handle callables"""
+    try:
+        column = obj.agg(func)
+    except:  # noqa: E722
+        column = func(obj)
+    return column
+
+
+def _process_maybe_string(func: str, obj):
+    """Function to handle pandas string functions"""
+    # treat as a pandas approved string function
+    # https://pandas.pydata.org/docs/user_guide/groupby.html#built-in-aggregation-methods
+    return obj.agg(func)
+
+
+def _process_within_dict(mutator, obj):
+    """Handle str/callables within a dictionary"""
+    if isinstance(mutator, str):
+        return _process_maybe_string(func=mutator, obj=obj)
+    return _process_maybe_callable(func=mutator, obj=obj)
diff --git a/mkdocs/api/functions.md b/mkdocs/api/functions.md
index fba17cd..e7d043e 100644
--- a/mkdocs/api/functions.md
+++ b/mkdocs/api/functions.md
@@ -56,6 +56,7 @@
         - shuffle
         - sort_column_value_order
         - sort_naturally
+        - summarise
         - take_first
         - then
         - to_datetime


Test patch:
diff --git a/tests/functions/test_summarise.py b/tests/functions/test_summarise.py
new file mode 100644
index 0000000..a768453
--- /dev/null
+++ b/tests/functions/test_summarise.py
@@ -0,0 +1,166 @@
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal, assert_series_equal
+
+
+@pytest.fixture
+def df_summarise():
+    data = {
+        "avg_jump": [3, 4, 1, 2, 3, 4],
+        "avg_run": [3, 4, 1, 3, 2, 4],
+        "combine_id": [100200, 100200, 101200, 101200, 102201, 103202],
+    }
+    return pd.DataFrame(data)
+
+
+def test_summarise_wrong_arg(df_summarise):
+    """
+    Raise if wrong arg is provided
+    """
+    with pytest.raises(
+        NotImplementedError, match="janitor.summarise is not supported for.+"
+    ):
+        df_summarise.summarise(1)
+
+
+def test_mutate_callable_series_unnamed(df_summarise):
+    """Test output for callable"""
+    with pytest.raises(
+        ValueError, match="Ensure the pandas Series object has a name"
+    ):
+        df_summarise.summarise(lambda df: df.mean())
+
+
+def test_mutate_callable_series(df_summarise):
+    """Test output for callable"""
+    expected = df_summarise.summarise(lambda df: df.mean().rename("mean"))
+    actual = df_summarise.mean().to_frame("mean")
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_by_callable_grp(df_summarise):
+    """Test output for a callable"""
+    grp = df_summarise.groupby("combine_id")
+    actual = df_summarise.summarise(lambda df: df.sum(), by=grp)
+    expected = grp.sum()
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_dict_df_str(df_summarise):
+    """Test output for a dictionary"""
+    actual = df_summarise.summarise({"avg_run": "mean"})
+    expected = df_summarise.agg({"avg_run": "mean"})
+    assert_series_equal(actual, expected)
+
+
+def test_summarise_dict_by_str(df_summarise):
+    """Test output for a dictionary"""
+    actual = df_summarise.summarise(
+        {"avg_run": "mean"}, by={"by": "combine_id"}
+    )
+    expected = df_summarise.groupby("combine_id").agg({"avg_run": "mean"})
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_dict_df_callable(df_summarise):
+    """Test output for a dictionary"""
+    actual = df_summarise.summarise({"avg_run": lambda df: df.sum()})
+    expected = df_summarise.agg({"avg_run": "sum"})
+    assert_series_equal(actual, expected)
+
+
+def test_summarise_dict_by_callable(df_summarise):
+    """Test output for a dictionary"""
+    actual = df_summarise.summarise(
+        {"avg_run": lambda df: df.sum()}, by="combine_id"
+    )
+    expected = df_summarise.groupby("combine_id").agg({"avg_run": "sum"})
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_dict_by_agg_callable(df_summarise):
+    """Test output for a dictionary"""
+    actual = df_summarise.summarise(
+        {"avg_run": lambda df: df.agg("sum")}, by="combine_id"
+    )
+    expected = df_summarise.groupby("combine_id").agg({"avg_run": "sum"})
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_dict_df_tuple(df_summarise):
+    """Test output for a dictionary"""
+    actual = (
+        df_summarise.summarise({"avg_run_sum": ("avg_run", "sum")})
+        .rename("avg_run")
+        .to_frame()
+    )
+    expected = df_summarise.agg(avg_run_sum=("avg_run", "sum"))
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_dict_by_tuple(df_summarise):
+    """Test output for a dictionary"""
+    actual = df_summarise.summarise(
+        {"avg_run_mean": ("avg_run", "mean")}, by="combine_id"
+    )
+    expected = df_summarise.groupby("combine_id").agg(
+        avg_run_mean=("avg_run", "mean")
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_tuple_count_not_eq_2(df_summarise):
+    """Raise error if length of tuple is not 2"""
+    with pytest.raises(ValueError, match="the tuple has to be a length of 2"):
+        df_summarise.summarise(("avg_run",))
+
+
+def test_summarise_df_tuple(df_summarise):
+    "Test output for a tuple"
+    actual = df_summarise.summarise(("avg_run", "sum"))
+    expected = df_summarise.agg({"avg_run": "sum"})
+    assert_series_equal(actual, expected)
+
+
+def test_summarise_by_tuple(df_summarise):
+    """Test output for a tuple"""
+    actual = df_summarise.summarise(("avg_run", "mean"), by="combine_id")
+    expected = df_summarise.groupby("combine_id").agg({"avg_run": "mean"})
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_tuple_df_callable(df_summarise):
+    """Test output for a tuple"""
+    actual = df_summarise.summarise(("avg_run", lambda df: df.sum()))
+    expected = df_summarise.agg({"avg_run": "sum"})
+    assert_series_equal(actual, expected)
+
+
+def test_summarise_tuple_by_callable(df_summarise):
+    """Test output for a tuple"""
+    actual = df_summarise.summarise(
+        ("avg_run", lambda df: df.sum()), by="combine_id"
+    )
+    expected = df_summarise.groupby("combine_id").agg({"avg_run": "sum"})
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_tuple_by_callable_dataframe(df_summarise):
+    """Test output for a tuple"""
+    actual = df_summarise.summarise(
+        ("avg_run", lambda df: df.agg(["sum", "mean"])), by="combine_id"
+    )
+    expected = df_summarise.groupby("combine_id").agg(
+        {"avg_run": ["sum", "mean"]}
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_summarise_tuple_grouped_object(df_summarise):
+    """Test output for a tuple"""
+    grp = df_summarise.groupby("combine_id")
+    actual = df_summarise.summarise(
+        ("avg_run", lambda df: df.agg(["sum", "mean"])), by=grp
+    )
+    expected = grp.agg({"avg_run": ["sum", "mean"]})
+    assert_frame_equal(actual, expected)