5
5
from functools import reduce
6
6
import itertools
7
7
import re
8
- from typing import Callable , Dict , Mapping , Optional , Sequence , Union
8
+ from typing import Callable , Dict , Iterable , Mapping , Optional , Sequence , Union , cast
9
9
import warnings
10
10
11
11
import numpy as np
12
12
13
+ from pandas ._libs .lib import is_list_like
13
14
from pandas ._typing import Label , StorageOptions
14
15
from pandas .util ._decorators import doc
15
16
16
17
from pandas .core .dtypes import missing
17
18
from pandas .core .dtypes .common import is_float , is_scalar
18
- from pandas .core .dtypes .generic import ABCIndex
19
19
20
20
from pandas import DataFrame , Index , MultiIndex , PeriodIndex
21
21
from pandas .core import generic
@@ -31,7 +31,13 @@ class ExcelCell:
31
31
__slots__ = __fields__
32
32
33
33
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 ,
35
41
):
36
42
self .row = row
37
43
self .col = col
@@ -425,7 +431,7 @@ class ExcelFormatter:
425
431
Format string for floating point numbers
426
432
cols : sequence, optional
427
433
Columns to write
428
- header : boolean or list of string , default True
434
+ header : boolean or sequence of str , default True
429
435
Write out column names. If a list of string is given it is
430
436
assumed to be aliases for the column names
431
437
index : boolean, default True
@@ -524,16 +530,15 @@ def _format_value(self, val):
524
530
)
525
531
return val
526
532
527
- def _format_header_mi (self ):
533
+ def _format_header_mi (self ) -> Iterable [ ExcelCell ] :
528
534
if self .columns .nlevels > 1 :
529
535
if not self .index :
530
536
raise NotImplementedError (
531
537
"Writing to Excel with MultiIndex columns and no "
532
538
"index ('index'=False) is not yet implemented."
533
539
)
534
540
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 ):
537
542
return
538
543
539
544
columns = self .columns
@@ -549,28 +554,30 @@ def _format_header_mi(self):
549
554
550
555
if self .merge_cells :
551
556
# 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
+ )
555
564
556
565
for lnum , (spans , levels , level_codes ) in enumerate (
557
566
zip (level_lengths , columns .levels , columns .codes )
558
567
):
559
568
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
+ )
574
581
else :
575
582
# Format in legacy format with dots to indicate levels.
576
583
for i , values in enumerate (zip (* level_strs )):
@@ -579,9 +586,8 @@ def _format_header_mi(self):
579
586
580
587
self .rowcounter = lnum
581
588
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 :
585
591
coloffset = 0
586
592
587
593
if self .index :
@@ -590,17 +596,11 @@ def _format_header_regular(self):
590
596
coloffset = len (self .df .index [0 ])
591
597
592
598
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 ):
602
602
raise ValueError (
603
- f"Writing { len (self .columns )} cols " # type: ignore[arg-type]
603
+ f"Writing { len (self .columns )} cols "
604
604
f"but got { len (self .header )} aliases"
605
605
)
606
606
else :
@@ -611,7 +611,7 @@ def _format_header_regular(self):
611
611
self .rowcounter , colindex + coloffset , colname , self .header_style
612
612
)
613
613
614
- def _format_header (self ):
614
+ def _format_header (self ) -> Iterable [ ExcelCell ] :
615
615
if isinstance (self .columns , MultiIndex ):
616
616
gen = self ._format_header_mi ()
617
617
else :
@@ -633,15 +633,14 @@ def _format_header(self):
633
633
self .rowcounter += 1
634
634
return itertools .chain (gen , gen2 )
635
635
636
- def _format_body (self ):
636
+ def _format_body (self ) -> Iterable [ ExcelCell ] :
637
637
if isinstance (self .df .index , MultiIndex ):
638
638
return self ._format_hierarchical_rows ()
639
639
else :
640
640
return self ._format_regular_rows ()
641
641
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 :
645
644
self .rowcounter += 1
646
645
647
646
# output index and index_label?
@@ -678,9 +677,8 @@ def _format_regular_rows(self):
678
677
679
678
yield from self ._generate_body (coloffset )
680
679
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 :
684
682
self .rowcounter += 1
685
683
686
684
gcolidx = 0
@@ -723,40 +721,42 @@ def _format_hierarchical_rows(self):
723
721
fill_value = levels ._na_value ,
724
722
)
725
723
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
+ )
743
738
gcolidx += 1
744
739
745
740
else :
746
741
# Format hierarchical rows with non-merged values.
747
742
for indexcolvals in zip (* self .df .index ):
748
743
for idx , indexcolval in enumerate (indexcolvals ):
749
744
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 ,
754
749
)
755
750
gcolidx += 1
756
751
757
752
yield from self ._generate_body (gcolidx )
758
753
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 ]:
760
760
if self .styler is None :
761
761
styles = None
762
762
else :
@@ -773,7 +773,7 @@ def _generate_body(self, coloffset: int):
773
773
xlstyle = self .style_converter (";" .join (styles [i , colidx ]))
774
774
yield ExcelCell (self .rowcounter + i , colidx + coloffset , val , xlstyle )
775
775
776
- def get_formatted_cells (self ):
776
+ def get_formatted_cells (self ) -> Iterable [ ExcelCell ] :
777
777
for cell in itertools .chain (self ._format_header (), self ._format_body ()):
778
778
cell .val = self ._format_value (cell .val )
779
779
yield cell
0 commit comments