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 all commits
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
14 changes: 10 additions & 4 deletions pandas/io/excel/_odswriter.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,16 @@
from typing import (
Any,
DefaultDict,
Tuple,
cast,
)

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,9 +30,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 Expand Up @@ -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):
Expand Down
25 changes: 17 additions & 8 deletions pandas/io/excel/_openpyxl.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from typing import (
TYPE_CHECKING,
Any,
Tuple,
cast,
)

import numpy as np
Expand All @@ -13,6 +15,7 @@
ReadBuffer,
Scalar,
StorageOptions,
WriteExcelBuffer,
)
from pandas.compat._optional import import_optional_dependency

Expand All @@ -35,10 +38,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 +77,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 +220,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 +421,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 +461,7 @@ def write_cells(
self.sheets[sheet_name] = wks

if validate_freeze_panes(freeze_panes):
freeze_panes = cast(Tuple[int, int], freeze_panes)
wks.freeze_panes = wks.cell(
row=freeze_panes[0] + 1, column=freeze_panes[1] + 1
)
Expand Down
27 changes: 18 additions & 9 deletions pandas/io/excel/_xlsxwriter.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@
from typing import Any

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 Down Expand Up @@ -170,10 +174,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 +205,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
36 changes: 25 additions & 11 deletions pandas/io/excel/_xlwt.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,16 @@
from typing import (
TYPE_CHECKING,
Any,
Tuple,
cast,
)

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 +30,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 +63,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 +72,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 +89,7 @@ def write_cells(
self.sheets[sheet_name] = wks

if validate_freeze_panes(freeze_panes):
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])
Expand Down Expand Up @@ -111,7 +123,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 +162,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