From 386c3bc9c5c7cbbd1c5d0cb0a760fc2073445ace Mon Sep 17 00:00:00 2001 From: phofl Date: Fri, 12 Nov 2021 21:35:59 +0100 Subject: [PATCH 1/4] TYP: type excel util module --- pandas/io/excel/_util.py | 47 +++++++++++++++++++++++++++++++++++----- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/pandas/io/excel/_util.py b/pandas/io/excel/_util.py index 66a66fbbcd78a..5ab36e0079794 100644 --- a/pandas/io/excel/_util.py +++ b/pandas/io/excel/_util.py @@ -2,9 +2,13 @@ from typing import ( Any, + Callable, + Iterable, MutableMapping, + overload, ) +from pandas._typing import Scalar from pandas.compat._optional import import_optional_dependency from pandas.core.dtypes.common import ( @@ -31,7 +35,7 @@ def register_writer(klass): _writers[engine_name] = klass -def get_default_engine(ext, mode="reader"): +def get_default_engine(ext: str, mode: str = "reader") -> str: """ Return the default reader/writer for the given extension. @@ -73,7 +77,7 @@ def get_default_engine(ext, mode="reader"): return _default_readers[ext] -def get_writer(engine_name): +def get_writer(engine_name: str): try: return _writers[engine_name] except KeyError as err: @@ -145,7 +149,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: Callable) -> Callable: + ... + + +@overload +def maybe_convert_usecols(usecols: None) -> None: + ... + + +def maybe_convert_usecols( + usecols: str | list[int] | list[str] | Callable | None, +) -> None | list[int] | list[str] | Callable: """ Convert `usecols` into a compatible format for parsing in `parsers.py`. @@ -191,7 +217,9 @@ def validate_freeze_panes(freeze_panes): return False -def fill_mi_header(row, control_row): +def fill_mi_header( + row: list[Scalar], control_row: list[bool] +) -> tuple[list[Scalar], list[bool]]: """ Forward fill blank entries in row but only inside the same parent index. @@ -224,7 +252,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[Scalar], index_col: int | list[int] +) -> tuple[Scalar | None, list[Scalar]]: """ Pop the header name for MultiIndex parsing. @@ -243,7 +273,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 From bbeb329d1a2d411ea5edd5a469b464f4e18b6c1a Mon Sep 17 00:00:00 2001 From: phofl Date: Wed, 22 Dec 2021 12:42:03 +0100 Subject: [PATCH 2/4] Improve typing --- pandas/io/excel/_base.py | 2 +- pandas/io/excel/_util.py | 46 +++++++++++++++++++++++++++++----------- 2 files changed, 35 insertions(+), 13 deletions(-) diff --git a/pandas/io/excel/_base.py b/pandas/io/excel/_base.py index a5db36cee4254..2ff3360d0b808 100644 --- a/pandas/io/excel/_base.py +++ b/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 diff --git a/pandas/io/excel/_util.py b/pandas/io/excel/_util.py index 5ab36e0079794..9824d30eaaffa 100644 --- a/pandas/io/excel/_util.py +++ b/pandas/io/excel/_util.py @@ -1,14 +1,17 @@ from __future__ import annotations from typing import ( + TYPE_CHECKING, Any, Callable, + Hashable, Iterable, + Literal, MutableMapping, + Type, overload, ) -from pandas._typing import Scalar from pandas.compat._optional import import_optional_dependency from pandas.core.dtypes.common import ( @@ -16,10 +19,15 @@ is_list_like, ) -_writers: MutableMapping[str, str] = {} +if TYPE_CHECKING: + from pandas.io.excel._base import ExcelWriter + ExcelWriter_t = Type[ExcelWriter] -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. @@ -32,6 +40,8 @@ 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 @@ -77,7 +87,7 @@ def get_default_engine(ext: str, mode: str = "reader") -> str: return _default_readers[ext] -def get_writer(engine_name: str): +def get_writer(engine_name: str) -> ExcelWriter_t: try: return _writers[engine_name] except KeyError as err: @@ -160,7 +170,9 @@ def maybe_convert_usecols(usecols: list[str]) -> list[str]: @overload -def maybe_convert_usecols(usecols: Callable) -> Callable: +def maybe_convert_usecols( + usecols: Callable[[Hashable], object] +) -> Callable[[Hashable], object]: ... @@ -170,8 +182,8 @@ def maybe_convert_usecols(usecols: None) -> None: def maybe_convert_usecols( - usecols: str | list[int] | list[str] | Callable | None, -) -> None | list[int] | list[str] | Callable: + usecols: str | list[int] | list[str] | Callable[[Hashable], object] | None, +) -> None | list[int] | list[str] | Callable[[Hashable], object]: """ Convert `usecols` into a compatible format for parsing in `parsers.py`. @@ -200,7 +212,17 @@ def maybe_convert_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 @@ -218,8 +240,8 @@ def validate_freeze_panes(freeze_panes): def fill_mi_header( - row: list[Scalar], control_row: list[bool] -) -> tuple[list[Scalar], list[bool]]: + 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. @@ -253,8 +275,8 @@ def fill_mi_header( def pop_header_name( - row: list[Scalar], index_col: int | list[int] -) -> tuple[Scalar | None, list[Scalar]]: + row: list[Hashable], index_col: int | list[int] +) -> tuple[Hashable | None, list[Hashable]]: """ Pop the header name for MultiIndex parsing. From b40fa5a42c7e5f64113395450926d1ac14851a8c Mon Sep 17 00:00:00 2001 From: phofl Date: Wed, 22 Dec 2021 16:46:16 +0100 Subject: [PATCH 3/4] Adress comments --- pandas/io/excel/_util.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/pandas/io/excel/_util.py b/pandas/io/excel/_util.py index 9824d30eaaffa..60b63a3d08394 100644 --- a/pandas/io/excel/_util.py +++ b/pandas/io/excel/_util.py @@ -8,7 +8,7 @@ Iterable, Literal, MutableMapping, - Type, + TypeVar, overload, ) @@ -22,7 +22,8 @@ if TYPE_CHECKING: from pandas.io.excel._base import ExcelWriter - ExcelWriter_t = Type[ExcelWriter] + ExcelWriter_t = type[ExcelWriter] + usecols_func = TypeVar("usecols_func", bound=Callable[[Hashable], object]) _writers: MutableMapping[str, ExcelWriter_t] = {} @@ -45,7 +46,7 @@ def register_writer(klass: ExcelWriter_t) -> None: _writers[engine_name] = klass -def get_default_engine(ext: str, mode: str = "reader") -> str: +def get_default_engine(ext: str, mode: Literal["reader", "writer"] = "reader") -> str: """ Return the default reader/writer for the given extension. @@ -170,9 +171,7 @@ def maybe_convert_usecols(usecols: list[str]) -> list[str]: @overload -def maybe_convert_usecols( - usecols: Callable[[Hashable], object] -) -> Callable[[Hashable], object]: +def maybe_convert_usecols(usecols: usecols_func) -> usecols_func: ... @@ -182,8 +181,8 @@ def maybe_convert_usecols(usecols: None) -> None: def maybe_convert_usecols( - usecols: str | list[int] | list[str] | Callable[[Hashable], object] | None, -) -> None | list[int] | list[str] | Callable[[Hashable], object]: + 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`. From f0835f735b9945b820c9800a9015055ea8fdc5c5 Mon Sep 17 00:00:00 2001 From: phofl Date: Wed, 22 Dec 2021 22:27:11 +0100 Subject: [PATCH 4/4] Change list to sequence --- pandas/io/excel/_util.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pandas/io/excel/_util.py b/pandas/io/excel/_util.py index 60b63a3d08394..6be5ef0f64e16 100644 --- a/pandas/io/excel/_util.py +++ b/pandas/io/excel/_util.py @@ -8,6 +8,7 @@ Iterable, Literal, MutableMapping, + Sequence, TypeVar, overload, ) @@ -274,7 +275,7 @@ def fill_mi_header( def pop_header_name( - row: list[Hashable], index_col: int | list[int] + row: list[Hashable], index_col: int | Sequence[int] ) -> tuple[Hashable | None, list[Hashable]]: """ Pop the header name for MultiIndex parsing.