ENH: adding autofilter when writing to excel (#61194)

Author: jarent-nvidia

doc/source/user_guide/io.rst

@@ -3745,6 +3745,7 @@ The look and feel of Excel worksheets created from pandas can be modified using
 
 * ``float_format`` : Format string for floating point numbers (default ``None``).
 * ``freeze_panes`` : A tuple of two integers representing the bottommost row and rightmost column to freeze. Each of these parameters is one-based, so (1, 1) will freeze the first row and first column (default ``None``).
+* ``autofilter`` : A boolean indicating whether to add automatic filters to all columns (default ``False``).
 
 .. note::
 

doc/source/whatsnew/v3.0.0.rst

@@ -202,6 +202,7 @@ Other enhancements
 - :class:`Holiday` has gained the constructor argument and field ``exclude_dates`` to exclude specific datetimes from a custom holiday calendar (:issue:`54382`)
 - :class:`Rolling` and :class:`Expanding` now support ``nunique`` (:issue:`26958`)
 - :class:`Rolling` and :class:`Expanding` now support aggregations ``first`` and ``last`` (:issue:`33155`)
+- :func:`DataFrame.to_excel` has a new ``autofilter`` parameter to add automatic filters to all columns  (:issue:`61194`)
 - :func:`read_parquet` accepts ``to_pandas_kwargs`` which are forwarded to :meth:`pyarrow.Table.to_pandas` which enables passing additional keywords to customize the conversion to pandas, such as ``maps_as_pydicts`` to read the Parquet map data type as python dictionaries (:issue:`56842`)
 - :func:`to_numeric` on big integers converts to ``object`` datatype with python integers when not coercing. (:issue:`51295`)
 - :meth:`.DataFrameGroupBy.transform`, :meth:`.SeriesGroupBy.transform`, :meth:`.DataFrameGroupBy.agg`, :meth:`.SeriesGroupBy.agg`, :meth:`.SeriesGroupBy.apply`, :meth:`.DataFrameGroupBy.apply` now support ``kurt`` (:issue:`40139`)
@@ -232,7 +233,6 @@ Other enhancements
 - Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
 - Switched wheel upload to **PyPI Trusted Publishing** (OIDC) for release-tag pushes in ``wheels.yml``. (:issue:`61718`)
--
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_300.notable_bug_fixes:

pandas/core/generic.py

@@ -2180,6 +2180,7 @@ def to_excel(
         freeze_panes: tuple[int, int] | None = None,
         storage_options: StorageOptions | None = None,
         engine_kwargs: dict[str, Any] | None = None,
+        autofilter: bool = False,
     ) -> None:
         """
         Write {klass} to an Excel sheet.
@@ -2312,6 +2313,7 @@ def to_excel(
             index_label=index_label,
             merge_cells=merge_cells,
             inf_rep=inf_rep,
+            autofilter=autofilter,
         )
         formatter.write(
             excel_writer,

pandas/io/excel/_odswriter.py

@@ -99,6 +99,7 @@ def _write_cells(
         startrow: int = 0,
         startcol: int = 0,
         freeze_panes: tuple[int, int] | None = None,
+        autofilter_range: str | None = None,
     ) -> None:
         """
         Write the frame cells using odf

pandas/io/excel/_openpyxl.py

@@ -449,6 +449,7 @@ def _write_cells(
         startrow: int = 0,
         startcol: int = 0,
         freeze_panes: tuple[int, int] | None = None,
+        autofilter_range: str | None = None,
     ) -> None:
         # Write the frame cells using openpyxl.
         sheet_name = self._get_sheet_name(sheet_name)
@@ -532,6 +533,9 @@ def _write_cells(
                             for k, v in style_kwargs.items():
                                 setattr(xcell, k, v)
 
+        if autofilter_range:
+            wks.auto_filter.ref = autofilter_range
+
 
 class OpenpyxlReader(BaseExcelReader["Workbook"]):
     @doc(storage_options=_shared_docs["storage_options"])

pandas/io/excel/_xlsxwriter.py

@@ -245,6 +245,7 @@ def _write_cells(
         startrow: int = 0,
         startcol: int = 0,
         freeze_panes: tuple[int, int] | None = None,
+        autofilter_range: str | None = None,
     ) -> None:
         # Write the frame cells using xlsxwriter.
         sheet_name = self._get_sheet_name(sheet_name)
@@ -282,3 +283,6 @@ def _write_cells(
                 )
             else:
                 wks.write(startrow + cell.row, startcol + cell.col, val, style)
+
+        if autofilter_range:
+            wks.autofilter(autofilter_range)

pandas/io/formats/excel.py

@@ -532,6 +532,8 @@ class ExcelFormatter:
         Defaults to ``CSSToExcelConverter()``.
         It should have signature css_declarations string -> excel style.
         This is only called for body cells.
+    autofilter : bool, default False
+        If True, add automatic filters to all columns
     """
 
     max_rows = 2**20
@@ -549,6 +551,7 @@ def __init__(
         merge_cells: ExcelWriterMergeCells = False,
         inf_rep: str = "inf",
         style_converter: Callable | None = None,
+        autofilter: bool = False,
     ) -> None:
         self.rowcounter = 0
         self.na_rep = na_rep
@@ -584,6 +587,7 @@ def __init__(
             raise ValueError(f"Unexpected value for {merge_cells=}.")
         self.merge_cells = merge_cells
         self.inf_rep = inf_rep
+        self.autofilter = autofilter
 
     def _format_value(self, val):
         if is_scalar(val) and missing.isna(val):
@@ -873,6 +877,16 @@ def get_formatted_cells(self) -> Iterable[ExcelCell]:
             cell.val = self._format_value(cell.val)
             yield cell
 
+    def _num2excel(self, index: int) -> str:
+        """
+        Convert 0-based column index to Excel column name.
+        """
+        column_name = ""
+        while index > 0 or not column_name:
+            index, remainder = divmod(index, 26)
+            column_name = chr(65 + remainder) + column_name
+        return column_name
+
     @doc(storage_options=_shared_docs["storage_options"])
     def write(
         self,
@@ -916,6 +930,13 @@ def write(
                 f"Max sheet size is: {self.max_rows}, {self.max_cols}"
             )
 
+        if self.autofilter:
+            start = f"{self._num2excel(startcol)}{startrow + 1}"
+            end = f"{self._num2excel(startcol + num_cols)}{startrow + num_rows + 1}"
+            autofilter_range = f"{start}:{end}"
+        else:
+            autofilter_range = None
+
         if engine_kwargs is None:
             engine_kwargs = {}
 
@@ -938,6 +959,7 @@ def write(
                 startrow=startrow,
                 startcol=startcol,
                 freeze_panes=freeze_panes,
+                autofilter_range=autofilter_range,
             )
         finally:
             # make sure to close opened file handles

pandas/tests/io/excel/test_style.py

@@ -350,3 +350,20 @@ def test_format_hierarchical_rows_periodindex(merge_cells):
             assert isinstance(cell.val, Timestamp), (
                 "Period should be converted to Timestamp"
             )
+
+
+@pytest.mark.parametrize("engine", ["xlsxwriter", "openpyxl"])
+def test_autofilter(engine, tmp_excel):
+    # GH 61194
+    df = DataFrame.from_dict([{"A": 1, "B": 2, "C": 3}, {"A": 4, "B": 5, "C": 6}])
+
+    with ExcelWriter(tmp_excel, engine=engine) as writer:
+        df.to_excel(writer, autofilter=True, index=False)
+
+    openpyxl = pytest.importorskip("openpyxl")  # test loading only with openpyxl
+    with contextlib.closing(openpyxl.load_workbook(tmp_excel)) as wb:
+        ws = wb.active
+
+        assert ws.auto_filter.ref is not None
+        print(ws.auto_filter.ref)
+        assert ws.auto_filter.ref == "A1:D3"