type subset (pandas-dev#41433)

Co-authored-by: JHM Darbyshire (iMac) <[email protected]>

Author: attack68

pandas/io/formats/style.py

@@ -43,6 +43,7 @@
     CSSProperties,
     CSSStyles,
     StylerRenderer,
+    Subset,
     Tooltips,
     maybe_convert_css_to_tuples,
     non_reducing_slice,
@@ -545,7 +546,7 @@ def _apply(
         self,
         func: Callable[..., Styler],
         axis: Axis | None = 0,
-        subset=None,
+        subset: Subset | None = None,
         **kwargs,
     ) -> Styler:
         subset = slice(None) if subset is None else subset
@@ -590,7 +591,7 @@ def apply(
         self,
         func: Callable[..., Styler],
         axis: Axis | None = 0,
-        subset=None,
+        subset: Subset | None = None,
         **kwargs,
     ) -> Styler:
         """
@@ -651,7 +652,9 @@ def apply(
         )
         return self
 
-    def _applymap(self, func: Callable, subset=None, **kwargs) -> Styler:
+    def _applymap(
+        self, func: Callable, subset: Subset | None = None, **kwargs
+    ) -> Styler:
         func = partial(func, **kwargs)  # applymap doesn't take kwargs?
         if subset is None:
             subset = pd.IndexSlice[:]
@@ -660,7 +663,9 @@ def _applymap(self, func: Callable, subset=None, **kwargs) -> Styler:
         self._update_ctx(result)
         return self
 
-    def applymap(self, func: Callable, subset=None, **kwargs) -> Styler:
+    def applymap(
+        self, func: Callable, subset: Subset | None = None, **kwargs
+    ) -> Styler:
         """
         Apply a CSS-styling function elementwise.
 
@@ -707,7 +712,7 @@ def where(
         cond: Callable,
         value: str,
         other: str | None = None,
-        subset=None,
+        subset: Subset | None = None,
         **kwargs,
     ) -> Styler:
         """
@@ -1061,7 +1066,7 @@ def hide_index(self) -> Styler:
         self.hidden_index = True
         return self
 
-    def hide_columns(self, subset) -> Styler:
+    def hide_columns(self, subset: Subset) -> Styler:
         """
         Hide columns from rendering.
 
@@ -1093,7 +1098,7 @@ def background_gradient(
         low: float = 0,
         high: float = 0,
         axis: Axis | None = 0,
-        subset=None,
+        subset: Subset | None = None,
         text_color_threshold: float = 0.408,
         vmin: float | None = None,
         vmax: float | None = None,
@@ -1239,7 +1244,7 @@ def background_gradient(
         )
         return self
 
-    def set_properties(self, subset=None, **kwargs) -> Styler:
+    def set_properties(self, subset: Subset | None = None, **kwargs) -> Styler:
         """
         Set defined CSS-properties to each ``<td>`` HTML element within the given
         subset.
@@ -1331,7 +1336,7 @@ def css(x):
 
     def bar(
         self,
-        subset=None,
+        subset: Subset | None = None,
         axis: Axis | None = 0,
         color="#d65f5f",
         width: float = 100,
@@ -1417,7 +1422,7 @@ def bar(
     def highlight_null(
         self,
         null_color: str = "red",
-        subset: IndexLabel | None = None,
+        subset: Subset | None = None,
         props: str | None = None,
     ) -> Styler:
         """
@@ -1462,7 +1467,7 @@ def f(data: DataFrame, props: str) -> np.ndarray:
 
     def highlight_max(
         self,
-        subset: IndexLabel | None = None,
+        subset: Subset | None = None,
         color: str = "yellow",
         axis: Axis | None = 0,
         props: str | None = None,
@@ -1511,7 +1516,7 @@ def f(data: FrameOrSeries, props: str) -> np.ndarray:
 
     def highlight_min(
         self,
-        subset: IndexLabel | None = None,
+        subset: Subset | None = None,
         color: str = "yellow",
         axis: Axis | None = 0,
         props: str | None = None,
@@ -1560,7 +1565,7 @@ def f(data: FrameOrSeries, props: str) -> np.ndarray:
 
     def highlight_between(
         self,
-        subset: IndexLabel | None = None,
+        subset: Subset | None = None,
         color: str = "yellow",
         axis: Axis | None = 0,
         left: Scalar | Sequence | None = None,
@@ -1667,7 +1672,7 @@ def highlight_between(
 
     def highlight_quantile(
         self,
-        subset: IndexLabel | None = None,
+        subset: Subset | None = None,
         color: str = "yellow",
         axis: Axis | None = 0,
         q_left: float = 0.0,

pandas/io/formats/style_render.py

@@ -55,6 +55,7 @@ class CSSDict(TypedDict):
 
 
 CSSStyles = List[CSSDict]
+Subset = Union[slice, Sequence, Index]
 
 
 class StylerRenderer:
@@ -402,7 +403,7 @@ def _translate_body(
     def format(
         self,
         formatter: ExtFormatter | None = None,
-        subset: slice | Sequence[Any] | None = None,
+        subset: Subset | None = None,
         na_rep: str | None = None,
         precision: int | None = None,
         decimal: str = ".",
@@ -772,7 +773,7 @@ def _maybe_wrap_formatter(
         return lambda x: na_rep if isna(x) else func_2(x)
 
 
-def non_reducing_slice(slice_):
+def non_reducing_slice(slice_: Subset):
     """
     Ensure that a slice doesn't reduce to a Series or Scalar.
 
@@ -809,7 +810,9 @@ def pred(part) -> bool:
             # slice(a, b, c)
             slice_ = [slice_]  # to tuplize later
     else:
-        slice_ = [part if pred(part) else [part] for part in slice_]
+        # error: Item "slice" of "Union[slice, Sequence[Any]]" has no attribute
+        # "__iter__" (not iterable) -> is specifically list_like in conditional
+        slice_ = [p if pred(p) else [p] for p in slice_]  # type: ignore[union-attr]
     return tuple(slice_)