fix(tables): keep th row headers and carry rowspan down

crawl4ai/table_extraction.py

@@ -210,6 +210,20 @@ def is_data_table(self, table: etree.Element, **kwargs) -> bool:
         threshold = kwargs.get("table_score_threshold", self.table_score_threshold)
         return score >= threshold
 
+    @staticmethod
+    def _span(cell: etree.Element, name: str) -> int:
+        """`colspan`/`rowspan` as a positive count.
+
+        The value comes from the document, so it can be absent, empty, `0`, or
+        not a number at all. A table must not be turned into an error by any of
+        those, and a browser treats them all as "no span", so they mean 1 here.
+        """
+        try:
+            span = int(cell.get(name, 1))
+        except (TypeError, ValueError):
+            return 1
+        return span if span > 0 else 1
+
     def extract_table_data(self, table: etree.Element) -> Dict[str, Any]:
         """
         Extract structured data from a table element.
@@ -232,30 +246,64 @@ def extract_table_data(self, table: etree.Element) -> Dict[str, Any]:
 
         # Extract headers with colspan handling
         headers = []
+        header_row = None
         thead_rows = table.xpath(".//thead/tr")
         if thead_rows:
             header_cells = thead_rows[0].xpath(".//th")
             for cell in header_cells:
                 text = cell.text_content().strip()
-                colspan = int(cell.get("colspan", 1))
+                colspan = self._span(cell, "colspan")
                 headers.extend([text] * colspan)
         else:
             # Check first row for headers
             first_row = table.xpath(".//tr[1]")
             if first_row:
-                for cell in first_row[0].xpath(".//th|.//td"):
+                header_row = first_row[0]
+                for cell in header_row.xpath(".//th|.//td"):
                     text = cell.text_content().strip()
-                    colspan = int(cell.get("colspan", 1))
+                    colspan = self._span(cell, "colspan")
                     headers.extend([text] * colspan)
 
-        # Extract rows with colspan handling
+        # Extract rows with colspan and rowspan handling.
+        #
+        # `th` counts as a body cell: a `<th scope="row">` is the key column of
+        # a documentation table, and reading only `td` dropped it and shifted
+        # every other cell one place left.
+        #
+        # A `rowspan` cell occupies its column in the rows below it as well, so
+        # its value is carried down; without that the rows under it shift left
+        # by one and no longer line up with the header.
         rows = []
+        carried: Dict[int, List[Any]] = {}
         for row in table.xpath(".//tr[not(ancestor::thead)]"):
+            # The first row is the header when there is no thead, and it is in
+            # this set too. Emitting it again would repeat it as data.
+            if header_row is not None and row is header_row:
+                continue
             row_data = []
-            for cell in row.xpath(".//td"):
+            column = 0
+            cells = row.xpath(".//th|.//td")
+            index = 0
+            while index < len(cells) or column in carried:
+                if column in carried:
+                    text, remaining = carried[column]
+                    row_data.append(text)
+                    if remaining > 1:
+                        carried[column] = [text, remaining - 1]
+                    else:
+                        del carried[column]
+                    column += 1
+                    continue
+                cell = cells[index]
+                index += 1
                 text = cell.text_content().strip()
-                colspan = int(cell.get("colspan", 1))
-                row_data.extend([text] * colspan)
+                colspan = self._span(cell, "colspan")
+                rowspan = self._span(cell, "rowspan")
+                for _ in range(colspan):
+                    row_data.append(text)
+                    if rowspan > 1:
+                        carried[column] = [text, rowspan - 1]
+                    column += 1
             if row_data:
                 rows.append(row_data)
 

tests/test_table_extraction_alignment.py

@@ -0,0 +1,159 @@
+"""Row alignment in DefaultTableExtraction.
+
+`result.tables[i]["rows"]` has to match the grid a browser renders. Two shapes
+used to come out misaligned: a `<th scope="row">` in the body was dropped, and
+a `rowspan` cell did not occupy the rows below it.
+
+Fixes: https://github.com/unclecode/crawl4ai/issues/2258
+"""
+
+import os
+import sys
+
+import pytest
+from lxml import html as lhtml
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from crawl4ai.table_extraction import DefaultTableExtraction
+
+
+def _extract(markup: str) -> dict:
+    tree = lhtml.fromstring(markup)
+    table = tree if tree.tag == "table" else tree.xpath("//table")[0]
+    return DefaultTableExtraction().extract_table_data(table)
+
+
+class TestRowAlignment:
+    def test_a_th_row_header_stays_in_the_first_column(self):
+        result = _extract(
+            """<table>
+            <tr><th></th><th>Feature A</th><th>Feature B</th></tr>
+            <tr><th scope="row">Item 1</th><td>yes</td><td>yes</td></tr>
+            <tr><th scope="row">Item 2</th><td>no</td><td>yes</td></tr>
+            </table>"""
+        )
+
+        assert result["headers"] == ["", "Feature A", "Feature B"]
+        assert result["rows"] == [
+            ["Item 1", "yes", "yes"],
+            ["Item 2", "no", "yes"],
+        ]
+
+    def test_the_header_row_is_not_repeated_as_data(self):
+        # Without a thead the first row is the header, and it is also in the
+        # set of body rows. Now that `th` counts as a body cell, it has to be
+        # skipped explicitly rather than by having no `td` in it.
+        result = _extract(
+            """<table>
+            <tr><th>A</th><th>B</th></tr>
+            <tr><td>1</td><td>2</td></tr>
+            </table>"""
+        )
+
+        assert result["headers"] == ["A", "B"]
+        assert result["rows"] == [["1", "2"]]
+
+    def test_a_rowspan_cell_fills_the_rows_it_covers(self):
+        result = _extract(
+            """<table>
+            <thead><tr><th>Group</th><th>Option X</th><th>Option Y</th></tr></thead>
+            <tbody>
+            <tr><td rowspan="2">G1</td><td>x1</td><td>y1</td></tr>
+            <tr><td>x2</td><td>y2</td></tr>
+            </tbody></table>"""
+        )
+
+        assert result["rows"] == [
+            ["G1", "x1", "y1"],
+            ["G1", "x2", "y2"],
+        ]
+
+    def test_a_rowspan_runs_out_after_the_rows_it_declared(self):
+        result = _extract(
+            """<table>
+            <thead><tr><th>Group</th><th>Value</th></tr></thead>
+            <tbody>
+            <tr><td rowspan="2">G1</td><td>a</td></tr>
+            <tr><td>b</td></tr>
+            <tr><td>G2</td><td>c</td></tr>
+            </tbody></table>"""
+        )
+
+        assert result["rows"] == [["G1", "a"], ["G1", "b"], ["G2", "c"]]
+
+    def test_a_rowspan_in_a_later_column_lands_in_that_column(self):
+        result = _extract(
+            """<table>
+            <thead><tr><th>A</th><th>B</th><th>C</th></tr></thead>
+            <tbody>
+            <tr><td>a1</td><td rowspan="2">shared</td><td>c1</td></tr>
+            <tr><td>a2</td><td>c2</td></tr>
+            </tbody></table>"""
+        )
+
+        assert result["rows"] == [
+            ["a1", "shared", "c1"],
+            ["a2", "shared", "c2"],
+        ]
+
+    def test_colspan_still_repeats_across_the_columns_it_covers(self):
+        result = _extract(
+            """<table>
+            <thead><tr><th>A</th><th>B</th><th>C</th></tr></thead>
+            <tbody><tr><td colspan="2">wide</td><td>c</td></tr></tbody></table>"""
+        )
+
+        assert result["rows"] == [["wide", "wide", "c"]]
+
+    def test_a_cell_with_both_spans_covers_the_whole_block(self):
+        result = _extract(
+            """<table>
+            <thead><tr><th>A</th><th>B</th><th>C</th></tr></thead>
+            <tbody>
+            <tr><td rowspan="2" colspan="2">block</td><td>c1</td></tr>
+            <tr><td>c2</td></tr>
+            </tbody></table>"""
+        )
+
+        assert result["rows"] == [
+            ["block", "block", "c1"],
+            ["block", "block", "c2"],
+        ]
+
+    @pytest.mark.parametrize("rowspan", ["0", "-1", "", "nonsense"])
+    def test_a_rowspan_that_is_not_a_positive_number_is_ignored(self, rowspan):
+        # lxml hands back whatever the document said. A table must not be
+        # turned into an error, and no row may be carried into.
+        markup = f"""<table>
+        <thead><tr><th>A</th><th>B</th></tr></thead>
+        <tbody>
+        <tr><td rowspan="{rowspan}">x</td><td>a</td></tr>
+        <tr><td>y</td><td>b</td></tr>
+        </tbody></table>"""
+
+        assert _extract(markup)["rows"] == [["x", "a"], ["y", "b"]]
+
+    @pytest.mark.parametrize("colspan", ["0", "-1", "", "nonsense"])
+    def test_a_colspan_that_is_not_a_positive_number_still_yields_one_cell(
+        self, colspan
+    ):
+        # Zero columns would drop the cell entirely and shift the rest left,
+        # which is the same misalignment this change exists to fix.
+        markup = f"""<table>
+        <thead><tr><th>A</th><th>B</th></tr></thead>
+        <tbody><tr><td colspan="{colspan}">x</td><td>a</td></tr></tbody></table>"""
+
+        assert _extract(markup)["rows"] == [["x", "a"]]
+
+    def test_a_plain_table_is_unchanged(self):
+        result = _extract(
+            """<table>
+            <thead><tr><th>A</th><th>B</th></tr></thead>
+            <tbody><tr><td>1</td><td>2</td></tr><tr><td>3</td><td>4</td></tr></tbody></table>"""
+        )
+
+        assert result["headers"] == ["A", "B"]
+        assert result["rows"] == [["1", "2"], ["3", "4"]]
+        assert result["metadata"]["column_count"] == 2
+        assert result["metadata"]["row_count"] == 2