Skip to content

TYP: Add types for excel writer classes #45246

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jan 10, 2022
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions pandas/io/excel/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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/
Expand Down Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions pandas/io/excel/_odswriter.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
)

import pandas._libs.json as json
from pandas._typing import StorageOptions
from pandas._typing import StorageOptions, FilePath, WriteExcelBuffer
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isort complaints

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Had some issues with pre commit yesterday, will fix


from pandas.io.excel._base import ExcelWriter
from pandas.io.excel._util import (
Expand All @@ -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,
Expand Down
24 changes: 16 additions & 8 deletions pandas/io/excel/_openpyxl.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
ReadBuffer,
Scalar,
StorageOptions,
WriteExcelBuffer,
)
from pandas.compat._optional import import_optional_dependency

Expand All @@ -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,
Expand Down Expand Up @@ -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.
"""
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure if its been discussed, but i think we should move away from the assert-for-mypy pattern, since with cast/ignore mypy will alert us when it is no longer necessary

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will use cast then

wks.freeze_panes = wks.cell(
row=freeze_panes[0] + 1, column=freeze_panes[1] + 1
)
Expand Down
23 changes: 14 additions & 9 deletions pandas/io/excel/_xlsxwriter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand Down
35 changes: 24 additions & 11 deletions pandas/io/excel/_xlwt.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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,
Expand Down Expand Up @@ -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.
"""
Expand All @@ -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)

Expand All @@ -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])
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down