Polars spreadsheet type changes

examples/polars/materialization/requirements.txt

@@ -1,2 +1,4 @@
 polars
 sf-hamilton
+xlsx2csv  # for excel data loader
+xlsxwriter  # Excel export requires 'xlsxwriter'

hamilton/plugins/polars_extensions.py

@@ -22,17 +22,6 @@
 except ImportError:
     Workbook = Type
 
-
-from polars.type_aliases import (
-    ColumnFormatDict,
-    ColumnTotalsDefinition,
-    ColumnWidthsDefinition,
-    ConditionalFormatDict,
-    OneOrMoreDataTypes,
-    RowTotalsDefinition,
-    SelectorType,
-)
-
 try:
     import polars as pl
     from polars import PolarsDataType
@@ -633,18 +622,32 @@ class PolarsSpreadsheetWriter(DataSaver):
     Should map to https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.DataFrame.write_excel.html
     """
 
+    # importing here because this is where it's used. Can move later.
+    # but yeah the polars type aliases weren't resolving well in python 3.9
+    # so stripped/reduced them appropriately.
+    from polars.datatypes import DataType, DataTypeClass
+    from polars.type_aliases import ColumnTotalsDefinition, RowTotalsDefinition
+
     workbook: Union[Workbook, BytesIO, Path, str]
     worksheet: Union[str, None] = None
     # kwargs:
     position: Union[Tuple[int, int], str] = "A1"
     table_style: Union[str, Dict[str, Any], None] = None
     table_name: Union[str, None] = None
-    column_formats: Union[ColumnFormatDict, None] = None
-    dtype_formats: Union[Dict[OneOrMoreDataTypes, str], None] = None
-    conditional_formats: Union[ConditionalFormatDict, None] = None
+    column_formats: Union[
+        Mapping[Union[str, Tuple[str, ...]], Union[str, Mapping[str, str]]], None
+    ] = None
+    dtype_formats: Union[Dict[Union[DataType, DataTypeClass], str], None] = None
+    conditional_formats: Union[
+        Mapping[
+            Union[str, Collection[str]],
+            Union[str, Union[Mapping[str, Any], Sequence[Union[str, Mapping[str, Any]]]]],
+        ],
+        None,
+    ] = None
     header_format: Union[Dict[str, Any], None] = None
     column_totals: Union[ColumnTotalsDefinition, None] = None
-    column_widths: Union[ColumnWidthsDefinition, None] = None
+    column_widths: Union[Mapping[str, Union[Tuple[str, ...], int]], int, None] = None
     row_totals: Union[RowTotalsDefinition, None] = None
     row_heights: Union[Dict[Union[int, Tuple[int, ...]], int], int, None] = None
     sparklines: Union[Dict[str, Union[Sequence[str], Dict[str, Any]]], None] = None
@@ -653,7 +656,7 @@ class PolarsSpreadsheetWriter(DataSaver):
     include_header: bool = True
     autofilter: bool = True
     autofit: bool = False
-    hidden_columns: Union[Sequence[str], SelectorType, None] = None
+    hidden_columns: Union[Sequence[str], str, None] = None
     hide_gridlines: bool = None
     sheet_zoom: Union[int, None] = None
     freeze_panes: Union[

tests/plugins/test_polars_extensions.py

@@ -1,5 +1,7 @@
+import io
 import pathlib
 import sys
+import typing
 
 import polars as pl  # isort: skip
 import pytest  # isort: skip
@@ -168,3 +170,9 @@ def test_polars_spreadsheet(df: pl.DataFrame, tmp_path: pathlib.Path) -> None:
     assert write_kwargs["include_header"] is True
     assert "raise_if_empty" in read_kwargs
     assert read_kwargs["raise_if_empty"] is True
+
+
+def test_getting_type_hints_spreadsheetwriter():
+    """Tests that types can be resolved at run time."""
+    type_hints = typing.get_type_hints(PolarsSpreadsheetWriter)
+    assert type_hints["workbook"] == typing.Union[typing.Type, io.BytesIO, pathlib.Path, str]