quantile add

Author: attack68

pandas/io/formats/style.py

@@ -1769,6 +1769,68 @@ def highlight_range(
             stop=stop,
         )
 
+    def highlight_quantile(
+        self,
+        subset: Optional[IndexLabel] = None,
+        color: str = "yellow",
+        q_low: float = 0.0,
+        q_high: float = 1.0,
+        axis: Optional[Axis] = 0,
+        props: Optional[str] = None,
+    ) -> Styler:
+        """
+        Highlight values defined by inclusion in given quantile by shading the
+        background, or otherwise.
+
+        Parameters
+        ----------
+        subset : IndexSlice, default None
+            A valid slice for ``data`` to limit the style application to.
+        color : str, default 'yellow'
+            Background color added to CSS style.
+        q_low : float, default 0
+            Left bound for the target quantile range (exclusive if not 0).
+        q_high : float, default 1
+            Right bound for the target quantile range (inclusive)
+        axis : {0 or 'index', 1 or 'columns', None}, default 0
+            Apply to each column (``axis=0`` or ``'index'``), to each row
+            (``axis=1`` or ``'columns'``), or to the entire DataFrame at once
+            with ``axis=None``.
+        props : str, default None
+            CSS properties to use for highlighting. If ``props`` is given, ``color``
+            is not used.
+
+        Returns
+        -------
+        self : Styler
+
+        See Also
+        --------
+        Styler.highlight_max:
+        Styler.highlight_min:
+        Styler.highlight_range:
+
+        Notes
+        -----
+        This function only works with consistent ``dtypes`` within the ``subset``.
+        For example a mixture of datetime-like and float items will raise errors.
+
+        This method uses ``pandas.qcut`` to implement the quantile labelling of data
+        values.
+        """
+        if props is None:
+            props = f"background-color: {color};"
+        return self.apply(
+            _Builtins._highlight_func,
+            axis=axis,
+            subset=subset,
+            props=props,
+            highlight="quantile",
+            q_low=q_low,
+            q_high=q_high,
+            axis_=axis,
+        )
+
     @classmethod
     def from_custom_template(cls, searchpath, name):
         """
@@ -1893,8 +1955,32 @@ def _highlight_func(
             in_l = data >= l if l is not None else np.full_like(data, True, dtype=bool)
             in_h = data <= h if h is not None else np.full_like(data, True, dtype=bool)
             return np.where(in_l & in_h, props, "")
+        elif highlight == "quantile":
+            l, h = kwargs.get("q_low"), kwargs.get("q_high")
+            if l > 0:
+                q, tgt_label = [0, l, h], 1
+            else:
+                q, tgt_label = [0, h], 0
+            # if l > 0 and h < 1:
+            #     q, tgt_label = [0, l, h, 1], 1
+            # elif l > 0 and h == 1:
+            #     q, tgt_label = [0, l, h], 1
+            # elif l == 0 and h < 1:
+            #     q, tgt_label = [0, h, 1], 0
+            # elif l == 0 and h == 1:
+            #     q, tgt_label = [0, 1], 0
+            # else:
+            #     raise ValueError("'q_low' and 'q_high' must be in range [0,1]")
+            if kwargs.get("axis_") is None:
+                shape = data.values.shape
+                labels = pd.qcut(data.values.ravel(), q=q, labels=False).reshape(shape)
+            else:
+                labels = pd.qcut(data, q=q, labels=False)
+            return np.where(labels == tgt_label, props, "")
         else:
-            raise ValueError("'highlight' must be one of 'max', 'min', 'null', 'range'")
+            raise ValueError(
+                "'highlight' must be one of 'max', 'min', 'null', 'range', 'quantile'."
+            )
 
 
 class _Tooltips: