Update src/extractors.py

src/extractors.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
+
 import pandas as pd
 
-from schema import ExtractionResult, DealSnapshot, ClassifiedTable
+from schema import ExtractionResult, DealSnapshot, ClassifiedTable, ReadResult
 from normalize import clean_text, find_first
 from table_classifier import classify_table
 
+
 def extract_snapshot_from_text(text: str) -> DealSnapshot:
     t = clean_text(text)
 
-    # VERY IMPORTANT: 先用 deterministic regex 抽取你最常用字段
     property_name = find_first([
         r"Property Name[:\s]+(.+?)(?:\s{2,}|Address|Location|$)",
         r"Offering Memorandum[:\s]+(.+?)(?:\s{2,}|$)",
@@ -74,28 +75,44 @@ def extract_snapshot_from_text(text: str) -> DealSnapshot:
 
     return snapshot
 
-def extract_tables(read) -> list[ClassifiedTable]:
-    out = []
-    for tb in getattr(read, "tables", []) or []:
-        page = int(tb.get("page", -1))
-        rows = tb.get("rows", []) or []
+
+def extract_tables(read: ReadResult) -> list[ClassifiedTable]:
+    out: list[ClassifiedTable] = []
+
+    for tb in (read.tables or []):
+        # ✅ tb is TableExtract (pydantic model), NOT dict
+        page = int(tb.page) if tb.page is not None else -1
+        rows = tb.rows or []
         if not rows:
             continue
 
-        headers = rows[0]
+        headers = rows[0] if len(rows) >= 1 else []
         data_rows = rows[1:] if len(rows) > 1 else []
+
         table_type = classify_table(headers=headers, rows=data_rows)
 
-        # make df
-        df = pd.DataFrame(data_rows, columns=[h or f"col_{i}" for i, h in enumerate(headers)])
-        out.append(ClassifiedTable(page=page, table_type=table_type, df=df.to_dict(orient="records")))
+        # DataFrame -> json-serializable rows (list[dict])
+        if headers:
+            cols = [(h.strip() if isinstance(h, str) and h.strip() else f"col_{i}") for i, h in enumerate(headers)]
+            df = pd.DataFrame(data_rows, columns=cols)
+        else:
+            df = pd.DataFrame(data_rows)
+
+        out.append(
+            ClassifiedTable(
+                page=page,
+                table_type=table_type,
+                rows=df.to_dict(orient="records"),  # ✅ schema field is rows
+            )
+        )
+
     return out
 
-def extract_deal_snapshot(filename: str, read) -> ExtractionResult:
+
+def extract_deal_snapshot(filename: str, read: ReadResult) -> ExtractionResult:
     snapshot = extract_snapshot_from_text(read.text or "")
     tables = extract_tables(read)
 
-    # NOTE: streamlit 展示时需要 dataframe，所以 dict->df 在 UI 里做
     return ExtractionResult(
         filename=filename,
         snapshot=snapshot,