Skip to content

Commit 6320e98

Browse files
bashtagebashtage
committed
TST: Add tests for styler
2 parents 3085bb9 + 28fd624 commit 6320e98

24 files changed

+1416
-563
lines changed

‎pandas-stubs/_typing.pyi

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ from typing import (
2020
Optional,
2121
Protocol,
2222
Sequence,
23+
TypedDict,
2324
TypeVar,
2425
Union,
2526
)
@@ -244,4 +245,14 @@ FileWriteMode = Literal[
244245
]
245246
ColspaceArgType = str | int | Sequence[int | str] | Mapping[Hashable, str | int]
246247

248+
class StyleExportDict(TypedDict, total=False):
249+
apply: Any
250+
table_attributes: Any
251+
table_styles: Any
252+
hide_index: bool
253+
hide_columns: bool
254+
hide_index_names: bool
255+
hide_column_names: bool
256+
css: dict[str, str | int]
257+
247258
__all__ = ["npt", "type_t"]

‎pandas-stubs/core/frame.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1216,7 +1216,7 @@ class DataFrame(NDFrame, OpsMixin):
12161216
by: _str | ListLike | None = ...,
12171217
ax: PlotAxes | None = ...,
12181218
fontsize: float | _str | None = ...,
1219-
rot: int = ...,
1219+
rot: float = ...,
12201220
grid: _bool = ...,
12211221
figsize: tuple[float, float] | None = ...,
12221222
layout: tuple[int, int] | None = ...,

‎pandas-stubs/core/generic.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ class NDFrame(PandasObject, indexing.IndexingMixin):
160160
def to_sql(
161161
self,
162162
name: _str,
163-
con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
163+
con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
164164
schema: _str | None = ...,
165165
if_exists: Literal["fail", "replace", "append"] = ...,
166166
index: _bool = ...,

‎pandas-stubs/core/groupby/generic.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ class DataFrameGroupBy(GroupBy):
171171
grouped: DataFrame,
172172
subplots: bool = ...,
173173
column: str | Sequence | None = ...,
174-
fontsize: int | str = ...,
174+
fontsize: float | str = ...,
175175
rot: float = ...,
176176
grid: bool = ...,
177177
ax: PlotAxes | None = ...,

‎pandas-stubs/io/formats/style.pyi

Lines changed: 90 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,12 @@
11
from typing import (
22
Any,
33
Callable,
4-
Hashable,
54
Literal,
65
Sequence,
7-
TypeVar,
86
overload,
97
)
108

9+
from matplotlib.colors import Colormap
1110
import numpy as np
1211
from pandas.core.frame import DataFrame
1312
from pandas.core.series import Series
@@ -20,20 +19,21 @@ from pandas._typing import (
2019
Level,
2120
Scalar,
2221
WriteBuffer,
22+
WriteExcelBuffer,
2323
npt,
2424
)
2525

26+
from pandas.io.excel import ExcelWriter
2627
from pandas.io.formats.style_render import (
2728
CSSProperties,
2829
CSSStyles,
2930
ExtFormatter,
31+
StyleExportDict,
3032
StylerRenderer,
3133
Subset,
3234
)
3335

34-
_StylerT = TypeVar("_StylerT", bound=Styler)
35-
36-
class Styler(StylerRenderer):
36+
class Styler(StylerRenderer[Styler]):
3737
def __init__(
3838
self,
3939
data: DataFrame | Series,
@@ -59,46 +59,91 @@ class Styler(StylerRenderer):
5959
) -> Styler: ...
6060
def to_excel(
6161
self,
62-
excel_writer,
62+
excel_writer: FilePath | WriteExcelBuffer | ExcelWriter,
6363
sheet_name: str = ...,
6464
na_rep: str = ...,
6565
float_format: str | None = ...,
66-
columns: Sequence[Hashable] | None = ...,
67-
header: Sequence[Hashable] | bool = ...,
66+
columns: list[HashableT] | None = ...,
67+
header: list[HashableT] | bool = ...,
6868
index: bool = ...,
6969
index_label: IndexLabel | None = ...,
7070
startrow: int = ...,
7171
startcol: int = ...,
72-
engine: str | None = ...,
72+
engine: Literal["openpyxl", "xlsxwriter"] | None = ...,
7373
merge_cells: bool = ...,
7474
encoding: str | None = ...,
7575
inf_rep: str = ...,
7676
verbose: bool = ...,
7777
freeze_panes: tuple[int, int] | None = ...,
78+
# TODO: Listed in docs but not in function decl
79+
# storage_options: StorageOptions = ...,
7880
) -> None: ...
81+
@overload
7982
def to_latex(
8083
self,
81-
buf: FilePath | WriteBuffer[str] | None = ...,
84+
buf: FilePath | WriteBuffer[str],
8285
*,
8386
column_format: str | None = ...,
8487
position: str | None = ...,
85-
position_float: str | None = ...,
88+
position_float: Literal["centering", "raggedleft", "raggedright"] | None = ...,
8689
hrules: bool | None = ...,
87-
clines: str | None = ...,
90+
clines: Literal["all;data", "all;index", "skip-last;data", "skip-last;index"]
91+
| None = ...,
8892
label: str | None = ...,
89-
caption: str | tuple | None = ...,
93+
caption: str | tuple[str, str] | None = ...,
94+
sparse_index: bool | None = ...,
95+
sparse_columns: bool | None = ...,
96+
multirow_align: Literal["c", "t", "b", "naive"] | None = ...,
97+
multicol_align: Literal["r", "c", "l", "naive-l", "naive-r"] | None = ...,
98+
siunitx: bool = ...,
99+
environment: str | None = ...,
100+
encoding: str | None = ...,
101+
convert_css: bool = ...,
102+
) -> None: ...
103+
@overload
104+
def to_latex(
105+
self,
106+
buf: None = ...,
107+
*,
108+
column_format: str | None = ...,
109+
position: str | None = ...,
110+
position_float: Literal["centering", "raggedleft", "raggedright"] | None = ...,
111+
hrules: bool | None = ...,
112+
clines: Literal["all;data", "all;index", "skip-last;data", "skip-last;index"]
113+
| None = ...,
114+
label: str | None = ...,
115+
caption: str | tuple[str, str] | None = ...,
90116
sparse_index: bool | None = ...,
91117
sparse_columns: bool | None = ...,
92-
multirow_align: str | None = ...,
93-
multicol_align: str | None = ...,
118+
multirow_align: Literal["c", "t", "b", "naive"] | None = ...,
119+
multicol_align: Literal["r", "c", "l", "naive-l", "naive-r"] | None = ...,
94120
siunitx: bool = ...,
95121
environment: str | None = ...,
96122
encoding: str | None = ...,
97123
convert_css: bool = ...,
98-
): ...
124+
) -> str: ...
125+
@overload
126+
def to_html(
127+
self,
128+
buf: FilePath | WriteBuffer[str],
129+
*,
130+
table_uuid: str | None = ...,
131+
table_attributes: str | None = ...,
132+
sparse_index: bool | None = ...,
133+
sparse_columns: bool | None = ...,
134+
bold_headers: bool = ...,
135+
caption: str | None = ...,
136+
max_rows: int | None = ...,
137+
max_columns: int | None = ...,
138+
encoding: str | None = ...,
139+
doctype_html: bool = ...,
140+
exclude_styles: bool = ...,
141+
**kwargs: Any,
142+
) -> None: ...
143+
@overload
99144
def to_html(
100145
self,
101-
buf: FilePath | WriteBuffer[str] | None = ...,
146+
buf: None = ...,
102147
*,
103148
table_uuid: str | None = ...,
104149
table_attributes: str | None = ...,
@@ -112,7 +157,7 @@ class Styler(StylerRenderer):
112157
doctype_html: bool = ...,
113158
exclude_styles: bool = ...,
114159
**kwargs: Any,
115-
): ...
160+
) -> str: ...
116161
def set_td_classes(self, classes: DataFrame) -> Styler: ...
117162
def __copy__(self) -> Styler: ...
118163
def __deepcopy__(self, memo) -> Styler: ...
@@ -135,7 +180,7 @@ class Styler(StylerRenderer):
135180
) -> Styler: ...
136181
def apply_index(
137182
self,
138-
func: Callable[[Series], npt.NDArray[np.str_]],
183+
func: Callable[[Series], npt.NDArray[np.str_] | list[str] | Series[str]],
139184
axis: int | str = ...,
140185
level: Level | list[Level] | None = ...,
141186
**kwargs: Any,
@@ -153,8 +198,8 @@ class Styler(StylerRenderer):
153198
# def where(self, cond: Callable, value: str, other: str | None = ..., subset: Subset | None = ..., **kwargs) -> Styler: ...
154199
# def set_precision(self, precision: int) -> StylerRenderer: ...
155200
def set_table_attributes(self, attributes: str) -> Styler: ...
156-
def export(self) -> dict[str, Any]: ...
157-
def use(self, styles: dict[str, Any]) -> Styler: ...
201+
def export(self) -> StyleExportDict: ...
202+
def use(self, styles: StyleExportDict) -> Styler: ...
158203
uuid: Any # Incomplete
159204
def set_uuid(self, uuid: str) -> Styler: ...
160205
caption: Any # Incomplete
@@ -174,6 +219,7 @@ class Styler(StylerRenderer):
174219
) -> Styler: ...
175220
# def set_na_rep(self, na_rep: str) -> StylerRenderer: ...
176221
# def hide_index(self, subset: Subset | None = ..., level: Level | list[Level] | None = ..., names: bool = ...,) -> Styler: ...
222+
# TODO: Not in docs? Should it be?
177223
def hide_columns(
178224
self,
179225
subset: Subset | None = ...,
@@ -189,26 +235,38 @@ class Styler(StylerRenderer):
189235
) -> Styler: ...
190236
def background_gradient(
191237
self,
192-
cmap: str = ...,
238+
cmap: str | Colormap = ...,
193239
low: float = ...,
194240
high: float = ...,
195241
axis: Axis | None = ...,
196242
subset: Subset | None = ...,
197243
text_color_threshold: float = ...,
198244
vmin: float | None = ...,
199245
vmax: float | None = ...,
200-
gmap: Sequence | None = ...,
246+
gmap: Sequence[float]
247+
| Sequence[Sequence[float]]
248+
| npt.NDArray
249+
| DataFrame
250+
| Series
251+
| None = ...,
201252
) -> Styler: ...
202253
def text_gradient(
203254
self,
204-
cmap: str = ...,
255+
cmap: str | Colormap = ...,
205256
low: float = ...,
206257
high: float = ...,
207258
axis: Axis | None = ...,
208259
subset: Subset | None = ...,
260+
# In docs but not in function declaration
261+
# text_color_threshold: float
209262
vmin: float | None = ...,
210263
vmax: float | None = ...,
211-
gmap: Sequence | None = ...,
264+
gmap: Sequence[float]
265+
| Sequence[Sequence[float]]
266+
| npt.NDArray
267+
| DataFrame
268+
| Series
269+
| None = ...,
212270
) -> Styler: ...
213271
def set_properties(
214272
self, subset: Subset | None = ..., **kwargs: str | int
@@ -218,11 +276,13 @@ class Styler(StylerRenderer):
218276
subset: Subset | None = ...,
219277
axis: Axis | None = ...,
220278
*,
221-
color: str | list | tuple | None = ...,
222-
cmap: Any | None = ...,
279+
color: str | list[str] | tuple[str, str] | None = ...,
280+
cmap: str | Colormap = ...,
223281
width: float = ...,
224282
height: float = ...,
225-
align: str | float | Callable = ...,
283+
align: Literal["left", "right", "zero", "mid", "mean"]
284+
| float
285+
| Callable[[Series | npt.NDArray | DataFrame], float] = ...,
226286
vmin: float | None = ...,
227287
vmax: float | None = ...,
228288
props: str = ...,
@@ -276,10 +336,10 @@ class Styler(StylerRenderer):
276336
searchpath: str | list[str],
277337
html_table: str | None = ...,
278338
html_style: str | None = ...,
279-
) -> _StylerT: ...
339+
) -> type[Styler]: ...
280340
def pipe(
281341
self,
282-
func: Callable[[Styler], Styler] | tuple[Callable[[Styler], Styler], str],
342+
func: Callable | tuple[Callable, str],
283343
*args: Any,
284344
**kwargs: Any,
285345
) -> Styler: ...

‎pandas-stubs/io/formats/style_render.pyi

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
from typing import (
22
Any,
33
Callable,
4+
Generic,
45
Literal,
56
Optional,
6-
Sequence,
77
TypedDict,
8+
TypeVar,
89
Union,
910
)
1011

@@ -26,11 +27,23 @@ class CSSDict(TypedDict):
2627
selector: str
2728
props: CSSProperties
2829

30+
class StyleExportDict(TypedDict, total=False):
31+
apply: Any
32+
table_attributes: Any
33+
table_styles: Any
34+
hide_index: bool
35+
hide_columns: bool
36+
hide_index_names: bool
37+
hide_column_names: bool
38+
css: dict[str, str | int]
39+
2940
CSSStyles = list[CSSDict]
3041
Subset = Union[slice, list[HashableT], Index]
3142

32-
class StylerRenderer:
33-
loader: jinja2.loaders.BaseLoader
43+
_StylerT = TypeVar("_StylerT", bound=StylerRenderer)
44+
45+
class StylerRenderer(Generic[_StylerT]):
46+
loader: jinja2.loaders.PackageLoader
3447
env: jinja2.environment.Environment
3548
template_html: jinja2.environment.Template
3649
template_html_table: jinja2.environment.Template
@@ -46,7 +59,7 @@ class StylerRenderer:
4659
thousands: str | None = ...,
4760
escape: str | None = ...,
4861
hyperlinks: Literal["html", "latex"] | None = ...,
49-
) -> StylerRenderer: ...
62+
) -> _StylerT: ...
5063
def format_index(
5164
self,
5265
formatter: ExtFormatter | None = ...,
@@ -58,4 +71,4 @@ class StylerRenderer:
5871
thousands: str | None = ...,
5972
escape: str | None = ...,
6073
hyperlinks: Literal["html", "latex"] | None = ...,
61-
) -> StylerRenderer: ...
74+
) -> _StylerT: ...

0 commit comments

Comments
 (0)