Added autofilter parameter to Excel output for both xlsxwriter and openpyxl engines.

Author: Upabjojr

pandas/core/generic.py

@@ -2145,6 +2145,7 @@ def to_excel(
         inf_rep="inf",
         verbose=True,
         freeze_panes=None,
+        autofilter=False,
         storage_options: StorageOptions = None,
     ) -> None:
         """
@@ -2210,6 +2211,9 @@ def to_excel(
         freeze_panes : tuple of int (length 2), optional
             Specifies the one-based bottommost row and rightmost column that
             is to be frozen.
+        autofilter : bool, default False
+            Specifies whether filters should be added to the columns in the
+            spreadsheet.
         {storage_options}
 
             .. versionadded:: 1.2.0
@@ -2286,6 +2290,7 @@ def to_excel(
             startrow=startrow,
             startcol=startcol,
             freeze_panes=freeze_panes,
+            autofilter=autofilter,
             engine=engine,
             storage_options=storage_options,
         )

pandas/io/excel/_base.py

@@ -871,7 +871,8 @@ def engine(self):
 
     @abc.abstractmethod
     def write_cells(
-        self, cells, sheet_name=None, startrow=0, startcol=0, freeze_panes=None
+        self, cells, sheet_name=None, startrow=0, startcol=0,
+        freeze_panes=None, autofilter=None
     ):
         """
         Write given formatted cells into Excel an excel sheet

pandas/io/excel/_odswriter.py

@@ -62,6 +62,7 @@ def write_cells(
         startrow: int = 0,
         startcol: int = 0,
         freeze_panes: tuple[int, int] | None = None,
+        autofilter: bool = False,
     ) -> None:
         """
         Write the frame cells using odf

pandas/io/excel/_openpyxl.py

@@ -417,7 +417,8 @@ def _convert_to_protection(cls, protection_dict):
         return Protection(**protection_dict)
 
     def write_cells(
-        self, cells, sheet_name=None, startrow=0, startcol=0, freeze_panes=None
+        self, cells, sheet_name=None, startrow=0, startcol=0,
+        freeze_panes=None, autofilter=False
     ):
         # Write the frame cells using openpyxl.
         sheet_name = self._get_sheet_name(sheet_name)
@@ -454,6 +455,9 @@ def write_cells(
                 row=freeze_panes[0] + 1, column=freeze_panes[1] + 1
             )
 
+        max_row = 0
+        max_col = 0
+
         for cell in cells:
             xcell = wks.cell(
                 row=startrow + cell.row + 1, column=startcol + cell.col + 1
@@ -501,6 +505,15 @@ def write_cells(
                             for k, v in style_kwargs.items():
                                 setattr(xcell, k, v)
 
+            if startrow + cell.row > max_row:
+                max_row = startrow + cell.row
+            if startcol + cell.col > max_col:
+                max_col = startcol + cell.col
+
+        if autofilter:
+            from openpyxl.utils import get_column_letter
+            wks.auto_filter.ref = "A1:" + get_column_letter(max_col + 1) + str(max_row)
+
 
 class OpenpyxlReader(BaseExcelReader):
     def __init__(

pandas/io/excel/_xlsxwriter.py

@@ -208,7 +208,8 @@ def save(self):
         return self.book.close()
 
     def write_cells(
-        self, cells, sheet_name=None, startrow=0, startcol=0, freeze_panes=None
+        self, cells, sheet_name=None, startrow=0, startcol=0, freeze_panes=None,
+        autofilter=False
     ):
         # Write the frame cells using xlsxwriter.
         sheet_name = self._get_sheet_name(sheet_name)
@@ -224,6 +225,9 @@ def write_cells(
         if validate_freeze_panes(freeze_panes):
             wks.freeze_panes(*(freeze_panes))
 
+        max_row = 0
+        max_col = 0
+
         for cell in cells:
             val, fmt = self._value_with_fmt(cell.val)
 
@@ -248,3 +252,11 @@ def write_cells(
                 )
             else:
                 wks.write(startrow + cell.row, startcol + cell.col, val, style)
+
+            if startrow + cell.row > max_row:
+                max_row = startrow + cell.row
+            if startcol + cell.col > max_col:
+                max_col = startcol + cell.col
+
+        if autofilter:
+            wks.autofilter(0, 0, max_row, max_col)

pandas/io/excel/_xlwt.py

@@ -66,7 +66,8 @@ def save(self):
             self.book.save(self.handles.handle)
 
     def write_cells(
-        self, cells, sheet_name=None, startrow=0, startcol=0, freeze_panes=None
+        self, cells, sheet_name=None, startrow=0, startcol=0,
+        freeze_panes=None, autofilter=False,
     ):
 
         sheet_name = self._get_sheet_name(sheet_name)

pandas/io/formats/excel.py

@@ -786,6 +786,7 @@ def write(
         startrow=0,
         startcol=0,
         freeze_panes=None,
+        autofilter=False,
         engine=None,
         storage_options: StorageOptions = None,
     ):
@@ -801,6 +802,9 @@ def write(
         freeze_panes : tuple of integer (length 2), default None
             Specifies the one-based bottommost row and rightmost column that
             is to be frozen
+        autofilter : bool, default False
+            Specifies whether filters should be added to the columns in the
+            spreadsheet.
         engine : string, default None
             write engine to use if writer is a path - you can also set this
             via the options ``io.excel.xlsx.writer``, ``io.excel.xls.writer``,
@@ -843,6 +847,7 @@ def write(
                 startrow=startrow,
                 startcol=startcol,
                 freeze_panes=freeze_panes,
+                autofilter=autofilter,
             )
         finally:
             # make sure to close opened file handles

pandas/tests/io/excel/test_writers.py

@@ -1270,6 +1270,13 @@ def test_freeze_panes(self, path):
         result = pd.read_excel(path, index_col=0)
         tm.assert_frame_equal(result, expected)
 
+    def test_autofilter(self, path):
+        expected = DataFrame([[1, 2], [3, 4]], columns=["col1", "col2"])
+        expected.to_excel(path, "Sheet1", autofilter=True)
+
+        result = pd.read_excel(path, index_col=0)
+        tm.assert_frame_equal(result, expected)
+
     def test_path_path_lib(self, engine, ext):
         df = tm.makeDataFrame()
         writer = partial(df.to_excel, engine=engine)