Repository: pyjanitor-devs/pyjanitor

Issue:
suggestion - Add `names_expand` to `pivot_wider`
# Brief Description

I would like to propose adding a `names_expand` parameter to `pivot_wider` - this expands the columns to all possible combinations (useful for categorical columns).  ``df.pivot_wider(index, names_from, values_from, names_expand = True/False)``



# Example API

```python
# adapted from https://www.tidyverse.org/blog/2022/02/tidyr-1-2-0/
import pandas as pd
weekdays = ("Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun")
daily = {'day': pd.Categorical(("Tue", "Thu", "Fri", "Mon"), categories = weekdays),
              'value' : (2,3,1,5)}

daily = pd.DataFrame(daily, index = [0,0,0,0])

daily
   day  value
0  Tue      2
0  Thu      3
0  Fri      1
0  Mon      5

 daily.dtypes
day      category
value       int64
dtype: object

 daily.day
0    Tue
0    Thu
0    Fri
0    Mon
Name: day, dtype: category
Categories (7, object): ['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun']


# pivot only shows the existing values from the `day` column
daily.pivot_wider(index = None, names_from = 'day', values_from = 'value')
   Tue  Thu  Fri  Mon
0    2    3    1    5

# one efficient option is with reindexing on the columns
# which is what `names_expand` will use under the hood
daily.pivot_wider(index = None, names_from = 'day', values_from = 'value').reindex(columns = weekdays)

   Mon  Tue  Wed  Thu  Fri  Sat  Sun
0    5    2  NaN    3    1  NaN  NaN

# note however that we had to grab the weekdays to do this
# instead of using the categories within `names_from`
# if there were other columns, then we'd have to set_index, reindex, then reset_index (delaying the reset_index - doing the operations once before resetting would be ideal)
# and if there were multiple `names_from`, reindexing on one level won't suffice.
```



Discussion:
I'm in favour of this idea! 

Resolving patch:
diff --git a/janitor/functions/pivot.py b/janitor/functions/pivot.py
--- a/janitor/functions/pivot.py
+++ b/janitor/functions/pivot.py
@@ -10,6 +10,7 @@
 from pandas.api.types import (
     is_list_like,
     is_string_dtype,
+    is_categorical_dtype,
 )
 from pandas.core.dtypes.concat import concat_compat
 
@@ -1176,6 +1177,7 @@ def _final_frame_longer(
     df = {**index, **outcome, **values}
 
     df = pd.DataFrame(df, copy=False, index=df_index)
+    df_index = None
 
     if sort_by_appearance:
         df = _sort_by_appearance_for_melt(df=df, len_index=len_index)
@@ -1198,6 +1200,9 @@ def pivot_wider(
     flatten_levels: Optional[bool] = True,
     names_sep: str = "_",
     names_glue: str = None,
+    reset_index: bool = True,
+    names_expand: bool = False,
+    index_expand: bool = False,
 ) -> pd.DataFrame:
     """
     Reshapes data from *long* to *wide* form.
@@ -1222,6 +1227,7 @@ def pivot_wider(
     at the start of each label in the columns.
 
 
+
     Example:
 
         >>> import pandas as pd
@@ -1292,9 +1298,16 @@ def pivot_wider(
         and uses python's `str.format_map` under the hood.
         Simply create the string template,
         using the column labels in `names_from`,
-        and special `_value` as a placeholder
-        if there are multiple `values_from`.
+        and special `_value` as a placeholder for `values_from`.
         Applicable only if `flatten_levels` is `True`.
+    :param reset_index: Determines whether to restore `index`
+        as a column/columns. Applicable only if `index` is provided,
+        and `flatten_levels` is `True`. Default is `True`.
+    :param names_expand: Expand columns to show all the categories.
+        Applies only if `names_from` is a categorical column.
+        Default is `False`.
+    :param index_expand: Expand the index to show all the categories.
+        Applies only if `index` is a categorical column. Default is `False`.
     :returns: A pandas DataFrame that has been unpivoted from long to wide
         form.
     """
@@ -1309,6 +1322,9 @@ def pivot_wider(
         flatten_levels,
         names_sep,
         names_glue,
+        reset_index,
+        names_expand,
+        index_expand,
     )
 
 
@@ -1320,6 +1336,9 @@ def _computations_pivot_wider(
     flatten_levels: Optional[bool] = True,
     names_sep: str = "_",
     names_glue: str = None,
+    reset_index: bool = True,
+    names_expand: bool = False,
+    index_expand: bool = False,
 ) -> pd.DataFrame:
     """
     This is the main workhorse of the `pivot_wider` function.
@@ -1339,6 +1358,9 @@ def _computations_pivot_wider(
         flatten_levels,
         names_sep,
         names_glue,
+        reset_index,
+        names_expand,
+        index_expand,
     ) = _data_checks_pivot_wider(
         df,
         index,
@@ -1347,30 +1369,54 @@ def _computations_pivot_wider(
         flatten_levels,
         names_sep,
         names_glue,
+        reset_index,
+        names_expand,
+        index_expand,
     )
-    if flatten_levels:
-        # check dtype of `names_from` is string
-        names_from_all_strings = (
-            df.filter(names_from).agg(is_string_dtype).all().item()
-        )
-
-        # check dtype of columns
-        column_dtype = is_string_dtype(df.columns)
 
     df = df.pivot(  # noqa: PD010
         index=index, columns=names_from, values=values_from
     )
 
-    # an empty df is likely because
-    # there is no `values_from`
+    indexer = df.index
+    if index_expand and index:
+        any_categoricals = (indexer.get_level_values(name) for name in index)
+        any_categoricals = any(map(is_categorical_dtype, any_categoricals))
+        if any_categoricals:
+            indexer = _expand(indexer, retain_categories=True)
+            df = df.reindex(index=indexer)
+
+    indexer = df.columns
+    if names_expand:
+        any_categoricals = (
+            indexer.get_level_values(name) for name in names_from
+        )
+        any_categoricals = any(map(is_categorical_dtype, any_categoricals))
+        if any_categoricals:
+            retain_categories = True
+            if flatten_levels & (
+                (names_glue is not None)
+                | isinstance(indexer, pd.MultiIndex)
+                | ((index is not None) & reset_index)
+            ):
+                retain_categories = False
+            indexer = _expand(indexer, retain_categories=retain_categories)
+            df = df.reindex(columns=indexer)
+
+    indexer = None
     if any((df.empty, not flatten_levels)):
         return df
 
     if isinstance(df.columns, pd.MultiIndex):
-        if (not names_from_all_strings) or (not column_dtype):
-            new_columns = [tuple(map(str, entry)) for entry in df]
-        else:
-            new_columns = [entry for entry in df]
+        new_columns = df.columns
+        all_strings = (
+            new_columns.get_level_values(num)
+            for num in range(new_columns.nlevels)
+        )
+        all_strings = all(map(is_string_dtype, all_strings))
+        if not all_strings:
+            new_columns = (tuple(map(str, entry)) for entry in new_columns)
+
         if names_glue is not None:
             if ("_value" in names_from) and (None in df.columns.names):
                 warnings.warn(
@@ -1403,24 +1449,18 @@ def _computations_pivot_wider(
 
         df.columns = new_columns
     else:
-        if (not names_from_all_strings) or (not column_dtype):
-            df.columns = df.columns.astype(str)
         if names_glue is not None:
             try:
                 df.columns = [
                     names_glue.format_map({names_from[0]: entry})
-                    for entry in df
+                    for entry in df.columns
                 ]
             except KeyError as error:
                 raise KeyError(
                     f"{error} is not a column label in names_from."
                 ) from error
 
-    # if columns are of category type
-    # this returns columns to object dtype
-    # also, resetting index with category columns is not possible
-    df.columns = [*df.columns]
-    if index:
+    if index and reset_index:
         df = df.reset_index()
 
     if df.columns.names:
@@ -1437,6 +1477,9 @@ def _data_checks_pivot_wider(
     flatten_levels,
     names_sep,
     names_glue,
+    reset_index,
+    names_expand,
+    index_expand,
 ):
 
     """
@@ -1464,9 +1507,12 @@ def _data_checks_pivot_wider(
     if values_from is not None:
         if is_list_like(values_from):
             values_from = [*values_from]
-        values_from = _select_column_names(values_from, df)
-        if len(values_from) == 1:
-            values_from = values_from[0]
+        out = _select_column_names(values_from, df)
+        # hack to align with pd.pivot
+        if values_from == out[0]:
+            values_from = out[0]
+        else:
+            values_from = out
 
     check("flatten_levels", flatten_levels, [bool])
 
@@ -1476,6 +1522,10 @@ def _data_checks_pivot_wider(
     if names_glue is not None:
         check("names_glue", names_glue, [str])
 
+    check("reset_index", reset_index, [bool])
+    check("names_expand", names_expand, [bool])
+    check("index_expand", index_expand, [bool])
+
     return (
         df,
         index,
@@ -1484,4 +1534,51 @@ def _data_checks_pivot_wider(
         flatten_levels,
         names_sep,
         names_glue,
+        reset_index,
+        names_expand,
+        index_expand,
     )
+
+
+def _expand(indexer, retain_categories):
+    """
+    Expand Index to all categories.
+    Applies to categorical index, and used
+    in _computations_pivot_wider for scenarios where
+    names_expand and/or index_expand is True.
+    Categories are preserved where possible.
+    If `retain_categories` is False, a fastpath is taken
+    to generate all possible combinations.
+
+    Returns an Index.
+    """
+    if indexer.nlevels > 1:
+        names = indexer.names
+        if not retain_categories:
+            indexer = pd.MultiIndex.from_product(indexer.levels, names=names)
+        else:
+            indexer = [
+                indexer.get_level_values(n) for n in range(indexer.nlevels)
+            ]
+            indexer = [
+                pd.Categorical(
+                    values=arr.categories,
+                    categories=arr.categories,
+                    ordered=arr.ordered,
+                )
+                if is_categorical_dtype(arr)
+                else arr.unique()
+                for arr in indexer
+            ]
+            indexer = pd.MultiIndex.from_product(indexer, names=names)
+
+    else:
+        if not retain_categories:
+            indexer = indexer.categories
+        else:
+            indexer = pd.Categorical(
+                values=indexer.categories,
+                categories=indexer.categories,
+                ordered=indexer.ordered,
+            )
+    return indexer


Test patch:
diff --git a/tests/functions/test_pivot_wider.py b/tests/functions/test_pivot_wider.py
--- a/tests/functions/test_pivot_wider.py
+++ b/tests/functions/test_pivot_wider.py
@@ -135,6 +135,30 @@ def test_name_sep_wrong_type(df_checks_output):
         )
 
 
+def test_name_expand_wrong_type(df_checks_output):
+    """Raise TypeError if the wrong type is provided for `names_expand`."""
+    with pytest.raises(TypeError):
+        df_checks_output.pivot_wider(
+            index="name", names_from=["estimate", "variable"], names_expand=1
+        )
+
+
+def test_id_expand_wrong_type(df_checks_output):
+    """Raise TypeError if the wrong type is provided for `id_expand`."""
+    with pytest.raises(TypeError):
+        df_checks_output.pivot_wider(
+            index="name", names_from=["estimate", "variable"], id_expand=1
+        )
+
+
+def test_reset_index_wrong_type(df_checks_output):
+    """Raise TypeError if the wrong type is provided for `reset_index`."""
+    with pytest.raises(TypeError):
+        df_checks_output.pivot_wider(
+            index="name", names_from=["estimate", "variable"], reset_index=1
+        )
+
+
 def test_name_glue_wrong_type(df_checks_output):
     """Raise TypeError if the wrong type is provided for `names_glue`."""
     with pytest.raises(TypeError):
@@ -317,13 +341,38 @@ def test_index_names():
     assert_frame_equal(result, expected_output)
 
 
+def test_categorical():
+    """Test output for categorical column"""
+    df_in = pd.DataFrame(
+        {
+            "family": ["Kelly", "Kelly", "Quin", "Quin"],
+            "name": ["Mark", "Scott", "Tegan", "Sara"],
+            "n": pd.Categorical([1, 2, 1, 2]),
+        }
+    )
+    df_out = pd.DataFrame(
+        {
+            "family": ["Kelly", "Quin"],
+            1: ["Mark", "Tegan"],
+            2: ["Scott", "Sara"],
+        }
+    )
+
+    result = df_in.pivot_wider(
+        index="family",
+        names_from="n",
+        values_from="name",
+    )
+    assert_frame_equal(result, df_out)
+
+
 def test_names_glue():
     """Test output with `names_glue`"""
     df_in = pd.DataFrame(
         {
             "family": ["Kelly", "Kelly", "Quin", "Quin"],
             "name": ["Mark", "Scott", "Tegan", "Sara"],
-            "n": [1, 2, 1, 2],
+            "n": ["1", "2", "1", "2"],
         }
     )
     df_out = pd.DataFrame(
@@ -343,10 +392,9 @@ def test_names_glue():
     assert_frame_equal(result, df_out)
 
 
-def test_change_level_order(df_checks_output):
+def test_names_glue_multiple_levels(df_checks_output):
     """
-    Test output with `levels_order`,
-    while maintaining order from `names_from`.
+    Test output with names_glue for multiple levels.
     """
 
     df_out = pd.DataFrame(
@@ -367,8 +415,9 @@ def test_change_level_order(df_checks_output):
         names_from="variable",
         values_from=["estimate", "error"],
         names_glue="{variable}_{_value}",
+        reset_index=False,
     )
-    assert_frame_equal(result, df_out)
+    assert_frame_equal(result, df_out.set_index(["geoid", "name"]))
 
 
 def test_names_glue_single_column(df_checks_output):
@@ -421,3 +470,106 @@ def test_int_columns():
     )
 
     assert_frame_equal(result, df_out)
+
+
+@pytest.fixture
+def df_expand():
+    """pytest fixture"""
+    # adapted from
+    # https://github.com/tidyverse/tidyr/issues/770#issuecomment-993872495
+    return pd.DataFrame(
+        dict(
+            id=pd.Categorical(
+                values=(2, 1, 1, 2, 1), categories=(1, 2, 3), ordered=True
+            ),
+            year=(2018, 2018, 2019, 2020, 2020),
+            gender=pd.Categorical(
+                ("female", "male", "male", "female", "male")
+            ),
+            percentage=range(30, 80, 10),
+        ),
+        index=np.repeat([0], 5),
+    )
+
+
+def test_names_expand(df_expand):
+    """Test output if `names_expand`"""
+    actual = df_expand.pivot("year", "id", "percentage").reindex(
+        columns=pd.Categorical([1, 2, 3], ordered=True)
+    )
+    expected = df_expand.pivot_wider(
+        "year", "id", "percentage", names_expand=True, flatten_levels=False
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_names_expand_flatten_levels(df_expand):
+    """Test output if `names_expand`"""
+    actual = (
+        df_expand.pivot("year", "id", "percentage")
+        .reindex(columns=[1, 2, 3])
+        .rename_axis(columns=None)
+        .reset_index()
+    )
+    expected = df_expand.pivot_wider(
+        "year", "id", "percentage", names_expand=True, flatten_levels=True
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_index_expand(df_expand):
+    """Test output if `index_expand`"""
+    actual = df_expand.pivot("id", "year", "percentage").reindex(
+        pd.Categorical([1, 2, 3], ordered=True)
+    )
+    expected = df_expand.pivot_wider(
+        "id", "year", "percentage", index_expand=True, flatten_levels=False
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_index_expand_flatten_levels(df_expand):
+    """Test output if `index_expand`"""
+    actual = (
+        df_expand.pivot("id", "year", "percentage")
+        .reindex(pd.Categorical([1, 2, 3], ordered=True))
+        .rename_axis(columns=None)
+        .reset_index()
+    )
+    expected = df_expand.pivot_wider(
+        "id", "year", "percentage", index_expand=True
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_expand_multiple_levels(df_expand):
+    """Test output for names_expand for multiple names_from."""
+    expected = df_expand.pivot_wider(
+        "id",
+        ("year", "gender"),
+        "percentage",
+        names_expand=True,
+        flatten_levels=False,
+    )
+    actual = df_expand.complete("year", "gender", "id").pivot(
+        "id", ("year", "gender"), "percentage"
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_expand_multiple_levels_flatten_levels(df_expand):
+    """Test output for names_expand for multiple names_from."""
+    expected = df_expand.pivot_wider(
+        "id",
+        ("year", "gender"),
+        "percentage",
+        names_expand=True,
+        flatten_levels=True,
+    )
+    actual = (
+        df_expand.complete("year", "gender", "id")
+        .pivot("id", ("year", "gender"), "percentage")
+        .collapse_levels()
+        .reset_index()
+    )
+    assert_frame_equal(actual, expected)