From ddb186af7eee185567a6b969306210f95f15d775 Mon Sep 17 00:00:00 2001 From: MomIsBestFriend <> Date: Mon, 23 Dec 2019 12:24:39 +0200 Subject: [PATCH] TYP: Type annotations in pandas/io/formats/style.py --- pandas/io/formats/style.py | 192 +++++++++++++++++++++---------------- 1 file changed, 111 insertions(+), 81 deletions(-) diff --git a/pandas/io/formats/style.py b/pandas/io/formats/style.py index d3a12ccb770489..b5c2dff1d27b79 100644 --- a/pandas/io/formats/style.py +++ b/pandas/io/formats/style.py @@ -1,6 +1,5 @@ """ -Module for applying conditional formatting to -DataFrames and Series. +Module for applying conditional formatting to DataFrames and Series. """ from collections import defaultdict @@ -8,7 +7,18 @@ import copy from functools import partial from itertools import product -from typing import Any, Callable, DefaultDict, Dict, List, Optional, Sequence, Tuple +from typing import ( + Any, + Callable, + DefaultDict, + Dict, + Hashable, + List, + Optional, + Sequence, + Tuple, + Union, +) from uuid import uuid1 import numpy as np @@ -21,6 +31,7 @@ from pandas.core.dtypes.common import is_float import pandas as pd +from pandas._typing import Axis, FrameOrSeries from pandas.api.types import is_dict_like, is_list_like import pandas.core.common as com from pandas.core.generic import _shared_docs @@ -125,12 +136,12 @@ class Styler: def __init__( self, data, - precision=None, + precision: Optional[int] = None, table_styles=None, - uuid=None, - caption=None, - table_attributes=None, - cell_ids=True, + uuid: Optional[str] = None, + caption: Optional[str] = None, + table_attributes: Optional[str] = None, + cell_ids: bool = True, na_rep: Optional[str] = None, ): self.ctx: DefaultDict[Tuple[int, int], List[str]] = defaultdict(list) @@ -195,21 +206,21 @@ def _repr_html_(self): def to_excel( self, excel_writer, - sheet_name="Sheet1", - na_rep="", - float_format=None, - columns=None, - header=True, - index=True, - index_label=None, - startrow=0, - startcol=0, - engine=None, - merge_cells=True, - encoding=None, - inf_rep="inf", - verbose=True, - freeze_panes=None, + sheet_name: str = "Sheet1", + na_rep: str = "", + float_format: Optional[str] = None, + columns: Optional[Sequence[Hashable]] = None, + header: Union[Sequence[Hashable], bool] = True, + index: bool = True, + index_label: Optional[Union[Hashable, Sequence[Hashable]]] = None, + startrow: int = 0, + startcol: int = 0, + engine: Optional[str] = None, + merge_cells: bool = True, + encoding: Optional[str] = None, + inf_rep: str = "inf", + verbose: bool = True, + freeze_panes: Optional[Tuple[int, int]] = None, ): from pandas.io.formats.excel import ExcelFormatter @@ -254,7 +265,7 @@ def _translate(self): BLANK_CLASS = "blank" BLANK_VALUE = "" - def format_attr(pair): + def format_attr(pair) -> str: return "{key}={value}".format(**pair) # for sparsifying a MultiIndex @@ -422,7 +433,7 @@ def format_attr(pair): table_attributes=table_attr, ) - def format(self, formatter, subset=None, na_rep: Optional[str] = None): + def format(self, formatter, subset=None, na_rep: Optional[str] = None) -> "Styler": """ Format the text display value of cells. @@ -495,7 +506,7 @@ def format(self, formatter, subset=None, na_rep: Optional[str] = None): self._display_funcs[(i, j)] = formatter return self - def render(self, **kwargs): + def render(self, **kwargs) -> str: """ Render the built up styles to HTML. @@ -509,7 +520,7 @@ def render(self, **kwargs): Returns ------- - rendered : str + str The rendered HTML. Notes @@ -544,7 +555,7 @@ def render(self, **kwargs): d.update(kwargs) return self.template.render(**d) - def _update_ctx(self, attrs): + def _update_ctx(self, attrs) -> None: """ Update the state of the Styler. @@ -562,7 +573,7 @@ def _update_ctx(self, attrs): for pair in col.rstrip(";").split(";"): self.ctx[(i, j)].append(pair) - def _copy(self, deepcopy=False): + def _copy(self, deepcopy: bool = False): styler = Styler( self.data, precision=self.precision, @@ -611,7 +622,7 @@ def _compute(self): r = func(self)(*args, **kwargs) return r - def _apply(self, func, axis=0, subset=None, **kwargs): + def _apply(self, func, axis: Optional[Axis] = 0, subset=None, **kwargs) -> "Styler": subset = slice(None) if subset is None else subset subset = _non_reducing_slice(subset) data = self.data.loc[subset] @@ -644,7 +655,7 @@ def _apply(self, func, axis=0, subset=None, **kwargs): self._update_ctx(result) return self - def apply(self, func, axis=0, subset=None, **kwargs): + def apply(self, func, axis: Optional[Axis] = 0, subset=None, **kwargs) -> "Styler": """ Apply a function column-wise, row-wise, or table-wise. @@ -695,7 +706,7 @@ def apply(self, func, axis=0, subset=None, **kwargs): ) return self - def _applymap(self, func, subset=None, **kwargs): + def _applymap(self, func, subset=None, **kwargs) -> "Styler": func = partial(func, **kwargs) # applymap doesn't take kwargs? if subset is None: subset = pd.IndexSlice[:] @@ -704,7 +715,7 @@ def _applymap(self, func, subset=None, **kwargs): self._update_ctx(result) return self - def applymap(self, func, subset=None, **kwargs): + def applymap(self, func, subset=None, **kwargs) -> "Styler": """ Apply a function elementwise. @@ -733,7 +744,9 @@ def applymap(self, func, subset=None, **kwargs): ) return self - def where(self, cond, value, other=None, subset=None, **kwargs): + def where( + self, cond, value: str, other: Optional[str] = None, subset=None, **kwargs + ) -> "Styler": """ Apply a function elementwise. @@ -764,7 +777,6 @@ def where(self, cond, value, other=None, subset=None, **kwargs): -------- Styler.applymap """ - if other is None: other = "" @@ -772,7 +784,7 @@ def where(self, cond, value, other=None, subset=None, **kwargs): lambda val: value if cond(val) else other, subset=subset, **kwargs ) - def set_precision(self, precision): + def set_precision(self, precision: int) -> "Styler": """ Set the precision used to render. @@ -787,7 +799,7 @@ def set_precision(self, precision): self.precision = precision return self - def set_table_attributes(self, attributes): + def set_table_attributes(self, attributes: str) -> "Styler": """ Set the table attributes. @@ -811,7 +823,7 @@ def set_table_attributes(self, attributes): self.table_attributes = attributes return self - def export(self): + def export(self) -> List: """ Export the styles to applied to the current Styler. @@ -827,7 +839,7 @@ def export(self): """ return self._todo - def use(self, styles): + def use(self, styles: List) -> "Styler": """ Set the styles on the current Styler. @@ -835,7 +847,7 @@ def use(self, styles): Parameters ---------- - styles : list + styles : List List of style functions. Returns @@ -849,7 +861,7 @@ def use(self, styles): self._todo.extend(styles) return self - def set_uuid(self, uuid): + def set_uuid(self, uuid: str) -> "Styler": """ Set the uuid for a Styler. @@ -864,7 +876,7 @@ def set_uuid(self, uuid): self.uuid = uuid return self - def set_caption(self, caption): + def set_caption(self, caption: str) -> "Styler": """ Set the caption on a Styler. @@ -879,7 +891,7 @@ def set_caption(self, caption): self.caption = caption return self - def set_table_styles(self, table_styles): + def set_table_styles(self, table_styles: List) -> "Styler": """ Set the table styles on a Styler. @@ -887,7 +899,7 @@ def set_table_styles(self, table_styles): Parameters ---------- - table_styles : list + table_styles : List Each individual table_style should be a dictionary with ``selector`` and ``props`` keys. ``selector`` should be a CSS selector that the style will be applied to (automatically @@ -926,7 +938,7 @@ def set_na_rep(self, na_rep: str) -> "Styler": self.na_rep = na_rep return self - def hide_index(self): + def hide_index(self) -> "Styler": """ Hide any indices from rendering. @@ -939,7 +951,7 @@ def hide_index(self): self.hidden_index = True return self - def hide_columns(self, subset): + def hide_columns(self, subset) -> "Styler": """ Hide columns from rendering. @@ -965,10 +977,10 @@ def hide_columns(self, subset): # ----------------------------------------------------------------------- @staticmethod - def _highlight_null(v, null_color): + def _highlight_null(v, null_color: str) -> str: return f"background-color: {null_color}" if pd.isna(v) else "" - def highlight_null(self, null_color="red"): + def highlight_null(self, null_color: str = "red") -> "Styler": """ Shade the background ``null_color`` for missing values. @@ -986,14 +998,14 @@ def highlight_null(self, null_color="red"): def background_gradient( self, cmap="PuBu", - low=0, - high=0, - axis=0, + low: float = 0, + high: float = 0, + axis: Optional[Axis] = 0, subset=None, - text_color_threshold=0.408, + text_color_threshold: float = 0.408, vmin: Optional[float] = None, vmax: Optional[float] = None, - ): + ) -> "Styler": """ Color the background in a gradient style. @@ -1068,9 +1080,9 @@ def background_gradient( def _background_gradient( s, cmap="PuBu", - low=0, - high=0, - text_color_threshold=0.408, + low: float = 0, + high: float = 0, + text_color_threshold: float = 0.408, vmin: Optional[float] = None, vmax: Optional[float] = None, ): @@ -1081,8 +1093,7 @@ def _background_gradient( not isinstance(text_color_threshold, (float, int)) or not 0 <= text_color_threshold <= 1 ): - msg = "`text_color_threshold` must be a value from 0 to 1." - raise ValueError(msg) + raise ValueError("`text_color_threshold` must be a value from 0 to 1.") with _mpl(Styler.background_gradient) as (plt, colors): smin = np.nanmin(s.to_numpy()) if vmin is None else vmin @@ -1094,7 +1105,7 @@ def _background_gradient( # https://github.com/matplotlib/matplotlib/issues/5427 rgbas = plt.cm.get_cmap(cmap)(norm(s.to_numpy(dtype=float))) - def relative_luminance(rgba): + def relative_luminance(rgba) -> float: """ Calculate relative luminance of a color. @@ -1116,7 +1127,7 @@ def relative_luminance(rgba): ) return 0.2126 * r + 0.7152 * g + 0.0722 * b - def css(rgba): + def css(rgba) -> str: dark = relative_luminance(rgba) < text_color_threshold text_color = "#f1f1f1" if dark else "#000000" return f"background-color: {colors.rgb2hex(rgba)};color: {text_color};" @@ -1130,7 +1141,7 @@ def css(rgba): columns=s.columns, ) - def set_properties(self, subset=None, **kwargs): + def set_properties(self, subset=None, **kwargs) -> "Styler": """ Method to set one or more non-data dependent properties or each cell. @@ -1156,7 +1167,14 @@ def set_properties(self, subset=None, **kwargs): return self.applymap(f, subset=subset) @staticmethod - def _bar(s, align, colors, width=100, vmin=None, vmax=None): + def _bar( + s, + align: str, + colors, + width=100, + vmin: Optional[float] = None, + vmax: Optional[float] = None, + ): """ Draw bar chart in dataframe cells. """ @@ -1187,7 +1205,7 @@ def css_bar(start, end, color): css += f"{color} {e:.1f}%, transparent {e:.1f}%)" return css - def css(x): + def css(x) -> str: if pd.isna(x): return "" @@ -1211,13 +1229,13 @@ def css(x): def bar( self, subset=None, - axis=0, + axis: Optional[Axis] = 0, color="#d65f5f", - width=100, - align="left", - vmin=None, - vmax=None, - ): + width: float = 100, + align: str = "left", + vmin: Optional[float] = None, + vmax: Optional[float] = None, + ) -> "Styler": """ Draw bar chart in the cell backgrounds. @@ -1272,9 +1290,9 @@ def bar( color = [color[0], color[0]] elif len(color) > 2: raise ValueError( - "`color` must be string or a list-like" - " of length 2: [`color_neg`, `color_pos`]" - " (eg: color=['#d65f5f', '#5fba7d'])" + "`color` must be string or a list-like " + "of length 2: [`color_neg`, `color_pos`] " + "(eg: color=['#d65f5f', '#5fba7d'])" ) subset = _maybe_numeric_slice(self.data, subset) @@ -1292,7 +1310,9 @@ def bar( return self - def highlight_max(self, subset=None, color="yellow", axis=0): + def highlight_max( + self, subset=None, color: str = "yellow", axis: Optional[Axis] = 0 + ) -> "Styler": """ Highlight the maximum by shading the background. @@ -1312,7 +1332,9 @@ def highlight_max(self, subset=None, color="yellow", axis=0): """ return self._highlight_handler(subset=subset, color=color, axis=axis, max_=True) - def highlight_min(self, subset=None, color="yellow", axis=0): + def highlight_min( + self, subset=None, color: str = "yellow", axis: Optional[Axis] = 0 + ) -> "Styler": """ Highlight the minimum by shading the background. @@ -1334,7 +1356,13 @@ def highlight_min(self, subset=None, color="yellow", axis=0): subset=subset, color=color, axis=axis, max_=False ) - def _highlight_handler(self, subset=None, color="yellow", axis=None, max_=True): + def _highlight_handler( + self, + subset=None, + color: str = "yellow", + axis: Optional[Axis] = None, + max_: bool = True, + ) -> "Styler": subset = _non_reducing_slice(_maybe_numeric_slice(self.data, subset)) self.apply( self._highlight_extrema, color=color, axis=axis, subset=subset, max_=max_ @@ -1342,7 +1370,9 @@ def _highlight_handler(self, subset=None, color="yellow", axis=None, max_=True): return self @staticmethod - def _highlight_extrema(data, color="yellow", max_=True): + def _highlight_extrema( + data: FrameOrSeries, color: str = "yellow", max_: bool = True + ): """ Highlight the min or max in a Series or DataFrame. """ @@ -1459,7 +1489,7 @@ def pipe(self, func, *args, **kwargs): return com.pipe(self, func, *args, **kwargs) -def _is_visible(idx_row, idx_col, lengths): +def _is_visible(idx_row, idx_col, lengths) -> bool: """ Index -> {(idx_row, idx_col): bool}). """ @@ -1516,13 +1546,13 @@ def _maybe_wrap_formatter(formatter, na_rep: Optional[str]): elif callable(formatter): formatter_func = formatter else: - msg = f"Expected a template string or callable, got {formatter} instead" - raise TypeError(msg) + raise TypeError( + f"Expected a template string or callable, got {formatter} instead" + ) if na_rep is None: return formatter_func elif isinstance(na_rep, str): return lambda x: na_rep if pd.isna(x) else formatter_func(x) else: - msg = f"Expected a string, got {na_rep} instead" - raise TypeError(msg) + raise TypeError(f"Expected a string, got {na_rep} instead")