ENH: added optional caption and label arguments to DataFrame.to_latex()

pandas/core/generic.py

@@ -2781,17 +2781,20 @@ class      (index) object 'bird' 'bird' 'mammal' 'mammal'
     def to_latex(self, buf=None, columns=None, col_space=None, header=True,
                  index=True, na_rep='NaN', formatters=None, float_format=None,
                  sparsify=None, index_names=True, bold_rows=False,
-                 column_format=None, longtable=None, escape=None,
-                 encoding=None, decimal='.', multicolumn=None,
-                 multicolumn_format=None, multirow=None):
+                 column_format=None, longtable=None, caption=None,
+                 label=None, escape=None, encoding=None, decimal='.',
+                 multicolumn=None, multicolumn_format=None, multirow=None):
         r"""
-        Render an object to a LaTeX tabular environment table.
+        Render an object to a LaTeX tabular, longtable, or nested
+        table/tabular environment.
 
         Render an object to a tabular environment table. You can splice
         this into a LaTeX document. Requires \usepackage{booktabs}.
 
         .. versionchanged:: 0.20.2
-           Added to Series
+           Added to Series.
+        .. versionchanged:: 0.25.0
+           Added caption and label arguments.
 
         Parameters
         ----------
@@ -2831,6 +2834,13 @@ def to_latex(self, buf=None, columns=None, col_space=None, header=True,
             By default, the value will be read from the pandas config
             module. Use a longtable environment instead of tabular. Requires
             adding a \usepackage{longtable} to your LaTeX preamble.
+        caption : str, optional
+            The LaTeX caption to be placed inside \caption{} in the output.
+            .. versionadded:: 0.25.0
+        label : str, optional
+            The LaTeX label to be placed inside \label{} in the output.
+            This is used with \ref{} in the main .tex file.
+            .. versionadded:: 0.25.0
         escape : bool, optional
             By default, the value will be read from the pandas config
             module. When set to False prevents from escaping latex special
@@ -2908,7 +2918,8 @@ def to_latex(self, buf=None, columns=None, col_space=None, header=True,
                                        index_names=index_names,
                                        escape=escape, decimal=decimal)
         formatter.to_latex(column_format=column_format, longtable=longtable,
-                           encoding=encoding, multicolumn=multicolumn,
+                           caption=caption, label=label, encoding=encoding,
+                           multicolumn=multicolumn,
                            multicolumn_format=multicolumn_format,
                            multirow=multirow)
 

pandas/io/formats/format.py

@@ -676,15 +676,17 @@ def _join_multiline(self, *strcols):
             st = ed
         return '\n\n'.join(str_lst)
 
-    def to_latex(self, column_format=None, longtable=False, encoding=None,
-                 multicolumn=False, multicolumn_format=None, multirow=False):
+    def to_latex(self, column_format=None, longtable=False, caption=None,
+                 label=None, encoding=None, multicolumn=False,
+                 multicolumn_format=None, multirow=False):
         """
         Render a DataFrame to a LaTeX tabular/longtable environment output.
         """
 
         from pandas.io.formats.latex import LatexFormatter
         latex_renderer = LatexFormatter(self, column_format=column_format,
-                                        longtable=longtable,
+                                        longtable=longtable, caption=caption,
+                                        label=label,
                                         multicolumn=multicolumn,
                                         multicolumn_format=multicolumn_format,
                                         multirow=multirow)

pandas/io/formats/latex.py

@@ -34,19 +34,23 @@ class LatexFormatter(TableFormatter):
     """
 
     def __init__(self, formatter, column_format=None, longtable=False,
-                 multicolumn=False, multicolumn_format=None, multirow=False):
+                 caption=None, label=None, multicolumn=False,
+                 multicolumn_format=None, multirow=False):
         self.fmt = formatter
         self.frame = self.fmt.frame
         self.bold_rows = self.fmt.kwds.get('bold_rows', False)
         self.column_format = column_format
         self.longtable = longtable
+        self.caption = caption
+        self.label = label
         self.multicolumn = multicolumn
         self.multicolumn_format = multicolumn_format
         self.multirow = multirow
 
     def write_result(self, buf):
         """
-        Render a DataFrame to a LaTeX tabular/longtable environment output.
+        Render a DataFrame to a LaTeX tabular, longtable, or table/tabular
+        environment output.
         """
 
         # string representation of the columns
@@ -106,14 +110,12 @@ def pad_empties(x):
             raise AssertionError('column_format must be str or unicode, '
                                  'not {typ}'.format(typ=type(column_format)))
 
-        if not self.longtable:
-            buf.write('\\begin{{tabular}}{{{fmt}}}\n'
-                      .format(fmt=column_format))
-            buf.write('\\toprule\n')
+        if self.longtable:
+            self._write_longtable_begin(buf, column_format)
         else:
-            buf.write('\\begin{{longtable}}{{{fmt}}}\n'
-                      .format(fmt=column_format))
-            buf.write('\\toprule\n')
+            self._write_tabular_begin(buf, column_format)
+
+        buf.write('\\toprule\n')
 
         ilevels = self.frame.index.nlevels
         clevels = self.frame.columns.nlevels
@@ -164,11 +166,10 @@ def pad_empties(x):
             if self.multirow and i < len(strrows) - 1:
                 self._print_cline(buf, i, len(strcols))
 
-        if not self.longtable:
-            buf.write('\\bottomrule\n')
-            buf.write('\\end{tabular}\n')
+        if self.longtable:
+            self._write_longtable_end(buf)
         else:
-            buf.write('\\end{longtable}\n')
+            self._write_tabular_end(buf)
 
     def _format_multicolumn(self, row, ilevels):
         r"""
@@ -244,3 +245,74 @@ def _print_cline(self, buf, i, icol):
                           .format(cl=cl[1], icol=icol))
         # remove entries that have been written to buffer
         self.clinebuf = [x for x in self.clinebuf if x[0] != i]
+
+    def _write_tabular_begin(self, buf, column_format):
+        """
+        write the beginning of a tabular environment or
+        nested table/tabular environments including caption and label
+        """
+        if self.caption is None and self.label is None:
+            # then write output only in a tabular environment
+            pass
+        else:
+            # then write output in a nested table/tabular environment
+            if self.caption is None:
+                caption_ = ''
+            else:
+                caption_ = '\n\\caption{{{}}}'.format(self.caption)
+
+            if self.label is None:
+                label_ = ''
+            else:
+                label_ = '\n\\label{{{}}}'.format(self.label)
+
+            buf.write('\\begin{{table}}\n\\centering{}{}\n'.format(
+                caption_,
+                label_
+            ))
+
+        buf.write('\\begin{{tabular}}{{{fmt}}}\n'.format(fmt=column_format))
+
+    def _write_longtable_begin(self, buf, column_format):
+        """
+        write the beginning of a longtable environment including caption and
+        label if provided by user
+        """
+        buf.write('\\begin{{longtable}}{{{fmt}}}\n'.format(fmt=column_format))
+
+        if self.caption is None and self.label is None:
+            pass
+        else:
+            if self.caption is None:
+                pass
+            else:
+                buf.write('\\caption{{{}}}'.format(self.caption))
+
+            if self.label is None:
+                pass
+            else:
+                buf.write('\\label{{{}}}'.format(self.label))
+
+            # a double-backslash is required at the end of the line
+            # as discussed here:
+            # https://tex.stackexchange.com/questions/219138
+            buf.write('\\\\\n')
+
+    def _write_tabular_end(self, buf):
+        """
+        write the end of a tabular environment or nested table/tabular
+        environment
+        """
+        buf.write('\\bottomrule\n')
+        buf.write('\\end{tabular}\n')
+        if self.caption is None and self.label is None:
+            pass
+        else:
+            buf.write('\\end{table}\n')
+
+    @staticmethod
+    def _write_longtable_end(buf):
+        """
+        write the end of a longtable environment
+        """
+        buf.write('\\end{longtable}\n')

pandas/tests/io/formats/test_to_latex.py

@@ -391,9 +391,7 @@ def test_to_latex_special_escape(self):
 """
         assert escaped_result == escaped_expected
 
-    def test_to_latex_longtable(self, frame):
-        frame.to_latex(longtable=True)
-
+    def test_to_latex_longtable(self):
         df = DataFrame({'a': [1, 2], 'b': ['b1', 'b2']})
         withindex_result = df.to_latex(longtable=True)
         withindex_expected = r"""\begin{longtable}{lrl}
@@ -442,6 +440,158 @@ def test_to_latex_longtable(self, frame):
         with3columns_result = df.to_latex(index=False, longtable=True)
         assert r"\multicolumn{3}" in with3columns_result
 
+    def test_to_latex_caption_label(self):
+        # GH 25436
+        the_caption = 'a table in a \\texttt{table/tabular} environment'
+        the_label = 'tab:table_tabular'
+
+        df = DataFrame({'a': [1, 2], 'b': ['b1', 'b2']})
+
+        # test when only the caption is provided
+        result_c = df.to_latex(
+            caption=the_caption
+        )
+
+        expected_c = r"""\begin{table}
+\centering
+\caption{a table in a \texttt{table/tabular} environment}
+\begin{tabular}{lrl}
+\toprule
+{} &  a &   b \\
+\midrule
+0 &  1 &  b1 \\
+1 &  2 &  b2 \\
+\bottomrule
+\end{tabular}
+\end{table}
+"""
+        assert result_c == expected_c
+
+        # test when only the label is provided
+        result_l = df.to_latex(
+            label=the_label
+        )
+
+        expected_l = r"""\begin{table}
+\centering
+\label{tab:table_tabular}
+\begin{tabular}{lrl}
+\toprule
+{} &  a &   b \\
+\midrule
+0 &  1 &  b1 \\
+1 &  2 &  b2 \\
+\bottomrule
+\end{tabular}
+\end{table}
+"""
+        assert result_l == expected_l
+
+        # test when the caption and the label are provided
+        result_cl = df.to_latex(
+            caption=the_caption,
+            label=the_label
+        )
+
+        expected_cl = r"""\begin{table}
+\centering
+\caption{a table in a \texttt{table/tabular} environment}
+\label{tab:table_tabular}
+\begin{tabular}{lrl}
+\toprule
+{} &  a &   b \\
+\midrule
+0 &  1 &  b1 \\
+1 &  2 &  b2 \\
+\bottomrule
+\end{tabular}
+\end{table}
+"""
+        assert result_cl == expected_cl
+
+    def test_to_latex_longtable_caption_label(self):
+        # GH 25436
+        the_caption = 'a table in a \\texttt{longtable} environment'
+        the_label = 'tab:longtable'
+
+        df = DataFrame({'a': [1, 2], 'b': ['b1', 'b2']})
+
+        # test when only the caption is provided
+        result_c = df.to_latex(
+            longtable=True,
+            caption=the_caption
+        )
+
+        expected_c = r"""\begin{longtable}{lrl}
+\caption{a table in a \texttt{longtable} environment}\\
+\toprule
+{} &  a &   b \\
+\midrule
+\endhead
+\midrule
+\multicolumn{3}{r}{{Continued on next page}} \\
+\midrule
+\endfoot
+
+\bottomrule
+\endlastfoot
+0 &  1 &  b1 \\
+1 &  2 &  b2 \\
+\end{longtable}
+"""
+        assert result_c == expected_c
+
+        # test when only the label is provided
+        result_l = df.to_latex(
+            longtable=True,
+            label=the_label,
+        )
+
+        expected_l = r"""\begin{longtable}{lrl}
+\label{tab:longtable}\\
+\toprule
+{} &  a &   b \\
+\midrule
+\endhead
+\midrule
+\multicolumn{3}{r}{{Continued on next page}} \\
+\midrule
+\endfoot
+
+\bottomrule
+\endlastfoot
+0 &  1 &  b1 \\
+1 &  2 &  b2 \\
+\end{longtable}
+"""
+        assert result_l == expected_l
+
+        # test when the caption and the label are provided
+        result_cl = df.to_latex(
+            longtable=True,
+            caption=the_caption,
+            label=the_label,
+        )
+
+        expected_cl = r"""\begin{longtable}{lrl}
+\caption{a table in a \texttt{longtable} environment}\label{tab:longtable}\\
+\toprule
+{} &  a &   b \\
+\midrule
+\endhead
+\midrule
+\multicolumn{3}{r}{{Continued on next page}} \\
+\midrule
+\endfoot
+
+\bottomrule
+\endlastfoot
+0 &  1 &  b1 \\
+1 &  2 &  b2 \\
+\end{longtable}
+"""
+        assert result_cl == expected_cl
+
     def test_to_latex_escape_special_chars(self):
         special_characters = ['&', '%', '$', '#', '_', '{', '}', '~', '^',
                               '\\']