Introduce header mode

Author: flbulgarelli

src/paper2table/readers/pdfplumber.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Literal, Optional
 
 import pandas as pd
 import pdfplumber
@@ -25,6 +25,9 @@ class TableSchema(BaseModel):
     """
 
     title: str
+
+    header_mode: Literal["all_pages", "first_page_only", "none"]
+
     first_page: int
     """
     1-based first page number where table is allocated
@@ -106,7 +109,8 @@ def read_schema_tables(pdf_path: str, schema: TablesSchema, pdf: pdfplumber.PDF)
             try:
                 dataframe = read_table(
                     table_fragment if table_fragment else [],
-                    column_mappings=table_schema.column_mappings,
+                    table_schema=table_schema,
+                    page=page,
                 )
                 tables.append(
                     DataFrameTableReader(
@@ -140,14 +144,22 @@ def to_dataframe(rows: TableFragment, column_names_hints: list[str]):
 def read_table(
     table_fragment: TableFragment,
     column_names_hints: list[str] = [],
-    column_mappings: Optional[ColumnMappings] = None,
+    table_schema: Optional[TableSchema] = None,
+    page: Optional[int] = None,
 ) -> pd.DataFrame:
     dataframe = to_dataframe(table_fragment, column_names_hints)
 
-    if column_mappings is not None:
-        dataframe = dataframe[column_mappings.keys()].rename(column_mappings)
-
-    dataframe = dataframe.rename(columns=lambda column: normalize_name(str(column)))
+    if table_schema is not None:
+        selected_column_names = list(table_schema.column_mappings.keys())
+        renamer = {(key): value for key, value in table_schema.column_mappings.items()}
+        dataframe = dataframe[selected_column_names].rename(columns=renamer)
+        if table_schema.header_mode == "all_pages" or (
+            table_schema.header_mode == "first_page_only"
+            and page == table_schema.first_page
+        ):
+            dataframe.drop([0], inplace=True)
+
+    dataframe.rename(columns=lambda column: normalize_name(str(column)), inplace=True)
     dataframe = dataframe.apply(
         lambda row: list(
             map(lambda v: v.replace("\n", " ") if type(v) == str else v, row)

tests/test_pdfplumber.py

@@ -155,11 +155,12 @@ def test_read_table_with_schema_that_matches_page():
             tables=[
                 TableSchema(
                     title="Plants",
+                    header_mode="all_pages",
                     first_page=1,
                     last_page=1,
                     column_mappings={
-                        "0": "vernacular_name",
-                        "1": "scientific_name",
+                        0: "vernacular_name",
+                        1: "scientific_name",
                     },
                 )
             ],
@@ -216,3 +217,102 @@ def test_read_table_with_schema_that_matches_page():
     result_dict = result.to_dict()
     assert result_dict["metadata"] == {"filename": "demo_table.pdf"}
     assert len(result_dict["tables"][0]["table_fragments"]) == 1
+
+
+def test_read_table_with_schema_without_headers():
+    result = read_tables(
+        "./tests/data/demo_table.pdf",
+        schema=TablesSchema(
+            tables=[
+                TableSchema(
+                    title="Plants",
+                    header_mode="none",
+                    first_page=1,
+                    last_page=1,
+                    column_mappings={
+                        0: "vernacular_name",
+                        1: "scientific_name",
+                    },
+                )
+            ],
+            citation="A citation",
+        ),
+    )
+
+    assert result.citation == "A citation"
+    assert len(result.tables) == 1
+    assert result.tables[0].title == "Plants"
+    assert result.tables[0].page == 1
+    assert result.tables[0].rows == [
+        {
+            "scientific_name": "scienti\x00c_name",
+            "vernacular_name": "common_name",
+        },
+        {
+            "vernacular_name": "Sun\x00ower",
+            "scientific_name": "Helianthus annuus",
+        },
+        {
+            "vernacular_name": "Rose",
+            "scientific_name": "Rosa gallica",
+        },
+        {
+            "vernacular_name": "Tulip",
+            "scientific_name": "Tulipa gesneriana",
+        },
+        {
+            "vernacular_name": "Lavender",
+            "scientific_name": "Lavandula angustifolia",
+        },
+        {
+            "vernacular_name": "Oak",
+            "scientific_name": "Quercus robur",
+        },
+        {
+            "vernacular_name": "Maple",
+            "scientific_name": "Acer saccharum",
+        },
+        {
+            "vernacular_name": "Dandelion",
+            "scientific_name": "Taraxacum o\x00cinale",
+        },
+        {
+            "vernacular_name": "Bamboo",
+            "scientific_name": "Bambusa vulgaris",
+        },
+        {
+            "vernacular_name": "Cactus (Prickly Pear)",
+            "scientific_name": "Opuntia \x00cus-indica",
+        },
+        {
+            "vernacular_name": "Coffee",
+            "scientific_name": "Coffea arabica",
+        },
+    ]
+    result_dict = result.to_dict()
+    assert result_dict["metadata"] == {"filename": "demo_table.pdf"}
+    assert len(result_dict["tables"][0]["table_fragments"]) == 1
+
+
+def test_read_table_with_schema_that_doesnt_matches_page():
+    result = read_tables(
+        "./tests/data/demo_table.pdf",
+        schema=TablesSchema(
+            tables=[
+                TableSchema(
+                    title="Plants",
+                    header_mode="all_pages",
+                    first_page=2,
+                    last_page=2,
+                    column_mappings={
+                        0: "vernacular_name",
+                        1: "scientific_name",
+                    },
+                )
+            ],
+            citation="A citation",
+        ),
+    )
+
+    assert result.citation == "A citation"
+    assert len(result.tables) == 0