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

Author: jarent-nvidia

doc/source/user_guide/io.rst

@@ -3717,6 +3717,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

@@ -2178,6 +2178,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.
@@ -2238,6 +2239,9 @@ def to_excel(
 
             .. versionadded:: {storage_options_versionadded}
         {extra_parameters}
+        autofilter : bool, default False
+            If True, add automatic filters to all columns.
+
         See Also
         --------
         to_csv : Write DataFrame to a comma-separated values (csv) file.
@@ -2310,6 +2314,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/_base.py

@@ -1204,6 +1204,7 @@ def _write_cells(
         startrow: int = 0,
         startcol: int = 0,
         freeze_panes: tuple[int, int] | None = None,
+        autofilter_range: str | None = None,
     ) -> None:
         """
         Write given formatted cells into Excel an excel sheet
@@ -1218,6 +1219,8 @@ def _write_cells(
         startcol : upper left cell column to dump data frame
         freeze_panes: int tuple of length 2
             contains the bottom-most row and right-most column to freeze
+        autofilter_range: str, default None
+            column ranges to add automatic filters to, for example "A1:D5"
         """
         raise NotImplementedError
 

pandas/io/excel/_odswriter.py

@@ -99,10 +99,15 @@ 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
         """
+
+        if autofilter_range:
+            raise ValueError("Autofilter is not supported with odf!")
+
         from odf.table import (
             Table,
             TableCell,

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,34 @@ 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.
+
+        Parameters
+        ----------
+        index : int
+            The numeric column index to convert to a Excel column name.
+
+        Returns
+        -------
+        column_name : str
+            The column name corresponding to the index.
+
+        Raises
+        ------
+        ValueError
+            Index is negative
+        """
+        if index < 0:
+            raise ValueError(f"Index cannot be negative: {index}")
+        column_name = ""
+        # while loop in case column name needs to be longer than 1 character
+        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 +948,46 @@ def write(
                 f"Max sheet size is: {self.max_rows}, {self.max_cols}"
             )
 
+        if self.autofilter:
+            # default offset for header row
+            startrowsoffset = 1
+            endrowsoffset = 1
+
+            if num_cols == 0:
+                indexoffset = 0
+            elif self.index:
+                indexoffset = 0
+                if isinstance(self.df.index, MultiIndex):
+                    if self.merge_cells:
+                        raise ValueError(
+                            "Excel filters merged cells by showing only the first row. "
+                            "'autofilter' and 'merge_cells' cannot "
+                            "be used simultaneously."
+                        )
+                    else:
+                        indexoffset = self.df.index.nlevels - 1
+
+                if isinstance(self.columns, MultiIndex):
+                    if self.merge_cells:
+                        raise ValueError(
+                            "Excel filters merged cells by showing only the first row. "
+                            "'autofilter' and 'merge_cells' cannot "
+                            "be used simultaneously."
+                        )
+                    else:
+                        startrowsoffset = self.columns.nlevels
+                        # multindex columns add a blank row between header and data
+                        endrowsoffset = self.columns.nlevels + 1
+            else:
+                # no index column
+                indexoffset = -1
+            start = f"{self._num2excel(startcol)}{startrow + startrowsoffset}"
+            autofilter_end_column = self._num2excel(startcol + num_cols + indexoffset)
+            end = f"{autofilter_end_column}{startrow + num_rows + endrowsoffset}"
+            autofilter_range = f"{start}:{end}"
+        else:
+            autofilter_range = None
+
         if engine_kwargs is None:
             engine_kwargs = {}
 
@@ -938,6 +1010,7 @@ def write(
                 startrow=startrow,
                 startcol=startcol,
                 freeze_panes=freeze_panes,
+                autofilter_range=autofilter_range,
             )
         finally:
             # make sure to close opened file handles

pandas/io/formats/style.py

@@ -579,6 +579,7 @@ def to_excel(
         verbose: bool = True,
         freeze_panes: tuple[int, int] | None = None,
         storage_options: StorageOptions | None = None,
+        autofilter: bool = False,
     ) -> None:
         from pandas.io.formats.excel import ExcelFormatter
 
@@ -592,6 +593,7 @@ def to_excel(
             index_label=index_label,
             merge_cells=merge_cells,
             inf_rep=inf_rep,
+            autofilter=autofilter,
         )
         formatter.write(
             excel_writer,