TYP: Add types for excel writer classes (#45246)

Author: phofl

pandas/io/excel/_base.py

@@ -533,7 +533,7 @@ def _workbook_class(self):
     def load_workbook(self, filepath_or_buffer):
         pass
 
-    def close(self):
+    def close(self) -> None:
         if hasattr(self, "book") and hasattr(self.book, "close"):
             # pyxlsb: opens a TemporaryFile
             # openpyxl: https://stackoverflow.com/questions/31416842/
@@ -1083,7 +1083,7 @@ def __init__(
         mode: str = "w",
         storage_options: StorageOptions = None,
         if_sheet_exists: str | None = None,
-        engine_kwargs: dict | None = None,
+        engine_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ):
         # validate that this engine can handle the extension

pandas/io/excel/_odswriter.py

@@ -5,10 +5,16 @@
 from typing import (
     Any,
     DefaultDict,
+    Tuple,
+    cast,
 )
 
 import pandas._libs.json as json
-from pandas._typing import StorageOptions
+from pandas._typing import (
+    FilePath,
+    StorageOptions,
+    WriteExcelBuffer,
+)
 
 from pandas.io.excel._base import ExcelWriter
 from pandas.io.excel._util import (
@@ -24,9 +30,9 @@ class ODSWriter(ExcelWriter):
 
     def __init__(
         self,
-        path: str,
+        path: FilePath | WriteExcelBuffer | ExcelWriter,
         engine: str | None = None,
-        date_format=None,
+        date_format: str | None = None,
         datetime_format=None,
         mode: str = "w",
         storage_options: StorageOptions = None,
@@ -88,7 +94,7 @@ def write_cells(
             self.sheets[sheet_name] = wks
 
         if validate_freeze_panes(freeze_panes):
-            assert freeze_panes is not None
+            freeze_panes = cast(Tuple[int, int], freeze_panes)
             self._create_freeze_panes(sheet_name, freeze_panes)
 
         for _ in range(startrow):

pandas/io/excel/_openpyxl.py

@@ -4,6 +4,8 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Tuple,
+    cast,
 )
 
 import numpy as np
@@ -13,6 +15,7 @@
     ReadBuffer,
     Scalar,
     StorageOptions,
+    WriteExcelBuffer,
 )
 from pandas.compat._optional import import_optional_dependency
 
@@ -35,10 +38,10 @@ class OpenpyxlWriter(ExcelWriter):
 
     def __init__(
         self,
-        path,
-        engine=None,
-        date_format=None,
-        datetime_format=None,
+        path: FilePath | WriteExcelBuffer | ExcelWriter,
+        engine: str | None = None,
+        date_format: str | None = None,
+        datetime_format: str | None = None,
         mode: str = "w",
         storage_options: StorageOptions = None,
         if_sheet_exists: str | None = None,
@@ -74,7 +77,7 @@ def __init__(
             if self.book.worksheets:
                 self.book.remove(self.book.worksheets[0])
 
-    def save(self):
+    def save(self) -> None:
         """
         Save workbook to disk.
         """
@@ -217,7 +220,7 @@ def _convert_to_stop(cls, stop_seq):
         return map(cls._convert_to_color, stop_seq)
 
     @classmethod
-    def _convert_to_fill(cls, fill_dict):
+    def _convert_to_fill(cls, fill_dict: dict[str, Any]):
         """
         Convert ``fill_dict`` to an openpyxl v2 Fill object.
 
@@ -418,8 +421,13 @@ 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: str | None = None,
+        startrow: int = 0,
+        startcol: int = 0,
+        freeze_panes: tuple[int, int] | None = None,
+    ) -> None:
         # Write the frame cells using openpyxl.
         sheet_name = self._get_sheet_name(sheet_name)
 
@@ -453,6 +461,7 @@ def write_cells(
             self.sheets[sheet_name] = wks
 
         if validate_freeze_panes(freeze_panes):
+            freeze_panes = cast(Tuple[int, int], freeze_panes)
             wks.freeze_panes = wks.cell(
                 row=freeze_panes[0] + 1, column=freeze_panes[1] + 1
             )

pandas/io/excel/_xlsxwriter.py

@@ -3,7 +3,11 @@
 from typing import Any
 
 import pandas._libs.json as json
-from pandas._typing import StorageOptions
+from pandas._typing import (
+    FilePath,
+    StorageOptions,
+    WriteExcelBuffer,
+)
 
 from pandas.io.excel._base import ExcelWriter
 from pandas.io.excel._util import (
@@ -170,10 +174,10 @@ class XlsxWriter(ExcelWriter):
 
     def __init__(
         self,
-        path,
-        engine=None,
-        date_format=None,
-        datetime_format=None,
+        path: FilePath | WriteExcelBuffer | ExcelWriter,
+        engine: str | None = None,
+        date_format: str | None = None,
+        datetime_format: str | None = None,
         mode: str = "w",
         storage_options: StorageOptions = None,
         if_sheet_exists: str | None = None,
@@ -201,15 +205,20 @@ def __init__(
 
         self.book = Workbook(self.handles.handle, **engine_kwargs)
 
-    def save(self):
+    def save(self) -> None:
         """
         Save workbook to disk.
         """
-        return self.book.close()
+        self.book.close()
 
     def write_cells(
-        self, cells, sheet_name=None, startrow=0, startcol=0, freeze_panes=None
-    ):
+        self,
+        cells,
+        sheet_name: str | None = None,
+        startrow: int = 0,
+        startcol: int = 0,
+        freeze_panes: tuple[int, int] | None = None,
+    ) -> None:
         # Write the frame cells using xlsxwriter.
         sheet_name = self._get_sheet_name(sheet_name)
 

pandas/io/excel/_xlwt.py

@@ -3,10 +3,16 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Tuple,
+    cast,
 )
 
 import pandas._libs.json as json
-from pandas._typing import StorageOptions
+from pandas._typing import (
+    FilePath,
+    StorageOptions,
+    WriteExcelBuffer,
+)
 
 from pandas.io.excel._base import ExcelWriter
 from pandas.io.excel._util import (
@@ -24,11 +30,11 @@ class XlwtWriter(ExcelWriter):
 
     def __init__(
         self,
-        path,
-        engine=None,
-        date_format=None,
-        datetime_format=None,
-        encoding=None,
+        path: FilePath | WriteExcelBuffer | ExcelWriter,
+        engine: str | None = None,
+        date_format: str | None = None,
+        datetime_format: str | None = None,
+        encoding: str | None = None,
         mode: str = "w",
         storage_options: StorageOptions = None,
         if_sheet_exists: str | None = None,
@@ -57,7 +63,7 @@ def __init__(
         self.fm_datetime = xlwt.easyxf(num_format_str=self.datetime_format)
         self.fm_date = xlwt.easyxf(num_format_str=self.date_format)
 
-    def save(self):
+    def save(self) -> None:
         """
         Save workbook to disk.
         """
@@ -66,8 +72,13 @@ 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: str | None = None,
+        startrow: int = 0,
+        startcol: int = 0,
+        freeze_panes: tuple[int, int] | None = None,
+    ) -> None:
 
         sheet_name = self._get_sheet_name(sheet_name)
 
@@ -78,6 +89,7 @@ def write_cells(
             self.sheets[sheet_name] = wks
 
         if validate_freeze_panes(freeze_panes):
+            freeze_panes = cast(Tuple[int, int], freeze_panes)
             wks.set_panes_frozen(True)
             wks.set_horz_split_pos(freeze_panes[0])
             wks.set_vert_split_pos(freeze_panes[1])
@@ -111,7 +123,7 @@ def write_cells(
 
     @classmethod
     def _style_to_xlwt(
-        cls, item, firstlevel: bool = True, field_sep=",", line_sep=";"
+        cls, item, firstlevel: bool = True, field_sep: str = ",", line_sep: str = ";"
     ) -> str:
         """
         helper which recursively generate an xlwt easy style string
@@ -150,7 +162,9 @@ def _style_to_xlwt(
             return item
 
     @classmethod
-    def _convert_to_style(cls, style_dict, num_format_str=None):
+    def _convert_to_style(
+        cls, style_dict, num_format_str: str | None = None
+    ) -> XFStyle:
         """
         converts a style_dict to an xlwt style object