Merge pull request #5 from HaishuoFang/fix/table-extraction

Fix table extraction from arXiv HTML papers

Author: timf34

src/arxiv2md/markdown.py

@@ -297,15 +297,33 @@ def _serialize_table(table: Tag, *, remove_inline_citations: bool = False) -> st
         return f"$$ {eqn_text} $$"
 
     rows = []
-    for row in table.find_all("tr", recursive=False):
-        cells = row.find_all(["th", "td"], recursive=False)
-        if not cells:
-            continue
-        values = []
-        for cell in cells:
-            cell_text = _cleanup_inline_text(_serialize_inline(cell, remove_inline_citations=remove_inline_citations)).replace("\n", "<br>")
-            values.append(cell_text)
-        rows.append(values)
+    # Find rows in tbody, thead, tfoot, or directly in table
+    # Handle nested structure where rows might be inside tbody/thead/tfoot
+    tbody_elements = table.find_all(["tbody", "thead", "tfoot"], recursive=False)
+    
+    if tbody_elements:
+        # Table has tbody/thead/tfoot structure - find rows within them
+        for tbody in tbody_elements:
+            for row in tbody.find_all("tr", recursive=False):
+                cells = row.find_all(["th", "td"], recursive=False)
+                if not cells:
+                    continue
+                values = []
+                for cell in cells:
+                    cell_text = _cleanup_inline_text(_serialize_inline(cell, remove_inline_citations=remove_inline_citations)).replace("\n", "<br>")
+                    values.append(cell_text)
+                rows.append(values)
+    else:
+        # Table has no tbody/thead/tfoot - find rows directly in table
+        for row in table.find_all("tr", recursive=False):
+            cells = row.find_all(["th", "td"], recursive=False)
+            if not cells:
+                continue
+            values = []
+            for cell in cells:
+                cell_text = _cleanup_inline_text(_serialize_inline(cell, remove_inline_citations=remove_inline_citations)).replace("\n", "<br>")
+                values.append(cell_text)
+            rows.append(values)
 
     if not rows:
         return ""
@@ -323,18 +341,40 @@ def _serialize_table(table: Tag, *, remove_inline_citations: bool = False) -> st
 
 
 def _serialize_figure(figure: Tag, *, remove_inline_citations: bool = False) -> str:
+    # Check if this is a table figure (ltx_table class)
+    figure_classes = " ".join(figure.get("class", []))
+    is_table_figure = "ltx_table" in figure_classes
+
     caption_tag = figure.find("figcaption")
     caption = _normalize_text(_serialize_inline(caption_tag, remove_inline_citations=remove_inline_citations)) if caption_tag else ""
-    img = figure.find("img")
-    src = img.get("src") if img else None
-    alt = img.get("alt") if img else None
 
     lines = []
-    if caption:
-        lines.append(f"Figure: {caption}")
-    if src:
-        image_label = alt or "Image"
-        lines.append(f"{image_label}: {src}")
+
+    if is_table_figure:
+        # Handle table figures - find and serialize the embedded table
+        # Note: fix_tabular_tables strips attributes, so search for any table element
+        table = figure.find("table")
+        if table:
+            table_md = _serialize_table(table, remove_inline_citations=remove_inline_citations)
+            if caption:
+                lines.append(f"**{caption}**")
+            if table_md:
+                lines.append(table_md)
+        elif caption:
+            # Fallback if no table found but has caption
+            lines.append(f"Table: {caption}")
+    else:
+        # Handle regular image figures
+        img = figure.find("img")
+        src = img.get("src") if img else None
+        alt = img.get("alt") if img else None
+
+        if caption:
+            lines.append(f"Figure: {caption}")
+        if src:
+            image_label = alt or "Image"
+            lines.append(f"{image_label}: {src}")
+
     return "\n".join(lines).strip()
 
 

tests/test_markdown.py

@@ -26,3 +26,90 @@ def test_math_and_tables_render() -> None:
     assert "| 1 | 2 |" in markdown
     assert "$$" in markdown
     assert "E = mc^2" in markdown
+
+
+def test_table_with_tbody() -> None:
+    """Test that tables with tbody/thead/tfoot structure are correctly converted."""
+    html = """
+    <table class="ltx_tabular">
+        <tbody>
+            <tr><th>Model</th><th>Accuracy</th></tr>
+            <tr><td>Llama-7B</td><td>70.12</td></tr>
+            <tr><td>Llama-13B</td><td>72.39</td></tr>
+        </tbody>
+    </table>
+    """
+
+    markdown = convert_fragment_to_markdown(html)
+
+    # Should contain table structure
+    assert "| Model | Accuracy |" in markdown
+    assert "| --- | --- |" in markdown
+    assert "| Llama-7B | 70.12 |" in markdown
+    assert "| Llama-13B | 72.39 |" in markdown
+
+
+def test_table_with_thead_tbody() -> None:
+    """Test that tables with thead and tbody are correctly converted."""
+    html = """
+    <table class="ltx_tabular">
+        <thead>
+            <tr><th>Method</th><th>Result</th></tr>
+        </thead>
+        <tbody>
+            <tr><td>Prune SW</td><td>0.0%</td></tr>
+            <tr><td>Prune Non-SW</td><td>68.5%</td></tr>
+        </tbody>
+    </table>
+    """
+
+    markdown = convert_fragment_to_markdown(html)
+
+    # Should contain table structure
+    assert "| Method | Result |" in markdown
+    assert "| Prune SW | 0.0% |" in markdown
+    assert "| Prune Non-SW | 68.5% |" in markdown
+
+
+def test_table_inside_figure() -> None:
+    """Test that tables wrapped in figure elements (ltx_table) are correctly converted."""
+    html = """
+    <figure class="ltx_table" id="S3.T1">
+        <table class="ltx_tabular ltx_centering ltx_guessed_headers ltx_align_middle">
+            <thead class="ltx_thead">
+                <tr class="ltx_tr">
+                    <th class="ltx_td ltx_align_left ltx_th ltx_th_column">Model</th>
+                    <th class="ltx_td ltx_align_center ltx_th ltx_th_column">Arc-c</th>
+                    <th class="ltx_td ltx_align_center ltx_th ltx_th_column">Arc-e</th>
+                </tr>
+            </thead>
+            <tbody class="ltx_tbody">
+                <tr class="ltx_tr">
+                    <td class="ltx_td ltx_align_left">Original</td>
+                    <td class="ltx_td ltx_align_center">41.81</td>
+                    <td class="ltx_td ltx_align_center">75.29</td>
+                </tr>
+                <tr class="ltx_tr">
+                    <td class="ltx_td ltx_align_left">Prune SW</td>
+                    <td class="ltx_td ltx_align_center">19.80</td>
+                    <td class="ltx_td ltx_align_center">39.60</td>
+                </tr>
+            </tbody>
+        </table>
+        <figcaption class="ltx_caption ltx_centering">
+            <span class="ltx_tag ltx_tag_table">Table 1: </span>
+            <span class="ltx_text ltx_font_bold">Super Weight Importance</span>.
+            Pruning the super weight significantly impairs quality.
+        </figcaption>
+    </figure>
+    """
+
+    markdown = convert_fragment_to_markdown(html)
+
+    # Should contain the caption
+    assert "Table 1:" in markdown
+    assert "Super Weight Importance" in markdown
+    # Should contain the actual table data
+    assert "| Model | Arc-c | Arc-e |" in markdown
+    assert "| Original | 41.81 | 75.29 |" in markdown
+    assert "| Prune SW | 19.80 | 39.60 |" in markdown