Repository: pyjanitor-devs/pyjanitor (python)

Issue:
extend the `select` function to `pd.Series`
At the moment, `jn.select` applies only to DataFrames. Extending to pd.Series would be handy, especially when selecting MultiIndexes.

Pull request:
[ENH] Add support for `pd.Series.select`


# PR Description

Please describe the changes proposed in the pull request:

- Add support for `pd.Series.select`


**This PR resolves #1394 .**

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6278eca..889c2cc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 
+-  [ENH] Added support for pd.Series.select - Issue #1394 @samukweku
+
 ## [v0.30.0] - 2024-12-04
 
 ## [v0.29.2] - 2024-09-28
@@ -28,7 +30,7 @@
 -   [ENH] Added a `clean_names` method for polars - it can be used to clean the column names, or clean column values . Issue #1343 @samukweku
 -   [ENH] Improved performance for non-equi joins when using numba - @samukweku PR #1341
 -   [ENH] pandas Index,Series, DataFrame now supported in the `complete` method. - PR #1369 @samukweku
--   [ENH] Improve performance for `first/last` in \`conditional_join, when the join columns in the right dataframe are sorted. - PR #1382 @samukweku
+-   [ENH] Improve performance for `first/last` in conditional_join, when the join columns in the right dataframe are sorted. - PR #1382 @samukweku
 
 ## [v0.27.0] - 2024-03-21
 
diff --git a/janitor/functions/select.py b/janitor/functions/select.py
index 3f2e90c..0c076fc 100644
--- a/janitor/functions/select.py
+++ b/janitor/functions/select.py
@@ -328,16 +328,17 @@ def select_rows(
 
 
 @pf.register_dataframe_method
+@pf.register_series_method
 @deprecated_alias(rows="index")
 def select(
-    df: pd.DataFrame,
+    df: pd.DataFrame | pd.Series,
     *args: tuple,
     index: Any = None,
     columns: Any = None,
     axis: str = "columns",
     invert: bool = False,
-) -> pd.DataFrame:
-    """Method-chainable selection of rows and columns.
+) -> pd.DataFrame | pd.Series:
+    """Method-chainable selection of rows and/or columns.
 
     It accepts a string, shell-like glob strings `(*string*)`,
     regex, slice, array-like object, or a list of the previous options.
@@ -345,7 +346,9 @@ def select(
     Selection on a MultiIndex on a level, or multiple levels,
     is possible with a dictionary.
 
-    This method does not mutate the original DataFrame.
+    This method does not mutate the original DataFrame or Series.
+
+    If the pandas object is a Series, selection is possible only on the index.
 
     Selection can be inverted with the `DropLabel` class.
 
@@ -366,6 +369,8 @@ def select(
         - 0.26.0
             - Added variable `args`, `invert` and `axis` parameters.
             - `rows` keyword deprecated in favour of `index`.
+        - 0.31.0
+            - Add support for pd.Series.
 
     Examples:
         >>> import pandas as pd
@@ -429,9 +434,8 @@ def select(
         ValueError: If args and index/columns are provided.
 
     Returns:
-        A pandas DataFrame with the specified rows and/or columns selected.
+        A pandas DataFrame or Series with the specified rows and/or columns selected.
     """  # noqa: E501
-
     if args:
         check("invert", invert, [bool])
         if (index is not None) or (columns is not None):
@@ -851,16 +855,19 @@ def _index_converter(arr, index):
 
 
 def _select(
-    df: pd.DataFrame,
+    df: pd.DataFrame | pd.Series,
     invert: bool = False,
     rows=None,
     columns=None,
-) -> pd.DataFrame:
+) -> pd.DataFrame | pd.Series:
     """
     Index DataFrame on the index or columns.
+    If it is a Series, indexing is only on the index.
 
-    Returns a DataFrame.
+    Returns a DataFrame or Series.
     """
+    if (columns is not None) and isinstance(df, pd.Series):
+        raise ValueError("columns axis is not supported for pd.Series.select")
     if rows is None:
         row_indexer = slice(None)
     else:
@@ -870,6 +877,8 @@ def _select(
             row_indexer[outcome] = False
         else:
             row_indexer = outcome
+    if isinstance(df, pd.Series):
+        return df.iloc[row_indexer]
     if columns is None:
         column_indexer = slice(None)
     else:
diff --git a/mkdocs.yml b/mkdocs.yml
index 428df15..6613c95 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -63,7 +63,7 @@ plugins:
             docstring_style: "google"
             docstring_options:
               trim_doctest_flags: true
-              show_if_no_docstring: false
+            show_if_no_docstring: false
             show_root_toc_entry: false
             show_root_heading: false
             show_submodules: true
diff --git a/mkdocs/api/biology.md b/mkdocs/api/biology.md
index 53b35fa..2522b5c 100644
--- a/mkdocs/api/biology.md
+++ b/mkdocs/api/biology.md
@@ -1,3 +1,6 @@
 # Biology
 
 ::: janitor.biology
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/chemistry.md b/mkdocs/api/chemistry.md
index a869ffd..6fed7df 100644
--- a/mkdocs/api/chemistry.md
+++ b/mkdocs/api/chemistry.md
@@ -1,3 +1,6 @@
 # Chemistry
 
 ::: janitor.chemistry
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/engineering.md b/mkdocs/api/engineering.md
index 520a5f8..5e537e4 100644
--- a/mkdocs/api/engineering.md
+++ b/mkdocs/api/engineering.md
@@ -1,3 +1,6 @@
 # Engineering
 
 ::: janitor.engineering
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/finance.md b/mkdocs/api/finance.md
index dbab6ed..291e6ef 100644
--- a/mkdocs/api/finance.md
+++ b/mkdocs/api/finance.md
@@ -1,3 +1,6 @@
 # Finance
 
 ::: janitor.finance
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/functions.md b/mkdocs/api/functions.md
index 5345b0e..445f173 100644
--- a/mkdocs/api/functions.md
+++ b/mkdocs/api/functions.md
@@ -2,6 +2,8 @@
 
 ::: janitor.functions
     options:
+      filters:
+      - "!^_"
       members:
         - add_columns
         - also
diff --git a/mkdocs/api/io.md b/mkdocs/api/io.md
index 29294ab..0ec736a 100644
--- a/mkdocs/api/io.md
+++ b/mkdocs/api/io.md
@@ -1,3 +1,6 @@
 # Input/Output (io)
 
 ::: janitor.io
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/math.md b/mkdocs/api/math.md
index 0d7a34b..c06e423 100644
--- a/mkdocs/api/math.md
+++ b/mkdocs/api/math.md
@@ -1,3 +1,6 @@
 # Math
 
 ::: janitor.math
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/ml.md b/mkdocs/api/ml.md
index a8fb80b..f712465 100644
--- a/mkdocs/api/ml.md
+++ b/mkdocs/api/ml.md
@@ -1,3 +1,6 @@
 # Machine Learning
 
 ::: janitor.ml
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/polars.md b/mkdocs/api/polars.md
index 8c58635..eaa6b25 100644
--- a/mkdocs/api/polars.md
+++ b/mkdocs/api/polars.md
@@ -2,6 +2,8 @@
 
 ::: janitor.polars
     options:
+      filters:
+      - "!^_"
       members:
         - clean_names
         - complete
diff --git a/mkdocs/api/timeseries.md b/mkdocs/api/timeseries.md
index a16fda5..43c77f3 100644
--- a/mkdocs/api/timeseries.md
+++ b/mkdocs/api/timeseries.md
@@ -1,3 +1,6 @@
 # Timeseries
 
 ::: janitor.timeseries
+    options:
+      filters:
+      - "!^_"
diff --git a/mkdocs/api/xarray.md b/mkdocs/api/xarray.md
index d7e81c1..3891342 100644
--- a/mkdocs/api/xarray.md
+++ b/mkdocs/api/xarray.md
@@ -1,3 +1,6 @@
 # XArray
 
 ::: janitor.xarray.functions
+    options:
+      filters:
+      - "!^_"


Test patch:
diff --git a/tests/functions/test_select.py b/tests/functions/test_select.py
index 605103f..114901b 100644
--- a/tests/functions/test_select.py
+++ b/tests/functions/test_select.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from pandas.testing import assert_frame_equal
+from pandas.testing import assert_frame_equal, assert_series_equal
 
 from janitor.functions.select import DropLabel
 
@@ -22,6 +22,22 @@ def dataframe():
     )
 
 
+def test_series_axis(dataframe):
+    """Raise if object is a  Series and axis is a column"""
+    with pytest.raises(
+        ValueError, match="columns axis is not supported for pd.Series.select"
+    ):
+        dataframe.iloc[:, 0].select("bar", axis="columns")
+
+
+def test_series_select(dataframe):
+    """Test output when selecting on a Series"""
+    series = dataframe.iloc[:, 0]
+    actual = series.select(index="bar")
+    expected = series.loc[["bar"]]
+    assert_series_equal(actual, expected)
+
+
 def test_args_and_rows_and_columns(dataframe):
     """
     Raise if args and rows/columns are provided.