From ce4a3bfd79749350faf65038d22bec2dc50749d4 Mon Sep 17 00:00:00 2001 From: Arnab Animesh Das Date: Sat, 25 Jan 2025 11:33:48 +0530 Subject: [PATCH] fix(python): Consider the original dtypes when selecting columns in the `write_excel` function --- py-polars/polars/dataframe/frame.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py index 30fe3691635e..17e234b4a57f 100644 --- a/py-polars/polars/dataframe/frame.py +++ b/py-polars/polars/dataframe/frame.py @@ -3473,6 +3473,12 @@ def write_excel( wb, ws, can_close = _xl_setup_workbook(workbook, worksheet) df, is_empty = self, not len(self) + # The _xl_setup_table_columns function in the below section + # converts all collection types (e.g. List, Struct, Object) to strings + # Hence, we need to store the original schema so that it can be used + # when selecting columns using column selectors based on datatypes + df_ori_dtypes = df.head(0) + # setup table format/columns fmt_cache = _XLFormatCache(wb) column_formats = column_formats or {} @@ -3540,7 +3546,7 @@ def write_excel( elif isinstance(hidden_columns, str): hidden = {hidden_columns} else: - hidden = set(_expand_selectors(df, hidden_columns)) + hidden = set(_expand_selectors(df_ori_dtypes, hidden_columns)) # Autofit section needs to be present above column_widths section # to ensure that parameters provided in the column_widths section @@ -3558,7 +3564,7 @@ def write_excel( column_widths = dict.fromkeys(df.columns, column_widths) else: column_widths = _expand_selector_dicts( # type: ignore[assignment] - df, column_widths, expand_keys=True, expand_values=False + df_ori_dtypes, column_widths, expand_keys=True, expand_values=False ) column_widths = _unpack_multi_column_dict(column_widths or {}) # type: ignore[assignment]