3838from pandas ._libs .tslib import format_array_from_datetime
3939from pandas ._libs .tslibs import NaT , Timedelta , Timestamp , iNaT
4040from pandas ._libs .tslibs .nattype import NaTType
41- from pandas ._typing import FilePathOrBuffer
41+ from pandas ._typing import FilePathOrBuffer , Label
4242from pandas .errors import AbstractMethodError
4343
4444from pandas .core .dtypes .common import (
7777 List [Callable ], Tuple [Callable , ...], Mapping [Union [str , int ], Callable ]
7878]
7979FloatFormatType = Union [str , Callable , "EngFormatter" ]
80+ ColspaceType = Mapping [Label , Union [str , int ]]
81+ ColspaceArgType = Union [
82+ str , int , Sequence [Union [str , int ]], Mapping [Label , Union [str , int ]],
83+ ]
8084
8185common_docstring = """
8286 Parameters
@@ -530,11 +534,13 @@ class DataFrameFormatter(TableFormatter):
530534 __doc__ = __doc__ if __doc__ else ""
531535 __doc__ += common_docstring + return_docstring
532536
537+ col_space : ColspaceType
538+
533539 def __init__ (
534540 self ,
535541 frame : "DataFrame" ,
536542 columns : Optional [Sequence [str ]] = None ,
537- col_space : Optional [Union [ str , int ] ] = None ,
543+ col_space : Optional [ColspaceArgType ] = None ,
538544 header : Union [bool , Sequence [str ]] = True ,
539545 index : bool = True ,
540546 na_rep : str = "NaN" ,
@@ -574,7 +580,27 @@ def __init__(
574580 )
575581 self .na_rep = na_rep
576582 self .decimal = decimal
577- self .col_space = col_space
583+ if col_space is None :
584+ self .col_space = {}
585+ elif isinstance (col_space , (int , str )):
586+ self .col_space = {"" : col_space }
587+ self .col_space .update ({column : col_space for column in self .frame .columns })
588+ elif isinstance (col_space , dict ):
589+ for column in col_space .keys ():
590+ if column not in self .frame .columns and column != "" :
591+ raise ValueError (
592+ f"Col_space is defined for an unknown column: { column } "
593+ )
594+ self .col_space = col_space
595+ else :
596+ col_space = cast (Sequence , col_space )
597+ if len (frame .columns ) != len (col_space ):
598+ raise ValueError (
599+ f"Col_space length({ len (col_space )} ) should match "
600+ f"DataFrame number of columns({ len (frame .columns )} )"
601+ )
602+ self .col_space = dict (zip (self .frame .columns , col_space ))
603+
578604 self .header = header
579605 self .index = index
580606 self .line_width = line_width
@@ -702,7 +728,7 @@ def _to_str_columns(self) -> List[List[str]]:
702728 """
703729 # this method is not used by to_html where self.col_space
704730 # could be a string so safe to cast
705- self . col_space = cast (int , self .col_space )
731+ col_space = { k : cast (int , v ) for k , v in self .col_space . items ()}
706732
707733 frame = self .tr_frame
708734 # may include levels names also
@@ -714,10 +740,7 @@ def _to_str_columns(self) -> List[List[str]]:
714740 for i , c in enumerate (frame ):
715741 fmt_values = self ._format_col (i )
716742 fmt_values = _make_fixed_width (
717- fmt_values ,
718- self .justify ,
719- minimum = (self .col_space or 0 ),
720- adj = self .adj ,
743+ fmt_values , self .justify , minimum = col_space .get (c , 0 ), adj = self .adj ,
721744 )
722745 stringified .append (fmt_values )
723746 else :
@@ -741,7 +764,7 @@ def _to_str_columns(self) -> List[List[str]]:
741764 for i , c in enumerate (frame ):
742765 cheader = str_columns [i ]
743766 header_colwidth = max (
744- self . col_space or 0 , * (self .adj .len (x ) for x in cheader )
767+ col_space . get ( c , 0 ) , * (self .adj .len (x ) for x in cheader )
745768 )
746769 fmt_values = self ._format_col (i )
747770 fmt_values = _make_fixed_width (
@@ -932,7 +955,7 @@ def _format_col(self, i: int) -> List[str]:
932955 formatter ,
933956 float_format = self .float_format ,
934957 na_rep = self .na_rep ,
935- space = self .col_space ,
958+ space = self .col_space . get ( frame . columns [ i ]) ,
936959 decimal = self .decimal ,
937960 )
938961
@@ -1025,7 +1048,7 @@ def show_col_idx_names(self) -> bool:
10251048 def _get_formatted_index (self , frame : "DataFrame" ) -> List [str ]:
10261049 # Note: this is only used by to_string() and to_latex(), not by
10271050 # to_html(). so safe to cast col_space here.
1028- self . col_space = cast (int , self .col_space )
1051+ col_space = { k : cast (int , v ) for k , v in self .col_space . items ()}
10291052 index = frame .index
10301053 columns = frame .columns
10311054 fmt = self ._get_formatter ("__index__" )
@@ -1043,7 +1066,7 @@ def _get_formatted_index(self, frame: "DataFrame") -> List[str]:
10431066 fmt_index = [
10441067 tuple (
10451068 _make_fixed_width (
1046- list (x ), justify = "left" , minimum = ( self . col_space or 0 ), adj = self .adj
1069+ list (x ), justify = "left" , minimum = col_space . get ( "" , 0 ), adj = self .adj ,
10471070 )
10481071 )
10491072 for x in fmt_index
0 commit comments