Backport PR pandas-dev#45138: ENH: add cline to Styler.to_latex

attack68 · meeseeksmachine · commit 7bb2d171c493 · 2022-01-05T23:30:43.000Z
diff --git a/doc/source/whatsnew/v1.4.0.rst b/doc/source/whatsnew/v1.4.0.rst
@@ -105,6 +105,7 @@ There are also some LaTeX specific enhancements:
 
   - :meth:`.Styler.to_latex` introduces keyword argument ``environment``, which also allows a specific "longtable" entry through a separate jinja2 template (:issue:`41866`).
   - Naive sparsification is now possible for LaTeX without the necessity of including the multirow package (:issue:`43369`)
+  - *cline* support has been added for MultiIndex row sparsification through a keyword argument (:issue:`45138`)
 
 .. _whatsnew_140.enhancements.pyarrow_csv_engine:
 
diff --git a/pandas/io/formats/style.py b/pandas/io/formats/style.py
@@ -494,6 +494,7 @@ def to_latex(
         position: str | None = None,
         position_float: str | None = None,
         hrules: bool | None = None,
+        clines: str | None = None,
         label: str | None = None,
         caption: str | tuple | None = None,
         sparse_index: bool | None = None,
@@ -542,6 +543,22 @@ def to_latex(
             Defaults to ``pandas.options.styler.latex.hrules``, which is `False`.
 
             .. versionchanged:: 1.4.0
+        clines : str, optional
+            Use to control adding \\cline commands for the index labels separation.
+            Possible values are:
+
+              - `None`: no cline commands are added (default).
+              - `"all;data"`: a cline is added for every index value extending the
+                width of the table, including data entries.
+              - `"all;index"`: as above with lines extending only the width of the
+                index entries.
+              - `"skip-last;data"`: a cline is added for each index value except the
+                last level (which is never sparsified), extending the widtn of the
+                table.
+              - `"skip-last;index"`: as above with lines extending only the width of the
+                index entries.
+
+            .. versionadded:: 1.4.0
         label : str, optional
             The LaTeX label included as: \\label{<label>}.
             This is used with \\ref{<label>} in the main .tex file.
@@ -911,6 +928,7 @@ def to_latex(
             environment=environment,
             convert_css=convert_css,
             siunitx=siunitx,
+            clines=clines,
         )
 
         encoding = encoding or get_option("styler.render.encoding")
diff --git a/pandas/io/formats/style_render.py b/pandas/io/formats/style_render.py
@@ -164,14 +164,16 @@ def _render_html(
             html_style_tpl=self.template_html_style,
         )
 
-    def _render_latex(self, sparse_index: bool, sparse_columns: bool, **kwargs) -> str:
+    def _render_latex(
+        self, sparse_index: bool, sparse_columns: bool, clines: str | None, **kwargs
+    ) -> str:
         """
         Render a Styler in latex format
         """
         self._compute()
 
         d = self._translate(sparse_index, sparse_columns, blank="")
-        self._translate_latex(d)
+        self._translate_latex(d, clines=clines)
 
         self.template_latex.globals["parse_wrap"] = _parse_latex_table_wrapping
         self.template_latex.globals["parse_table"] = _parse_latex_table_styles
@@ -257,13 +259,19 @@ def _translate(
         head = self._translate_header(sparse_cols, max_cols)
         d.update({"head": head})
 
+        # for sparsifying a MultiIndex and for use with latex clines
+        idx_lengths = _get_level_lengths(
+            self.index, sparse_index, max_rows, self.hidden_rows
+        )
+        d.update({"index_lengths": idx_lengths})
+
         self.cellstyle_map: DefaultDict[tuple[CSSPair, ...], list[str]] = defaultdict(
             list
         )
         self.cellstyle_map_index: DefaultDict[
             tuple[CSSPair, ...], list[str]
         ] = defaultdict(list)
-        body = self._translate_body(sparse_index, max_rows, max_cols)
+        body = self._translate_body(idx_lengths, max_rows, max_cols)
         d.update({"body": body})
 
         ctx_maps = {
@@ -515,7 +523,7 @@ def _generate_index_names_row(self, iter: tuple, max_cols: int, col_lengths: dic
 
         return index_names + column_blanks
 
-    def _translate_body(self, sparsify_index: bool, max_rows: int, max_cols: int):
+    def _translate_body(self, idx_lengths: dict, max_rows: int, max_cols: int):
         """
         Build each <tr> within table <body> as a list
 
@@ -537,11 +545,6 @@ def _translate_body(self, sparsify_index: bool, max_rows: int, max_cols: int):
         body : list
             The associated HTML elements needed for template rendering.
         """
-        # for sparsifying a MultiIndex
-        idx_lengths = _get_level_lengths(
-            self.index, sparsify_index, max_rows, self.hidden_rows
-        )
-
         rlabels = self.data.index.tolist()
         if not isinstance(self.data.index, MultiIndex):
             rlabels = [[x] for x in rlabels]
@@ -738,7 +741,7 @@ def _generate_body_row(
 
         return index_headers + data
 
-    def _translate_latex(self, d: dict) -> None:
+    def _translate_latex(self, d: dict, clines: str | None) -> None:
         r"""
         Post-process the default render dict for the LaTeX template format.
 
@@ -749,10 +752,10 @@ def _translate_latex(self, d: dict) -> None:
             or multirow sparsification (so that \multirow and \multicol work correctly).
         """
         index_levels = self.index.nlevels
-        visible_index_levels = index_levels - sum(self.hide_index_)
+        visible_index_level_n = index_levels - sum(self.hide_index_)
         d["head"] = [
             [
-                {**col, "cellstyle": self.ctx_columns[r, c - visible_index_levels]}
+                {**col, "cellstyle": self.ctx_columns[r, c - visible_index_level_n]}
                 for c, col in enumerate(row)
                 if col["is_visible"]
             ]
@@ -790,6 +793,39 @@ def _translate_latex(self, d: dict) -> None:
             body.append(row_body_headers + row_body_cells)
         d["body"] = body
 
+        # clines are determined from info on index_lengths and hidden_rows and input
+        # to a dict defining which row clines should be added in the template.
+        if clines not in [
+            None,
+            "all;data",
+            "all;index",
+            "skip-last;data",
+            "skip-last;index",
+        ]:
+            raise ValueError(
+                f"`clines` value of {clines} is invalid. Should either be None or one "
+                f"of 'all;data', 'all;index', 'skip-last;data', 'skip-last;index'."
+            )
+        elif clines is not None:
+            data_len = len(row_body_cells) if "data" in clines else 0
+
+            d["clines"] = defaultdict(list)
+            visible_row_indexes: list[int] = [
+                r for r in range(len(self.data.index)) if r not in self.hidden_rows
+            ]
+            visible_index_levels: list[int] = [
+                i for i in range(index_levels) if not self.hide_index_[i]
+            ]
+            for rn, r in enumerate(visible_row_indexes):
+                for lvln, lvl in enumerate(visible_index_levels):
+                    if lvl == index_levels - 1 and "skip-last" in clines:
+                        continue
+                    idx_len = d["index_lengths"].get((lvl, r), None)
+                    if idx_len is not None:  # i.e. not a sparsified entry
+                        d["clines"][rn + idx_len].append(
+                            f"\\cline{{{lvln+1}-{len(visible_index_levels)+data_len}}}"
+                        )
+
     def format(
         self,
         formatter: ExtFormatter | None = None,
diff --git a/pandas/io/formats/templates/latex_longtable.tpl b/pandas/io/formats/templates/latex_longtable.tpl
@@ -73,6 +73,10 @@
 {% for c in row %}{% if not loop.first %} & {% endif %}
   {%- if c.type == 'th' %}{{parse_header(c, multirow_align, multicol_align)}}{% else %}{{parse_cell(c.cellstyle, c.display_value, convert_css)}}{% endif %}
 {%- endfor %} \\
+{% if clines and clines[loop.index] | length > 0 %}
+  {%- for cline in clines[loop.index] %}{% if not loop.first %} {% endif %}{{ cline }}{% endfor %}
+
+{% endif %}
 {% endfor %}
 \end{longtable}
 {% raw %}{% endraw %}
diff --git a/pandas/io/formats/templates/latex_table.tpl b/pandas/io/formats/templates/latex_table.tpl
@@ -41,6 +41,10 @@
 {% for c in row %}{% if not loop.first %} & {% endif %}
   {%- if c.type == 'th' %}{{parse_header(c, multirow_align, multicol_align, False, convert_css)}}{% else %}{{parse_cell(c.cellstyle, c.display_value, convert_css)}}{% endif %}
 {%- endfor %} \\
+{% if clines and clines[loop.index] | length > 0 %}
+  {%- for cline in clines[loop.index] %}{% if not loop.first %} {% endif %}{{ cline }}{% endfor %}
+
+{% endif %}
 {% endfor %}
 {% set bottomrule = parse_table(table_styles, 'bottomrule') %}
 {% if bottomrule is not none %}
diff --git a/pandas/tests/io/formats/style/test_to_latex.py b/pandas/tests/io/formats/style/test_to_latex.py
@@ -876,3 +876,117 @@ def test_apply_index_hidden_levels():
         """
     )
     assert result == expected
+
+
+@pytest.mark.parametrize("clines", ["bad", "index", "skip-last", "all", "data"])
+def test_clines_validation(clines, styler):
+    msg = f"`clines` value of {clines} is invalid."
+    with pytest.raises(ValueError, match=msg):
+        styler.to_latex(clines=clines)
+
+
+@pytest.mark.parametrize(
+    "clines, exp",
+    [
+        ("all;index", "\n\\cline{1-1}"),
+        ("all;data", "\n\\cline{1-2}"),
+        ("skip-last;index", ""),
+        ("skip-last;data", ""),
+        (None, ""),
+    ],
+)
+@pytest.mark.parametrize("env", ["table", "longtable"])
+def test_clines_index(clines, exp, env):
+    df = DataFrame([[1], [2], [3], [4]])
+    result = df.style.to_latex(clines=clines, environment=env)
+    expected = f"""\
+0 & 1 \\\\{exp}
+1 & 2 \\\\{exp}
+2 & 3 \\\\{exp}
+3 & 4 \\\\{exp}
+"""
+    assert expected in result
+
+
+@pytest.mark.parametrize(
+    "clines, expected",
+    [
+        (
+            None,
+            dedent(
+                """\
+            \\multirow[c]{2}{*}{A} & X & 1 \\\\
+             & Y & 2 \\\\
+            \\multirow[c]{2}{*}{B} & X & 3 \\\\
+             & Y & 4 \\\\
+            """
+            ),
+        ),
+        (
+            "skip-last;index",
+            dedent(
+                """\
+            \\multirow[c]{2}{*}{A} & X & 1 \\\\
+             & Y & 2 \\\\
+            \\cline{1-2}
+            \\multirow[c]{2}{*}{B} & X & 3 \\\\
+             & Y & 4 \\\\
+            \\cline{1-2}
+            """
+            ),
+        ),
+        (
+            "skip-last;data",
+            dedent(
+                """\
+            \\multirow[c]{2}{*}{A} & X & 1 \\\\
+             & Y & 2 \\\\
+            \\cline{1-3}
+            \\multirow[c]{2}{*}{B} & X & 3 \\\\
+             & Y & 4 \\\\
+            \\cline{1-3}
+            """
+            ),
+        ),
+        (
+            "all;index",
+            dedent(
+                """\
+            \\multirow[c]{2}{*}{A} & X & 1 \\\\
+            \\cline{2-2}
+             & Y & 2 \\\\
+            \\cline{1-2} \\cline{2-2}
+            \\multirow[c]{2}{*}{B} & X & 3 \\\\
+            \\cline{2-2}
+             & Y & 4 \\\\
+            \\cline{1-2} \\cline{2-2}
+            """
+            ),
+        ),
+        (
+            "all;data",
+            dedent(
+                """\
+            \\multirow[c]{2}{*}{A} & X & 1 \\\\
+            \\cline{2-3}
+             & Y & 2 \\\\
+            \\cline{1-3} \\cline{2-3}
+            \\multirow[c]{2}{*}{B} & X & 3 \\\\
+            \\cline{2-3}
+             & Y & 4 \\\\
+            \\cline{1-3} \\cline{2-3}
+            """
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize("env", ["table"])
+def test_clines_multiindex(clines, expected, env):
+    # also tests simultaneously with hidden rows and a hidden multiindex level
+    midx = MultiIndex.from_product([["A", "-", "B"], [0], ["X", "Y"]])
+    df = DataFrame([[1], [2], [99], [99], [3], [4]], index=midx)
+    styler = df.style
+    styler.hide([("-", 0, "X"), ("-", 0, "Y")])
+    styler.hide(level=1)
+    result = styler.to_latex(clines=clines, environment=env)
+    assert expected in result
