TYP: type excel util module (#45014)

Author: phofl

pandas/io/excel/_base.py

@@ -943,7 +943,7 @@ def supported_extensions(self):
 
     @property
     @abc.abstractmethod
-    def engine(self):
+    def engine(self) -> str:
         """Name of engine."""
         pass
 

pandas/io/excel/_util.py

@@ -1,8 +1,16 @@
 from __future__ import annotations
 
 from typing import (
+    TYPE_CHECKING,
     Any,
+    Callable,
+    Hashable,
+    Iterable,
+    Literal,
     MutableMapping,
+    Sequence,
+    TypeVar,
+    overload,
 )
 
 from pandas.compat._optional import import_optional_dependency
@@ -12,10 +20,16 @@
     is_list_like,
 )
 
-_writers: MutableMapping[str, str] = {}
+if TYPE_CHECKING:
+    from pandas.io.excel._base import ExcelWriter
 
+    ExcelWriter_t = type[ExcelWriter]
+    usecols_func = TypeVar("usecols_func", bound=Callable[[Hashable], object])
 
-def register_writer(klass):
+_writers: MutableMapping[str, ExcelWriter_t] = {}
+
+
+def register_writer(klass: ExcelWriter_t) -> None:
     """
     Add engine to the excel writer registry.io.excel.
 
@@ -28,10 +42,12 @@ def register_writer(klass):
     if not callable(klass):
         raise ValueError("Can only register callables as engines")
     engine_name = klass.engine
+    # for mypy
+    assert isinstance(engine_name, str)
     _writers[engine_name] = klass
 
 
-def get_default_engine(ext, mode="reader"):
+def get_default_engine(ext: str, mode: Literal["reader", "writer"] = "reader") -> str:
     """
     Return the default reader/writer for the given extension.
 
@@ -73,7 +89,7 @@ def get_default_engine(ext, mode="reader"):
         return _default_readers[ext]
 
 
-def get_writer(engine_name):
+def get_writer(engine_name: str) -> ExcelWriter_t:
     try:
         return _writers[engine_name]
     except KeyError as err:
@@ -145,7 +161,29 @@ def _range2cols(areas: str) -> list[int]:
     return cols
 
 
-def maybe_convert_usecols(usecols):
+@overload
+def maybe_convert_usecols(usecols: str | list[int]) -> list[int]:
+    ...
+
+
+@overload
+def maybe_convert_usecols(usecols: list[str]) -> list[str]:
+    ...
+
+
+@overload
+def maybe_convert_usecols(usecols: usecols_func) -> usecols_func:
+    ...
+
+
+@overload
+def maybe_convert_usecols(usecols: None) -> None:
+    ...
+
+
+def maybe_convert_usecols(
+    usecols: str | list[int] | list[str] | usecols_func | None,
+) -> None | list[int] | list[str] | usecols_func:
     """
     Convert `usecols` into a compatible format for parsing in `parsers.py`.
 
@@ -174,7 +212,17 @@ def maybe_convert_usecols(usecols):
     return usecols
 
 
-def validate_freeze_panes(freeze_panes):
+@overload
+def validate_freeze_panes(freeze_panes: tuple[int, int]) -> Literal[True]:
+    ...
+
+
+@overload
+def validate_freeze_panes(freeze_panes: None) -> Literal[False]:
+    ...
+
+
+def validate_freeze_panes(freeze_panes: tuple[int, int] | None) -> bool:
     if freeze_panes is not None:
         if len(freeze_panes) == 2 and all(
             isinstance(item, int) for item in freeze_panes
@@ -191,7 +239,9 @@ def validate_freeze_panes(freeze_panes):
     return False
 
 
-def fill_mi_header(row, control_row):
+def fill_mi_header(
+    row: list[Hashable], control_row: list[bool]
+) -> tuple[list[Hashable], list[bool]]:
     """
     Forward fill blank entries in row but only inside the same parent index.
 
@@ -224,7 +274,9 @@ def fill_mi_header(row, control_row):
     return row, control_row
 
 
-def pop_header_name(row, index_col):
+def pop_header_name(
+    row: list[Hashable], index_col: int | Sequence[int]
+) -> tuple[Hashable | None, list[Hashable]]:
     """
     Pop the header name for MultiIndex parsing.
 
@@ -243,7 +295,12 @@ def pop_header_name(row, index_col):
         The original data row with the header name removed.
     """
     # Pop out header name and fill w/blank.
-    i = index_col if not is_list_like(index_col) else max(index_col)
+    if is_list_like(index_col):
+        assert isinstance(index_col, Iterable)
+        i = max(index_col)
+    else:
+        assert not isinstance(index_col, Iterable)
+        i = index_col
 
     header_name = row[i]
     header_name = None if header_name == "" else header_name