TYP: enhance ExcelWriter with generic type support and overloads (#1561)

* TYP: enhance `ExcelWriter` with generic type support and overloads

* Update pandas-stubs/io/excel/_base.pyi

Co-authored-by: Yi-Fan Wang <[email protected]>

* fix(types): narrow ExcelWriter path union to FilePath | BinaryIO

Co-authored-by: Yi-Fan Wang <[email protected]>

* fix(types): remove default value from ExcelWriter engine literal

Co-authored-by: Yi-Fan Wang <[email protected]>

* fix(types): remove default value from ExcelWriter engine literal

Co-authored-by: Yi-Fan Wang <[email protected]>

* fix(types): remove default value from ExcelWriter engine literal

Co-authored-by: Yi-Fan Wang <[email protected]>

* fix(types): update ExcelWriter path type to use BinaryIO instead of WriteExcelBuffer

* fix(types): refactor ExcelWriter constructor to use __new__ and improve parameter defaults

---------

Co-authored-by: Yi-Fan Wang <[email protected]>

Author: 9en9i

pandas-stubs/io/excel/_base.pyi

@@ -8,16 +8,25 @@ from collections.abc import (
 from types import TracebackType
 from typing import (
     Any,
+    BinaryIO,
+    Generic,
     Literal,
+    TypeAlias,
     overload,
 )
 
 from odf.opendocument import OpenDocument  # pyright: ignore[reportMissingTypeStubs]
 from openpyxl.workbook.workbook import Workbook
 from pandas.core.frame import DataFrame
 import pyxlsb.workbook  # pyright: ignore[reportMissingTypeStubs]
-from typing_extensions import Self
+from typing_extensions import (
+    Self,
+    TypeVar,
+)
 from xlrd.book import Book
+from xlsxwriter.workbook import (  # pyright: ignore[reportMissingTypeStubs]
+    Workbook as XlsxWorkbook,
+)
 
 from pandas._libs.lib import _NoDefaultDoNotUse
 from pandas._typing import (
@@ -32,7 +41,6 @@ from pandas._typing import (
     ReadBuffer,
     StorageOptions,
     UsecolsArgType,
-    WriteExcelBuffer,
 )
 
 @overload
@@ -209,26 +217,67 @@ def read_excel(
     engine_kwargs: dict[str, Any] | None = ...,
 ) -> DataFrame: ...
 
-class ExcelWriter:
-    def __init__(
-        self,
-        path: FilePath | WriteExcelBuffer | ExcelWriter,
-        engine: ExcelWriteEngine | Literal["auto"] | None = ...,
-        date_format: str | None = ...,
-        datetime_format: str | None = ...,
-        mode: Literal["w", "a"] = ...,
-        storage_options: StorageOptions = ...,
-        if_sheet_exists: ExcelWriterIfSheetExists | None = ...,
-        engine_kwargs: dict[str, Any] | None = ...,
-    ) -> None: ...
+ExcelWriteWorkbook: TypeAlias = Workbook | OpenDocument | XlsxWorkbook
+
+_WorkbookT = TypeVar("_WorkbookT", default=ExcelWriteWorkbook, bound=ExcelWriteWorkbook)
+
+class ExcelWriter(Generic[_WorkbookT]):
+    @overload
+    def __new__(
+        cls,
+        path: FilePath | BinaryIO,
+        engine: Literal["openpyxl"],
+        date_format: str | None = None,
+        datetime_format: str | None = None,
+        mode: Literal["w", "a"] = "w",
+        storage_options: StorageOptions = None,
+        if_sheet_exists: ExcelWriterIfSheetExists | None = None,
+        engine_kwargs: Mapping[str, Any] | None = None,
+    ) -> ExcelWriter[Workbook]: ...
+    @overload
+    def __new__(
+        cls,
+        path: FilePath | BinaryIO,
+        engine: Literal["odf"],
+        date_format: str | None = None,
+        datetime_format: str | None = None,
+        mode: Literal["w", "a"] = "w",
+        storage_options: StorageOptions = None,
+        if_sheet_exists: ExcelWriterIfSheetExists | None = None,
+        engine_kwargs: Mapping[str, Any] | None = None,
+    ) -> ExcelWriter[OpenDocument]: ...
+    @overload
+    def __new__(
+        cls,
+        path: FilePath | BinaryIO,
+        engine: Literal["xlsxwriter"],
+        date_format: str | None = None,
+        datetime_format: str | None = None,
+        mode: Literal["w", "a"] = "w",
+        storage_options: StorageOptions = None,
+        if_sheet_exists: ExcelWriterIfSheetExists | None = None,
+        engine_kwargs: Mapping[str, Any] | None = None,
+    ) -> ExcelWriter[XlsxWorkbook]: ...
+    @overload
+    def __new__(
+        cls,
+        path: FilePath | BinaryIO,
+        engine: Literal["auto"] | None = None,
+        date_format: str | None = None,
+        datetime_format: str | None = None,
+        mode: Literal["w", "a"] = "w",
+        storage_options: StorageOptions = None,
+        if_sheet_exists: ExcelWriterIfSheetExists | None = None,
+        engine_kwargs: Mapping[str, Any] | None = None,
+    ) -> ExcelWriter[ExcelWriteWorkbook]: ...
     @property
     def supported_extensions(self) -> tuple[str, ...]: ...
     @property
     def engine(self) -> ExcelWriteEngine: ...
     @property
     def sheets(self) -> dict[str, Any]: ...
     @property
-    def book(self) -> Workbook | OpenDocument: ...
+    def book(self) -> _WorkbookT: ...
     @property
     def date_format(self) -> str: ...
     @property

pandas-stubs/io/formats/style.pyi

@@ -13,9 +13,13 @@ from typing import (
 )
 
 from matplotlib.colors import Colormap
+from openpyxl.workbook.workbook import Workbook as OpenXlWorkbook
 from pandas.core.frame import DataFrame
 from pandas.core.series import Series
 from typing_extensions import Self
+from xlsxwriter.workbook import (  # pyright: ignore[reportMissingTypeStubs]
+    Workbook as XlsxWorkbook,
+)
 
 from pandas._typing import (
     Axis,
@@ -103,7 +107,9 @@ class Styler(StylerRenderer):
     ) -> Styler: ...
     def to_excel(
         self,
-        excel_writer: FilePath | WriteExcelBuffer | ExcelWriter,
+        excel_writer: (
+            FilePath | WriteExcelBuffer | ExcelWriter[OpenXlWorkbook | XlsxWorkbook]
+        ),
         sheet_name: str = "Sheet1",
         na_rep: str = "",
         float_format: str | None = None,

tests/test_io.py

@@ -13,6 +13,8 @@
 )
 
 import numpy as np
+from odf.opendocument import OpenDocument  # pyright: ignore[reportMissingTypeStubs]
+from openpyxl.workbook.workbook import Workbook as OpenXlWorkbook
 import pandas as pd
 from pandas import (
     DataFrame,
@@ -45,6 +47,9 @@
 import sqlalchemy.orm
 import sqlalchemy.orm.decl_api
 from typing_extensions import assert_type
+from xlsxwriter.workbook import (  # pyright: ignore[reportMissingTypeStubs]
+    Workbook as XlsxWorkbook,
+)
 
 from tests import (
     TYPE_CHECKING_INVALID_USAGE,
@@ -1188,7 +1193,10 @@ def test_excel_writer_engine() -> None:
 
     with ensure_clean(".xlsx") as path:
         with pd.ExcelWriter(path, engine="openpyxl") as ew:
-            check(assert_type(ew, pd.ExcelWriter), pd.ExcelWriter)
+            check(
+                assert_type(ew, pd.ExcelWriter[OpenXlWorkbook]),
+                pd.ExcelWriter[OpenXlWorkbook],
+            )
             DF.to_excel(ew, sheet_name="A")
             check(
                 assert_type(ew.engine, Literal["openpyxl", "odf", "xlsxwriter"]),
@@ -1197,7 +1205,10 @@ def test_excel_writer_engine() -> None:
 
     with ensure_clean(".ods") as path:
         with pd.ExcelWriter(path, engine="odf") as ew:
-            check(assert_type(ew, pd.ExcelWriter), pd.ExcelWriter)
+            check(
+                assert_type(ew, pd.ExcelWriter[OpenDocument]),
+                pd.ExcelWriter[OpenDocument],
+            )
             DF.to_excel(ew, sheet_name="A")
             check(
                 assert_type(ew.engine, Literal["openpyxl", "odf", "xlsxwriter"]),
@@ -1206,7 +1217,10 @@ def test_excel_writer_engine() -> None:
 
     with ensure_clean(".xlsx") as path:
         with pd.ExcelWriter(path, engine="xlsxwriter") as ew:
-            check(assert_type(ew, pd.ExcelWriter), pd.ExcelWriter)
+            check(
+                assert_type(ew, pd.ExcelWriter[XlsxWorkbook]),
+                pd.ExcelWriter[XlsxWorkbook],
+            )
             DF.to_excel(ew, sheet_name="A")
             check(
                 assert_type(ew.engine, Literal["openpyxl", "odf", "xlsxwriter"]),