Polars spreadsheet type changes (#780)

* Fixes up type hints and imports for polars spreadsheet writer

In python 3.9 at least things weren't working. The type hints for
the spreadsheet writer use type aliases that require other imports
to introspect apparently. That is `type_hints = typing.get_type_hints(cls)`
was failing on that class. Adds test for this.

* Adds missing deps to polars materializers example

Author: skrawcz

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]