Skip to content

Commit 7368b2a

Browse files
authored
CLN/TYP: pandas/io/formats/excel.py (#37862)
1 parent 7ceacb4 commit 7368b2a

File tree

1 file changed

+68
-68
lines changed

1 file changed

+68
-68
lines changed

pandas/io/formats/excel.py

+68-68
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,17 @@
55
from functools import reduce
66
import itertools
77
import re
8-
from typing import Callable, Dict, Mapping, Optional, Sequence, Union
8+
from typing import Callable, Dict, Iterable, Mapping, Optional, Sequence, Union, cast
99
import warnings
1010

1111
import numpy as np
1212

13+
from pandas._libs.lib import is_list_like
1314
from pandas._typing import Label, StorageOptions
1415
from pandas.util._decorators import doc
1516

1617
from pandas.core.dtypes import missing
1718
from pandas.core.dtypes.common import is_float, is_scalar
18-
from pandas.core.dtypes.generic import ABCIndex
1919

2020
from pandas import DataFrame, Index, MultiIndex, PeriodIndex
2121
from pandas.core import generic
@@ -31,7 +31,13 @@ class ExcelCell:
3131
__slots__ = __fields__
3232

3333
def __init__(
34-
self, row: int, col: int, val, style=None, mergestart=None, mergeend=None
34+
self,
35+
row: int,
36+
col: int,
37+
val,
38+
style=None,
39+
mergestart: Optional[int] = None,
40+
mergeend: Optional[int] = None,
3541
):
3642
self.row = row
3743
self.col = col
@@ -425,7 +431,7 @@ class ExcelFormatter:
425431
Format string for floating point numbers
426432
cols : sequence, optional
427433
Columns to write
428-
header : boolean or list of string, default True
434+
header : boolean or sequence of str, default True
429435
Write out column names. If a list of string is given it is
430436
assumed to be aliases for the column names
431437
index : boolean, default True
@@ -524,16 +530,15 @@ def _format_value(self, val):
524530
)
525531
return val
526532

527-
def _format_header_mi(self):
533+
def _format_header_mi(self) -> Iterable[ExcelCell]:
528534
if self.columns.nlevels > 1:
529535
if not self.index:
530536
raise NotImplementedError(
531537
"Writing to Excel with MultiIndex columns and no "
532538
"index ('index'=False) is not yet implemented."
533539
)
534540

535-
has_aliases = isinstance(self.header, (tuple, list, np.ndarray, ABCIndex))
536-
if not (has_aliases or self.header):
541+
if not (self._has_aliases or self.header):
537542
return
538543

539544
columns = self.columns
@@ -549,28 +554,30 @@ def _format_header_mi(self):
549554

550555
if self.merge_cells:
551556
# Format multi-index as a merged cells.
552-
for lnum in range(len(level_lengths)):
553-
name = columns.names[lnum]
554-
yield ExcelCell(lnum, coloffset, name, self.header_style)
557+
for lnum, name in enumerate(columns.names):
558+
yield ExcelCell(
559+
row=lnum,
560+
col=coloffset,
561+
val=name,
562+
style=self.header_style,
563+
)
555564

556565
for lnum, (spans, levels, level_codes) in enumerate(
557566
zip(level_lengths, columns.levels, columns.codes)
558567
):
559568
values = levels.take(level_codes)
560-
for i in spans:
561-
if spans[i] > 1:
562-
yield ExcelCell(
563-
lnum,
564-
coloffset + i + 1,
565-
values[i],
566-
self.header_style,
567-
lnum,
568-
coloffset + i + spans[i],
569-
)
570-
else:
571-
yield ExcelCell(
572-
lnum, coloffset + i + 1, values[i], self.header_style
573-
)
569+
for i, span_val in spans.items():
570+
spans_multiple_cells = span_val > 1
571+
yield ExcelCell(
572+
row=lnum,
573+
col=coloffset + i + 1,
574+
val=values[i],
575+
style=self.header_style,
576+
mergestart=lnum if spans_multiple_cells else None,
577+
mergeend=(
578+
coloffset + i + span_val if spans_multiple_cells else None
579+
),
580+
)
574581
else:
575582
# Format in legacy format with dots to indicate levels.
576583
for i, values in enumerate(zip(*level_strs)):
@@ -579,9 +586,8 @@ def _format_header_mi(self):
579586

580587
self.rowcounter = lnum
581588

582-
def _format_header_regular(self):
583-
has_aliases = isinstance(self.header, (tuple, list, np.ndarray, ABCIndex))
584-
if has_aliases or self.header:
589+
def _format_header_regular(self) -> Iterable[ExcelCell]:
590+
if self._has_aliases or self.header:
585591
coloffset = 0
586592

587593
if self.index:
@@ -590,17 +596,11 @@ def _format_header_regular(self):
590596
coloffset = len(self.df.index[0])
591597

592598
colnames = self.columns
593-
if has_aliases:
594-
# pandas\io\formats\excel.py:593: error: Argument 1 to "len"
595-
# has incompatible type "Union[Sequence[Optional[Hashable]],
596-
# bool]"; expected "Sized" [arg-type]
597-
if len(self.header) != len(self.columns): # type: ignore[arg-type]
598-
# pandas\io\formats\excel.py:602: error: Argument 1 to
599-
# "len" has incompatible type
600-
# "Union[Sequence[Optional[Hashable]], bool]"; expected
601-
# "Sized" [arg-type]
599+
if self._has_aliases:
600+
self.header = cast(Sequence, self.header)
601+
if len(self.header) != len(self.columns):
602602
raise ValueError(
603-
f"Writing {len(self.columns)} cols " # type: ignore[arg-type]
603+
f"Writing {len(self.columns)} cols "
604604
f"but got {len(self.header)} aliases"
605605
)
606606
else:
@@ -611,7 +611,7 @@ def _format_header_regular(self):
611611
self.rowcounter, colindex + coloffset, colname, self.header_style
612612
)
613613

614-
def _format_header(self):
614+
def _format_header(self) -> Iterable[ExcelCell]:
615615
if isinstance(self.columns, MultiIndex):
616616
gen = self._format_header_mi()
617617
else:
@@ -633,15 +633,14 @@ def _format_header(self):
633633
self.rowcounter += 1
634634
return itertools.chain(gen, gen2)
635635

636-
def _format_body(self):
636+
def _format_body(self) -> Iterable[ExcelCell]:
637637
if isinstance(self.df.index, MultiIndex):
638638
return self._format_hierarchical_rows()
639639
else:
640640
return self._format_regular_rows()
641641

642-
def _format_regular_rows(self):
643-
has_aliases = isinstance(self.header, (tuple, list, np.ndarray, ABCIndex))
644-
if has_aliases or self.header:
642+
def _format_regular_rows(self) -> Iterable[ExcelCell]:
643+
if self._has_aliases or self.header:
645644
self.rowcounter += 1
646645

647646
# output index and index_label?
@@ -678,9 +677,8 @@ def _format_regular_rows(self):
678677

679678
yield from self._generate_body(coloffset)
680679

681-
def _format_hierarchical_rows(self):
682-
has_aliases = isinstance(self.header, (tuple, list, np.ndarray, ABCIndex))
683-
if has_aliases or self.header:
680+
def _format_hierarchical_rows(self) -> Iterable[ExcelCell]:
681+
if self._has_aliases or self.header:
684682
self.rowcounter += 1
685683

686684
gcolidx = 0
@@ -723,40 +721,42 @@ def _format_hierarchical_rows(self):
723721
fill_value=levels._na_value,
724722
)
725723

726-
for i in spans:
727-
if spans[i] > 1:
728-
yield ExcelCell(
729-
self.rowcounter + i,
730-
gcolidx,
731-
values[i],
732-
self.header_style,
733-
self.rowcounter + i + spans[i] - 1,
734-
gcolidx,
735-
)
736-
else:
737-
yield ExcelCell(
738-
self.rowcounter + i,
739-
gcolidx,
740-
values[i],
741-
self.header_style,
742-
)
724+
for i, span_val in spans.items():
725+
spans_multiple_cells = span_val > 1
726+
yield ExcelCell(
727+
row=self.rowcounter + i,
728+
col=gcolidx,
729+
val=values[i],
730+
style=self.header_style,
731+
mergestart=(
732+
self.rowcounter + i + span_val - 1
733+
if spans_multiple_cells
734+
else None
735+
),
736+
mergeend=gcolidx if spans_multiple_cells else None,
737+
)
743738
gcolidx += 1
744739

745740
else:
746741
# Format hierarchical rows with non-merged values.
747742
for indexcolvals in zip(*self.df.index):
748743
for idx, indexcolval in enumerate(indexcolvals):
749744
yield ExcelCell(
750-
self.rowcounter + idx,
751-
gcolidx,
752-
indexcolval,
753-
self.header_style,
745+
row=self.rowcounter + idx,
746+
col=gcolidx,
747+
val=indexcolval,
748+
style=self.header_style,
754749
)
755750
gcolidx += 1
756751

757752
yield from self._generate_body(gcolidx)
758753

759-
def _generate_body(self, coloffset: int):
754+
@property
755+
def _has_aliases(self) -> bool:
756+
"""Whether the aliases for column names are present."""
757+
return is_list_like(self.header)
758+
759+
def _generate_body(self, coloffset: int) -> Iterable[ExcelCell]:
760760
if self.styler is None:
761761
styles = None
762762
else:
@@ -773,7 +773,7 @@ def _generate_body(self, coloffset: int):
773773
xlstyle = self.style_converter(";".join(styles[i, colidx]))
774774
yield ExcelCell(self.rowcounter + i, colidx + coloffset, val, xlstyle)
775775

776-
def get_formatted_cells(self):
776+
def get_formatted_cells(self) -> Iterable[ExcelCell]:
777777
for cell in itertools.chain(self._format_header(), self._format_body()):
778778
cell.val = self._format_value(cell.val)
779779
yield cell

0 commit comments

Comments
 (0)