Skip to content

Commit

Permalink
TYP: Add types for excel writer classes (pandas-dev#45246)
Browse files Browse the repository at this point in the history
  • Loading branch information
phofl authored Jan 10, 2022
1 parent 37c3343 commit d3c62ad
Show file tree
Hide file tree
Showing 5 changed files with 72 additions and 34 deletions.
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

0 comments on commit d3c62ad

Please sign in to comment.