From 27debda0698232aae15d7c9ecaa5fd1675a21cf1 Mon Sep 17 00:00:00 2001 From: Patrick Date: Thu, 6 Jan 2022 14:36:57 +0100 Subject: [PATCH 1/2] TYP: Add types for excel writer classes --- pandas/io/excel/_base.py | 4 ++-- pandas/io/excel/_odswriter.py | 6 +++--- pandas/io/excel/_openpyxl.py | 24 +++++++++++++++-------- pandas/io/excel/_xlsxwriter.py | 23 +++++++++++++--------- pandas/io/excel/_xlwt.py | 35 +++++++++++++++++++++++----------- 5 files changed, 59 insertions(+), 33 deletions(-) diff --git a/pandas/io/excel/_base.py b/pandas/io/excel/_base.py index b490244f7f396..bf2d6f0b27e41 100644 --- a/pandas/io/excel/_base.py +++ b/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 diff --git a/pandas/io/excel/_odswriter.py b/pandas/io/excel/_odswriter.py index d4fe3683c907e..933d38561a2e3 100644 --- a/pandas/io/excel/_odswriter.py +++ b/pandas/io/excel/_odswriter.py @@ -8,7 +8,7 @@ ) import pandas._libs.json as json -from pandas._typing import StorageOptions +from pandas._typing import StorageOptions, FilePath, WriteExcelBuffer from pandas.io.excel._base import ExcelWriter from pandas.io.excel._util import ( @@ -24,9 +24,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, diff --git a/pandas/io/excel/_openpyxl.py b/pandas/io/excel/_openpyxl.py index 27c03d4a74bc1..e74d4e1a0aa8f 100644 --- a/pandas/io/excel/_openpyxl.py +++ b/pandas/io/excel/_openpyxl.py @@ -13,6 +13,7 @@ ReadBuffer, Scalar, StorageOptions, + WriteExcelBuffer, ) from pandas.compat._optional import import_optional_dependency @@ -35,10 +36,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 +75,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 +218,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 +419,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 +459,8 @@ def write_cells( self.sheets[sheet_name] = wks if validate_freeze_panes(freeze_panes): + # for mypy + assert freeze_panes is not None wks.freeze_panes = wks.cell( row=freeze_panes[0] + 1, column=freeze_panes[1] + 1 ) diff --git a/pandas/io/excel/_xlsxwriter.py b/pandas/io/excel/_xlsxwriter.py index 06c73f2c6199e..122235bc547fd 100644 --- a/pandas/io/excel/_xlsxwriter.py +++ b/pandas/io/excel/_xlsxwriter.py @@ -3,7 +3,7 @@ from typing import Any import pandas._libs.json as json -from pandas._typing import StorageOptions +from pandas._typing import StorageOptions, FilePath, WriteExcelBuffer from pandas.io.excel._base import ExcelWriter from pandas.io.excel._util import ( @@ -170,10 +170,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 +201,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) diff --git a/pandas/io/excel/_xlwt.py b/pandas/io/excel/_xlwt.py index a74c03f330cd9..c8c0689d938a1 100644 --- a/pandas/io/excel/_xlwt.py +++ b/pandas/io/excel/_xlwt.py @@ -6,7 +6,11 @@ ) 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 +28,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 +61,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 +70,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 +87,8 @@ def write_cells( self.sheets[sheet_name] = wks if validate_freeze_panes(freeze_panes): + # for mypy + assert freeze_panes is not None wks.set_panes_frozen(True) wks.set_horz_split_pos(freeze_panes[0]) wks.set_vert_split_pos(freeze_panes[1]) @@ -111,7 +122,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 +161,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 From 24a921caa771a11dd5f9a608c5b0f9d3b6ab02b9 Mon Sep 17 00:00:00 2001 From: phofl Date: Mon, 10 Jan 2022 02:46:28 +0100 Subject: [PATCH 2/2] Adress review --- pandas/io/excel/_odswriter.py | 10 ++++++++-- pandas/io/excel/_openpyxl.py | 5 +++-- pandas/io/excel/_xlsxwriter.py | 6 +++++- pandas/io/excel/_xlwt.py | 5 +++-- 4 files changed, 19 insertions(+), 7 deletions(-) diff --git a/pandas/io/excel/_odswriter.py b/pandas/io/excel/_odswriter.py index 933d38561a2e3..0f31991ee29e9 100644 --- a/pandas/io/excel/_odswriter.py +++ b/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, FilePath, WriteExcelBuffer +from pandas._typing import ( + FilePath, + StorageOptions, + WriteExcelBuffer, +) from pandas.io.excel._base import ExcelWriter from pandas.io.excel._util import ( @@ -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): diff --git a/pandas/io/excel/_openpyxl.py b/pandas/io/excel/_openpyxl.py index e74d4e1a0aa8f..88a25d1c1e6ef 100644 --- a/pandas/io/excel/_openpyxl.py +++ b/pandas/io/excel/_openpyxl.py @@ -4,6 +4,8 @@ from typing import ( TYPE_CHECKING, Any, + Tuple, + cast, ) import numpy as np @@ -459,8 +461,7 @@ def write_cells( self.sheets[sheet_name] = wks if validate_freeze_panes(freeze_panes): - # for mypy - assert freeze_panes is not None + freeze_panes = cast(Tuple[int, int], freeze_panes) wks.freeze_panes = wks.cell( row=freeze_panes[0] + 1, column=freeze_panes[1] + 1 ) diff --git a/pandas/io/excel/_xlsxwriter.py b/pandas/io/excel/_xlsxwriter.py index 122235bc547fd..49c87732f1429 100644 --- a/pandas/io/excel/_xlsxwriter.py +++ b/pandas/io/excel/_xlsxwriter.py @@ -3,7 +3,11 @@ from typing import Any import pandas._libs.json as json -from pandas._typing import StorageOptions, FilePath, WriteExcelBuffer +from pandas._typing import ( + FilePath, + StorageOptions, + WriteExcelBuffer, +) from pandas.io.excel._base import ExcelWriter from pandas.io.excel._util import ( diff --git a/pandas/io/excel/_xlwt.py b/pandas/io/excel/_xlwt.py index c8c0689d938a1..1ada0eb25f81c 100644 --- a/pandas/io/excel/_xlwt.py +++ b/pandas/io/excel/_xlwt.py @@ -3,6 +3,8 @@ from typing import ( TYPE_CHECKING, Any, + Tuple, + cast, ) import pandas._libs.json as json @@ -87,8 +89,7 @@ def write_cells( self.sheets[sheet_name] = wks if validate_freeze_panes(freeze_panes): - # for mypy - assert freeze_panes is not None + 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])