Polars pivot_longer implementation (#1355)

Implemented `pivot_longer` for polars.
pyjanitor-devs · Jun 9, 2024 · a672fef · a672fef
1 parent 70fe127
commit a672fef
Show file tree

Hide file tree

Showing 8 changed files with 1,831 additions and 313 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
diff --git a/janitor/functions/clean_names.py b/janitor/functions/clean_names.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import unicodedata
-from typing import Optional, Union
 
 import pandas as pd
 import pandas_flavor as pf
@@ -18,9 +17,9 @@
 @deprecated_alias(preserve_original_columns="preserve_original_labels")
 def clean_names(
     df: pd.DataFrame,
-    axis: Union[str, None] = "columns",
-    column_names: Union[str, list] = None,
-    strip_underscores: Optional[Union[str, bool]] = None,
+    axis: str = "columns",
+    column_names: str | list = None,
+    strip_underscores: str | bool = None,
     case_type: str = "lower",
     remove_special: bool = False,
     strip_accents: bool = True,
@@ -170,14 +169,14 @@ def clean_names(
 
 
 def _clean_names(
-    obj: Union[pd.Index, pd.Series],
-    strip_underscores: Optional[Union[str, bool]] = None,
-    case_type: str = "lower",
-    remove_special: bool = False,
-    strip_accents: bool = False,
-    enforce_string: bool = False,
-    truncate_limit: int = None,
-) -> Union[pd.Index, pd.Series]:
+    obj: pd.Index | pd.Series,
+    strip_underscores: str | bool,
+    case_type: str,
+    remove_special: bool,
+    strip_accents: bool,
+    enforce_string: bool,
+    truncate_limit: int,
+) -> pd.Index | pd.Series:
     """
     Generic function to clean labels in a pandas object.
     """
@@ -202,9 +201,9 @@ def _clean_names(
 
 
 def _change_case(
-    obj: Union[pd.Index, pd.Series],
+    obj: pd.Index | pd.Series,
     case_type: str,
-) -> Union[pd.Index, pd.Series]:
+) -> pd.Index | pd.Series:
     """Change case of labels in obj."""
     case_types = {"preserve", "upper", "lower", "snake"}
     case_type = case_type.lower()
@@ -226,9 +225,7 @@ def _change_case(
     )
 
 
-def _normalize_1(
-    obj: Union[pd.Index, pd.Series]
-) -> Union[pd.Index, pd.Series]:
+def _normalize_1(obj: pd.Index | pd.Series) -> pd.Index | pd.Series:
     """Perform normalization of labels in obj."""
     FIXES = [(r"[ /:,?()\.-]", "_"), (r"['’]", ""), (r"[\xa0]", "_")]
     for search, replace in FIXES:
@@ -238,8 +235,8 @@ def _normalize_1(
 
 
 def _strip_accents(
-    obj: Union[pd.Index, pd.Series],
-) -> Union[pd.Index, pd.Series]:
+    obj: pd.Index | pd.Series,
+) -> pd.Index | pd.Series:
     """Remove accents from a label.
 
     Inspired from [StackOverflow][so].
@@ -258,9 +255,9 @@ def _strip_accents(
 
 
 def _strip_underscores_func(
-    obj: Union[pd.Index, pd.Series],
-    strip_underscores: Union[str, bool] = None,
-) -> Union[pd.Index, pd.Series]:
+    obj: pd.Index | pd.Series,
+    strip_underscores: str | bool = None,
+) -> pd.Index | pd.Series:
     """Strip underscores."""
     underscore_options = {None, "left", "right", "both", "l", "r", True}
     if strip_underscores not in underscore_options:

diff --git a/janitor/functions/pivot.py b/janitor/functions/pivot.py
@@ -98,7 +98,7 @@ def pivot_longer(
         6     setosa   Petal.Width    0.2
         7  virginica   Petal.Width    1.8
 
-        Split the column labels into parts:
+        Split the column labels into individual columns:
         >>> df.pivot_longer(
         ...     index = 'Species',
         ...     names_to = ('part', 'dimension'),
@@ -167,7 +167,7 @@ def pivot_longer(
         value        int64
         dtype: object
 
-        Use multiple `.value` to reshape dataframe:
+        Use multiple `.value` to reshape the dataframe:
         >>> df = pd.DataFrame(
         ...     [
         ...         {
@@ -265,16 +265,6 @@ def pivot_longer(
         ...             "Gin": [16, 200, 34],
         ...             "Vodka": [20, 33, 18],
         ...         },
-        ...         columns=[
-        ...             "City",
-        ...             "State",
-        ...             "Name",
-        ...             "Mango",
-        ...             "Orange",
-        ...             "Watermelon",
-        ...             "Gin",
-        ...             "Vodka",
-        ...         ],
         ...     )
         >>> df
               City    State      Name  Mango  Orange  Watermelon  Gin  Vodka