fix tests

pandas-dev · attack68 · May 24, 2021 · May 26, 2021 · May 29, 2021 · Jun 1, 2021
commit c9669424134451eca306cc602ceed121b6fcc1a9
diff --git a/pandas/core/generic.py b/pandas/core/generic.py
@@ -3122,6 +3122,7 @@ def to_latex(
         header=True,
         index=True,
         columns=None,
+        index_names=True,
         column_format=None,
         position=None,
         position_float=None,
@@ -3143,7 +3144,6 @@ def to_latex(
         multicolumn=None,
         multicolumn_format=None,
         multirow=None,
-        index_names=True,
         col_space=None,
         formatters=None,
         float_format=None,
@@ -3180,12 +3180,17 @@ def to_latex(
             ``pandas.options.styler.sparse.columns`` values.
 
             .. versionadded:: 1.4.0
-        index, header : bool
-            Whether to print index labels and/or column headers respectively.
+        index : bool
+            Whether to print index labels.
 
             .. versionchanged:: 1.4.0
+        header : bool or list of str
+            Whether to print column headers. If a list of strings is given is assumed
+            to be aliases for the column names.
         columns : list of label, optional
             The subset of columns to write. Writes all columns by default.
+        index_names : bool, default True
+            Remove the names of indexes before rendering.
         column_format : str, optional
             The LaTeX column specification placed in location:
 
@@ -3292,10 +3297,6 @@ def to_latex(
             .. deprecated:: 1.4.0
                Deprecated in favour of using arguments native to ``Styler.format``,
                such as ``precision``, ``decimal``, and ``thousands``.
-        index_names : bool, default True
-
-            .. deprecated:: 1.4.0
-               Remove the names of indexes before rendering.
         multicolumn : bool, default True
 
             .. deprecated:: 1.4.0
@@ -3366,10 +3367,29 @@ def to_latex(
             escape="latex" if escape else None,
         )
 
-        if not header:
+        if header is False:
             styler.hide_columns()
+        elif isinstance(header, (list, tuple)):
+            if len(header) != len(styler.columns):
+                raise ValueError(
+                    f"`header` gave {len(header)} aliases for {len(styler.columns)} "
+                    f"columns."
+                )
+            for i, val in enumerate(header):
+
+                def disp(x, v):
+                    return f"{v}"
+
+                styler._display_funcs_columns[(0, i)] = functools.partial(disp, v=val)
+            if isinstance(styler.columns, MultiIndex):
+                styler.hide_columns(
+                    level=[i for i in range(styler.columns.nlevels) if i != 0]
+                )
         if not index:
             styler.hide_index()
+        if not index_names:
+            styler.hide_index(names=True)
+            styler.hide_columns(names=True)
         if columns:
             hidden = [col for col in styler.columns if col not in columns]
             styler.hide_columns(hidden)

diff --git a/pandas/io/formats/style.py b/pandas/io/formats/style.py
@@ -780,6 +780,8 @@ def to_latex(
         )
 
         if column_format is not None:
+            if not isinstance(column_format, str):
+                raise ValueError("`column_format` must be a string.")
             # add more recent setting to table_styles
             obj.set_table_styles(
                 [{"selector": "column_format", "props": f":{column_format}"}],

diff --git a/pandas/tests/io/formats/test_to_latex.py b/pandas/tests/io/formats/test_to_latex.py
@@ -99,7 +99,7 @@ def test_to_latex_tabular_without_index(self):
     )
     def test_to_latex_bad_column_format(self, bad_column_format):
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        msg = r"column_format must be str or unicode"
+        msg = r"`column_format` must be a string"
         with pytest.raises(ValueError, match=msg):
             df.to_latex(column_format=bad_column_format)
 
@@ -197,72 +197,69 @@ def test_to_latex_midrule_location(self):
 class TestToLatexLongtable:
     def test_to_latex_empty_longtable(self):
         df = DataFrame()
-        result = df.to_latex(longtable=True)
+        result = df.to_latex(environment="longtable")
         expected = _dedent(
             r"""
             \begin{longtable}{l}
-            \toprule
-            Empty DataFrame
-            Columns: Index([], dtype='object')
-            Index: Index([], dtype='object') \\
+            \endfirsthead
+            \endhead
+            \multicolumn{1}{r}{Continued on next page} \\
+            \endfoot
+            \endlastfoot
             \end{longtable}
             """
         )
         assert result == expected
 
     def test_to_latex_longtable_with_index(self):
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        result = df.to_latex(longtable=True)
+        result = df.to_latex(environment="longtable", hrules=True)
         expected = _dedent(
             r"""
             \begin{longtable}{lrl}
             \toprule
-            {} &  a &   b \\
+            {} & {a} & {b} \\
             \midrule
             \endfirsthead
-
             \toprule
-            {} &  a &   b \\
+            {} & {a} & {b} \\
             \midrule
             \endhead
             \midrule
-            \multicolumn{3}{r}{{Continued on next page}} \\
+            \multicolumn{3}{r}{Continued on next page} \\
             \midrule
             \endfoot
-
             \bottomrule
             \endlastfoot
-            0 &  1 &  b1 \\
-            1 &  2 &  b2 \\
+            0 & 1 & b1 \\
+            1 & 2 & b2 \\
             \end{longtable}
             """
         )
         assert result == expected
 
     def test_to_latex_longtable_without_index(self):
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        result = df.to_latex(index=False, longtable=True)
+        result = df.to_latex(index=False, environment="longtable", hrules=True)
         expected = _dedent(
             r"""
             \begin{longtable}{rl}
             \toprule
-             a &  b \\
+            {a} & {b} \\
             \midrule
             \endfirsthead
-
             \toprule
-             a &  b \\
+            {a} & {b} \\
             \midrule
             \endhead
             \midrule
-            \multicolumn{2}{r}{{Continued on next page}} \\
+            \multicolumn{2}{r}{Continued on next page} \\
             \midrule
             \endfoot
-
             \bottomrule
             \endlastfoot
-             1 & b1 \\
-             2 & b2 \\
+            1 & b1 \\
+            2 & b2 \\
             \end{longtable}
             """
         )
@@ -277,21 +274,22 @@ def test_to_latex_longtable_without_index(self):
         ],
     )
     def test_to_latex_longtable_continued_on_next_page(self, df, expected_number):
-        result = df.to_latex(index=False, longtable=True)
+        result = df.to_latex(index=False, environment="longtable")
         assert fr"\multicolumn{{{expected_number}}}" in result
 
 
 class TestToLatexHeader:
     def test_to_latex_no_header_with_index(self):
         # GH 7124
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        result = df.to_latex(header=False)
+        result = df.to_latex(header=False, hrules=True)
         expected = _dedent(
             r"""
             \begin{tabular}{lrl}
             \toprule
-            0 &  1 &  b1 \\
-            1 &  2 &  b2 \\
+            \midrule
+            0 & 1 & b1 \\
+            1 & 2 & b2 \\
             \bottomrule
             \end{tabular}
             """
@@ -301,11 +299,12 @@ def test_to_latex_no_header_with_index(self):
     def test_to_latex_no_header_without_index(self):
         # GH 7124
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        result = df.to_latex(index=False, header=False)
+        result = df.to_latex(index=False, header=False, hrules=True)
         expected = _dedent(
             r"""
             \begin{tabular}{rl}
             \toprule
+            \midrule
             1 & b1 \\
             2 & b2 \\
             \bottomrule
@@ -317,15 +316,15 @@ def test_to_latex_no_header_without_index(self):
     def test_to_latex_specified_header_with_index(self):
         # GH 7124
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        result = df.to_latex(header=["AA", "BB"])
+        result = df.to_latex(header=["AA", "BB"], hrules=True)
         expected = _dedent(
             r"""
             \begin{tabular}{lrl}
             \toprule
-            {} & AA &  BB \\
+            {} & {AA} & {BB} \\
             \midrule
-            0 &  1 &  b1 \\
-            1 &  2 &  b2 \\
+            0 & 1 & b1 \\
+            1 & 2 & b2 \\
             \bottomrule
             \end{tabular}
             """
@@ -335,15 +334,15 @@ def test_to_latex_specified_header_with_index(self):
     def test_to_latex_specified_header_without_index(self):
         # GH 7124
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        result = df.to_latex(header=["AA", "BB"], index=False)
+        result = df.to_latex(header=["AA", "BB"], index=False, hrules=True)
         expected = _dedent(
             r"""
             \begin{tabular}{rl}
             \toprule
-            AA & BB \\
+            {AA} & {BB} \\
             \midrule
-             1 & b1 \\
-             2 & b2 \\
+            1 & b1 \\
+            2 & b2 \\
             \bottomrule
             \end{tabular}
             """
@@ -366,7 +365,7 @@ def test_to_latex_number_of_items_in_header_missmatch_raises(
     ):
         # GH 7124
         df = DataFrame({"a": [1, 2], "b": ["b1", "b2"]})
-        msg = f"Writing 2 cols but got {num_aliases} aliases"
+        msg = f"`header` gave {num_aliases} aliases for 2 columns."
         with pytest.raises(ValueError, match=msg):
             df.to_latex(header=header)