Skip to content

Commit

Permalink
Polars spreadsheet type changes (#780)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
skrawcz authored Mar 25, 2024
1 parent 31da460 commit 4858121
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 16 deletions.
2 changes: 2 additions & 0 deletions examples/polars/materialization/requirements.txt
Original file line number Diff line number Diff line change
@@ -1,2 +1,4 @@
polars
sf-hamilton
xlsx2csv # for excel data loader
xlsxwriter # Excel export requires 'xlsxwriter'
35 changes: 19 additions & 16 deletions hamilton/plugins/polars_extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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[
Expand Down
8 changes: 8 additions & 0 deletions tests/plugins/test_polars_extensions.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import io
import pathlib
import sys
import typing

import polars as pl # isort: skip
import pytest # isort: skip
Expand Down Expand Up @@ -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]

0 comments on commit 4858121

Please sign in to comment.