ENH: Styler.highlight_null can accepts subset argument (#31350)

Author: immaxchen

doc/source/whatsnew/v1.1.0.rst

@@ -64,6 +64,7 @@ Other enhancements
 ^^^^^^^^^^^^^^^^^^
 
 - :class:`Styler` may now render CSS more efficiently where multiple cells have the same styling (:issue:`30876`)
+- :meth:`Styler.highlight_null` now accepts ``subset`` argument (:issue:`31345`)
 - When writing directly to a sqlite connection :func:`to_sql` now supports the ``multi`` method (:issue:`29921`)
 - `OptionError` is now exposed in `pandas.errors` (:issue:`27553`)
 - :func:`timedelta_range` will now infer a frequency when passed ``start``, ``stop``, and ``periods`` (:issue:`32377`)

pandas/io/formats/style.py

@@ -1003,19 +1003,27 @@ def hide_columns(self, subset) -> "Styler":
     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: str = "red") -> "Styler":
+    def highlight_null(
+        self,
+        null_color: str = "red",
+        subset: Optional[Union[Label, Sequence[Label]]] = None,
+    ) -> "Styler":
         """
         Shade the background ``null_color`` for missing values.
 
         Parameters
         ----------
-        null_color : str
+        null_color : str, default 'red'
+        subset : label or list of labels, default None
+            A valid slice for ``data`` to limit the style application to.
+
+            .. versionadded:: 1.1.0
 
         Returns
         -------
         self : Styler
         """
-        self.applymap(self._highlight_null, null_color=null_color)
+        self.applymap(self._highlight_null, null_color=null_color, subset=subset)
         return self
 
     def background_gradient(

pandas/tests/io/formats/test_style.py

@@ -1091,6 +1091,23 @@ def test_highlight_null(self, null_color="red"):
         expected = {(0, 0): [""], (1, 0): ["background-color: red"]}
         assert result == expected
 
+    def test_highlight_null_subset(self):
+        # GH 31345
+        df = pd.DataFrame({"A": [0, np.nan], "B": [0, np.nan]})
+        result = (
+            df.style.highlight_null(null_color="red", subset=["A"])
+            .highlight_null(null_color="green", subset=["B"])
+            ._compute()
+            .ctx
+        )
+        expected = {
+            (0, 0): [""],
+            (1, 0): ["background-color: red"],
+            (0, 1): [""],
+            (1, 1): ["background-color: green"],
+        }
+        assert result == expected
+
     def test_nonunique_raises(self):
         df = pd.DataFrame([[1, 2]], columns=["A", "A"])
         with pytest.raises(ValueError):