From 9fe6ab1ba788092e8e1df3d12d3cb10be44b65c8 Mon Sep 17 00:00:00 2001
From: Alexander Beedie <alexander-beedie@users.noreply.github.com>
Date: Thu, 25 Jul 2024 17:29:27 +0400
Subject: [PATCH] fix(python): Improve default `write_excel` int/float format
 when using an explicit "table_style"

---
 py-polars/polars/dataframe/frame.py           |  1 +
 .../polars/io/spreadsheet/_write_utils.py     | 19 ++++++++++++-------
 py-polars/tests/unit/io/test_spreadsheet.py   |  2 +-
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py
index 6c26ed7feb2cb..cf23828c420b2 100644
--- a/py-polars/polars/dataframe/frame.py
+++ b/py-polars/polars/dataframe/frame.py
@@ -3186,6 +3186,7 @@ def write_excel(
             dtype_formats=dtype_formats,
             header_format=header_format,
             float_precision=float_precision,
+            table_style=table_style,
             row_totals=row_totals,
             sparklines=sparklines,
             formulas=formulas,
diff --git a/py-polars/polars/io/spreadsheet/_write_utils.py b/py-polars/polars/io/spreadsheet/_write_utils.py
index ab97658aa2065..a88a33074406f 100644
--- a/py-polars/polars/io/spreadsheet/_write_utils.py
+++ b/py-polars/polars/io/spreadsheet/_write_utils.py
@@ -48,8 +48,6 @@ def _cluster(iterable: Iterable[Any], n: int = 2) -> Iterable[Any]:
     Date: "yyyy-mm-dd;@",
     Time: "hh:mm:ss;@",
 }
-for tp in INTEGER_DTYPES:
-    _XL_DEFAULT_DTYPE_FORMATS_[tp] = _XL_DEFAULT_INTEGER_FORMAT_
 
 
 class _XLFormatCache:
@@ -331,6 +329,7 @@ def _xl_setup_table_columns(
     formulas: dict[str, str | dict[str, str]] | None = None,
     row_totals: RowTotalsDefinition | None = None,
     float_precision: int = 3,
+    table_style: dict[str, Any] | str | None = None,
 ) -> tuple[list[dict[str, Any]], dict[str | tuple[str, ...], str], DataFrame]:
     """Setup and unify all column-related formatting/defaults."""
 
@@ -429,12 +428,18 @@ def _map_str(s: Series) -> Series:
     fmt_default = format_cache.get({"valign": "vcenter"})
 
     # default float format
+    if table_style is None:
+        int_base_fmt = _XL_DEFAULT_INTEGER_FORMAT_
+        flt_base_fmt = _XL_DEFAULT_FLOAT_FORMAT_
+    else:
+        int_base_fmt = _XL_DEFAULT_INTEGER_FORMAT_.split(";", 1)[0]
+        flt_base_fmt = _XL_DEFAULT_FLOAT_FORMAT_.split(";", 1)[0]
+
+    for tp in INTEGER_DTYPES:
+        _XL_DEFAULT_DTYPE_FORMATS_[tp] = int_base_fmt
+
     zeros = "0" * float_precision
-    fmt_float = (
-        _XL_DEFAULT_INTEGER_FORMAT_
-        if not zeros
-        else _XL_DEFAULT_FLOAT_FORMAT_.replace(".000", f".{zeros}")
-    )
+    fmt_float = int_base_fmt if not zeros else flt_base_fmt.replace(".000", f".{zeros}")
 
     # assign default dtype formats
     for tp, fmt in _XL_DEFAULT_DTYPE_FORMATS_.items():
diff --git a/py-polars/tests/unit/io/test_spreadsheet.py b/py-polars/tests/unit/io/test_spreadsheet.py
index 744c029ae0599..d1c3aae421865 100644
--- a/py-polars/tests/unit/io/test_spreadsheet.py
+++ b/py-polars/tests/unit/io/test_spreadsheet.py
@@ -545,7 +545,7 @@ def test_read_excel_all_sheets_with_sheet_name(path_xlsx: Path, engine: str) ->
         # basic formatting
         {
             "autofit": True,
-            "table_style": "Table Style Light 16",
+            "table_style": "Table Style Dark 2",
             "column_totals": True,
             "float_precision": 0,
         },