Support header_rows in ascii.rst (reStructuredText) writer [swev-id: astropy__astropy-14182]

astropy/io/ascii/rst.py

@@ -57,10 +57,16 @@ class RST(FixedWidth):
     data_class = SimpleRSTData
     header_class = SimpleRSTHeader
 
-    def __init__(self):
-        super().__init__(delimiter_pad=None, bookend=False)
+    def __init__(self, header_rows=None):
+        super().__init__(
+            delimiter_pad=None, bookend=False, header_rows=header_rows
+        )
 
     def write(self, lines):
         lines = super().write(lines)
-        lines = [lines[1]] + lines + [lines[1]]
+        header_rows = getattr(self.header, "header_rows", ["name"])
+        n_hdr = len(header_rows)
+        if self.header.position_line is not None:
+            separator = lines[n_hdr]
+            lines = [separator] + lines + [separator]
         return lines

astropy/io/ascii/tests/test_rst.py

@@ -2,7 +2,9 @@
 
 from io import StringIO
 
+import astropy.units as u
 from astropy.io import ascii
+from astropy.table import MaskedColumn, QTable
 
 from .common import assert_almost_equal, assert_equal
 
@@ -185,3 +187,108 @@ def test_write_normal():
 ==== ========= ==== ====
 """,
     )
+
+
+def test_rst_header_rows_name_unit():
+    tbl = QTable(
+        {
+            "wave": [350, 950] * u.nm,
+            "response": [0.7, 1.2] * u.count,
+        }
+    )
+    out = StringIO()
+    ascii.write(tbl, out, format="rst", header_rows=["name", "unit"])
+    expected = (
+        "===== ========\n"
+        " wave response\n"
+        "   nm       ct\n"
+        "===== ========\n"
+        "350.0      0.7\n"
+        "950.0      1.2\n"
+        "===== ========\n"
+    )
+    assert_equal_splitlines(out.getvalue(), expected)
+
+
+def test_rst_header_rows_missing_units():
+    tbl = QTable(
+        {
+            "wave": [350, 950] * u.nm,
+            "response": [0.7, 1.2],
+        }
+    )
+    out = StringIO()
+    ascii.write(tbl, out, format="rst", header_rows=["name", "unit"])
+    expected = (
+        "===== ========\n"
+        " wave response\n"
+        "   nm         \n"
+        "===== ========\n"
+        "350.0      0.7\n"
+        "950.0      1.2\n"
+        "===== ========\n"
+    )
+    assert_equal_splitlines(out.getvalue(), expected)
+
+
+def test_rst_header_rows_unicode():
+    tbl = QTable(
+        {
+            "wave": [350, 950] * u.AA,
+            "flux": [2.3, 3.1] * u.Unit("erg / (s cm2)"),
+        }
+    )
+    out = StringIO()
+    ascii.write(tbl, out, format="rst", header_rows=["name", "unit"])
+    expected = (
+        "======== =============\n"
+        "    wave          flux\n"
+        "Angstrom erg / (cm2 s)\n"
+        "======== =============\n"
+        "   350.0           2.3\n"
+        "   950.0           3.1\n"
+        "======== =============\n"
+    )
+    assert_equal_splitlines(out.getvalue(), expected)
+
+
+def test_rst_header_rows_masked():
+    response = MaskedColumn([0.7, 1.2], mask=[False, True], unit=u.count)
+    tbl = QTable(
+        {
+            "wave": [350, 950] * u.nm,
+            "response": response,
+        }
+    )
+    out = StringIO()
+    ascii.write(tbl, out, format="rst", header_rows=["name", "unit"])
+    expected = (
+        "===== ========\n"
+        " wave response\n"
+        "   nm       ct\n"
+        "===== ========\n"
+        "350.0      0.7\n"
+        "950.0         \n"
+        "===== ========\n"
+    )
+    assert_equal_splitlines(out.getvalue(), expected)
+
+
+def test_rst_backward_compat_no_header_rows():
+    tbl = QTable(
+        {
+            "wave": [350, 950],
+            "response": [0.7, 1.2],
+        }
+    )
+    out = StringIO()
+    ascii.write(tbl, out, format="rst")
+    expected = (
+        "==== ========\n"
+        "wave response\n"
+        "==== ========\n"
+        " 350      0.7\n"
+        " 950      1.2\n"
+        "==== ========\n"
+    )
+    assert_equal_splitlines(out.getvalue(), expected)

docs/io/ascii/fixed_width_gallery.rst

@@ -286,6 +286,28 @@ Fixed Width Two Line
 ..
   EXAMPLE END
 
+..
+  EXAMPLE START
+  Custom Header Rows with RST
+
+::
+    >>> from astropy.table import QTable
+    >>> import astropy.units as u
+    >>> from astropy.io import ascii
+    >>> tbl = QTable({"wave": [350, 950] * u.nm,
+    ...               "response": [0.7, 1.2] * u.count})
+    >>> ascii.write(tbl, format="rst", header_rows=["name", "unit"])
+    ===== ========
+     wave response
+       nm       ct
+    ===== ========
+    350.0      0.7
+    950.0      1.2
+    ===== ========
+
+..
+  EXAMPLE END
+
 ..
   EXAMPLE START
   Reading a reStructuredText Table