Update extract_red_text.py

extract_red_text.py

@@ -1,335 +1,311 @@
 #!/usr/bin/env python3
-import re, json, sys
+import re
+import json
+import sys
 from docx import Document
 from docx.oxml.ns import qn
+from master_key import TABLE_SCHEMAS, HEADING_PATTERNS, PARAGRAPH_PATTERNS
 
-def is_red_font(run) -> bool:
-    """Return True if this run is coloured red-ish."""
+def is_red_font(run):
+    """Enhanced red font detection with better color checking"""
     col = run.font.color
     if col and col.rgb:
-        r,g,b = col.rgb[0], col.rgb[1], col.rgb[2]
-        if r>150 and g<100 and b<100 and (r-g)>30 and (r-b)>30:
+        r, g, b = col.rgb
+        if r > 150 and g < 100 and b < 100 and (r-g) > 30 and (r-b) > 30:
             return True
-    # fallback: raw <w:color w:val="XXXXXX"/>
     rPr = getattr(run._element, "rPr", None)
     if rPr is not None:
         clr = rPr.find(qn('w:color'))
         if clr is not None:
             val = clr.get(qn('w:val'))
-            if re.fullmatch(r"[0-9A-Fa-f]{6}", val):
-                rr,gg,bb = int(val[:2],16), int(val[2:4],16), int(val[4:],16)
-                if rr>150 and gg<100 and bb<100 and (rr-gg)>30 and (rr-bb)>30:
+            if val and re.fullmatch(r"[0-9A-Fa-f]{6}", val):
+                rr, gg, bb = int(val[:2], 16), int(val[2:4], 16), int(val[4:], 16)
+                if rr > 150 and gg < 100 and bb < 100 and (rr-gg) > 30 and (rr-bb) > 30:
                     return True
     return False
 
-# ─────────────────────────────────────────────────────────────────────────────
-# Your template, mapped 1:1 to doc.tables[0..18]
-MASTER_TABLES = [
-  # Table 0: Tick as appropriate (Mass, Maintenance, etc.)
-  {
-    "name": "Tick as appropriate",
-    "labels_on_row1": True,
-    "labels": [
-      "Mass", "Maintenance", "Basic Fatigue", "Advanced Fatigue",
-      "Entry Audit", "Initial Compliance Audit", "Compliance Audit",
-      "Spot Check", "Triggered Audit"
-    ]
-  },
-
-  # Table 1: Audit Information
-  {
-    "name": "Audit Information",
-    "labels_on_left": True,
-    "labels": [
-      "Date of Audit",
-      "Location of audit",
-      "Auditor name",
-      "Audit Matrix Identifier (Name or Number)",  # Corrected full label
-      "Auditor Exemplar Global Reg No.",
-      "expiry Date:",
-      "NHVR Auditor Registration Number",
-      "expiry Date:"  # Note: Duplicate label, might need special handling
-    ]
-  },
-
-  # Table 2: Operator Information (including contact details)
-  {
-    "name": "Operator Information",
-    "labels_on_left": True,
-    "skip_rows": ["Operator contact details", ""],  # Skip subheading and blank rows
-    "labels": [
-      "Operator name (Legal entity)",
-      "NHVAS Accreditation No. (If applicable)",
-      "Registered trading name/s",
-      "Australian Company Number",
-      "NHVAS Manual (Policies and Procedures) developed by",
-      "Operator business address",
-      "Operator Postal address",
-      "Email address",
-      "Operator Telephone Number"
-    ]
-  },
-
-  # Table 3: Attendance List
-  {
-    "name": "Attendance List (Names and Position Titles)",
-    "labels": ["Attendance List (Names and Position Titles)"]
-  },
-
-  # Table 4: Nature of the Operators Business
-  {
-    "name": "Nature of the Operators Business (Summary)",
-    "labels": [
-      "Nature of the Operators Business (Summary)",
-      "Accreditation Number:",
-      "Expiry Date:"
-    ]
-  },
-
-  # Table 5: Accreditation Vehicle Summary
-  {
-    "name": "Accreditation Vehicle Summary",
-    "labels_on_left": True,
-    "labels": [
-      "Number of powered vehicles",
-      "Number of trailing vehicles"
-    ]
-  },
-
-  # Table 6: Accreditation Driver Summary
-  {
-    "name": "Accreditation Driver Summary",
-    "labels_on_left": True,
-    "labels": [
-      "Number of drivers in BFM",
-      "Number of drivers in AFM"
-    ]
-  },
-
-  # Table 7: Compliance Codes
-  {
-    "name": "Compliance Codes",
-    "labels_on_row1": True,
-    "labels": ["V", "SFI", "NA", "NC", "NAP"]
-  },
-
-  # Table 8: Corrective Action Request Identification
-  {
-    "name": "Corrective Action Request Identification",
-    "labels_on_row1": True,
-    "labels": ["Title", "Abbreviation", "Description"]
-  },
-
-  # Table 9: MASS MANAGEMENT (Standards 1-8)
-  {
-    "name": "MASS MANAGEMENT",
-    "labels_on_left": True,
-    "labels": [
-      "Std 1. Responsibilities",
-      "Std 2. Vehicle Control",
-      "Std 3. Vehicle Use",
-      "Std 4. Records and Documentation",
-      "Std 5. Verification",
-      "Std 6. Internal Review",
-      "Std 7. Training and Education",
-      "Std 8. Maintenance of Suspension"
-    ]
-  },
-
-  # Table 10: Mass Management Summary of Audit findings (Standards 1-8)
-  {
-    "name": "Mass Management Summary of Audit findings",
-    "labels_on_left": True,
-    "labels": [
-      "Std 1. Responsibilities",
-      "Std 2. Vehicle Control",
-      "Std 3. Vehicle Use",
-      "Std 4. Records and Documentation",
-      "Std 5. Verification",
-      "Std 6. Internal Review",
-      "Std 7. Training and Education",
-      "Std 8. Maintenance of Suspension"
-    ]
-  },
-
-  # Table 11: Vehicle Registration Numbers of Records Examined
-  {
-    "name": "Vehicle Registration Numbers of Records Examined",
-    "labels_on_row1": True,
-    "labels": [
-      "No.", "Registration Number",
-      "Sub-contractor (Yes/No)",
-      "Sub-contracted Vehicles Statement of Compliance (Yes/No)",
-      "Weight Verification Records (Date Range)",
-      "RFS Suspension Certification # (N/A if not applicable)",
-      "Suspension System Maintenance (Date Range)",
-      "Trip Records (Date Range)",
-      "Fault Recording/ Reporting on Suspension System (Date Range)"
-    ]
-  },
-
-  # Table 12: Operator's Name (legal entity) - Signature block
-  {
-    "name": "Operator’s Name (legal entity)",
-    "labels": ["Operator’s Name (legal entity)"]
-  },
-
-  # Table 13: Non-conformance type
-  {
-    "name": "Non-conformance type (please tick)",
-    "labels": ["Un-conditional", "Conditional"]
-  },
-
-  # Table 14: Non-conformance Information
-  {
-    "name": "Non-conformance Information",
-    "labels_on_row1": True,
-    "labels": [
-      "Non-conformance agreed close out date",
-      "Module and Standard",
-      "Corrective Action Request (CAR) Number"
-    ]
-  },
-
-  # Table 15: Non-conformance and action taken
-  {
-    "name": "Non-conformance and action taken",
-    "labels_on_row1": True,
-    "labels": [
-      "Observed Non-conformance:",
-      "Corrective Action taken or to be taken by operator:",
-      "Operator or Representative Signature", "Position", "Date"
-    ]
-  },
-
-  # Table 16: Print Name / Auditor Reg Number
-  {
-    "name": "Print Name / Auditor Reg Number",
-    "labels_on_row1": True,
-    "labels": [
-      "Print Name",
-      "NHVR or Exemplar Global Auditor Registration Number"
-    ]
-  },
-
-  # Table 17: Audit Declaration
-  {
-    "name": "Audit Declaration",
-    "labels_on_left": True,
-    "labels": [
-      "Audit was conducted on",
-      "Unconditional CARs closed out on:",
-      "Conditional CARs to be closed out by:"
-    ]
-  },
-
-  # Table 18: print accreditation name
-  {
-    "name": "print accreditation name",
-    "labels": ["print accreditation name"]
-  },
-
-  # Table 19: Operator Declaration
-  {
-    "name": "Operator Declaration",
-    "labels_on_row1": True,
-    "labels": ["Print Name", "Position Title"]
-  }
-]
-
-def extract_red_text(path):
-    doc = Document(path)
-
-    # debug print
-    print(f"Found {len(doc.tables)} tables:")
-    for i,t in enumerate(doc.tables):
-        print(f"  Table#{i}: “{t.rows[0].cells[0].text.strip()[:30]}…”")
-    print()
-
-    out = {}
-    for ti, spec in enumerate(MASTER_TABLES):
-        if ti >= len(doc.tables):
-            break
-        tbl = doc.tables[ti]
-        name = spec["name"]
-
-        # prepare container & dedupe sets
-        collected = {lbl:[] for lbl in spec["labels"]}
-        seen = {lbl:set() for lbl in spec["labels"]}
-
-        # choose orientation
-        if spec.get("labels_on_row1"):
-            headers   = spec["labels"]
-            rows      = tbl.rows[1:]
-            col_mode  = True
-        elif spec.get("labels_on_left"):
-            headers   = spec["labels"]
-            # skip any unwanted header/subheading rows
-            rows = [
-                row for row in tbl.rows[1:]
-                if row.cells[0].text.strip() not in spec.get("skip_rows",[])
-            ]
-            col_mode = False
-        else:
-            headers   = [name]
-            rows      = tbl.rows
-            col_mode  = None
-
-        # scan each cell
-        for ri,row in enumerate(rows):
-            for ci,cell in enumerate(row.cells):
-                red = "".join(
-                    run.text for p in cell.paragraphs for run in p.runs
-                    if is_red_font(run)
-                ).strip()
-                if not red: continue
+def _prev_para_text(tbl):
+    """Get text from previous paragraph before table"""
+    prev = tbl._tbl.getprevious()
+    while prev is not None and not prev.tag.endswith("}p"):
+        prev = prev.getprevious()
+    if prev is None:
+        return ""
+    return "".join(node.text for node in prev.iter() if node.tag.endswith("}t") and node.text).strip()
+
+def normalize_text(text):
+    """Normalize text for better matching"""
+    return re.sub(r'\s+', ' ', text.strip())
+
+def fuzzy_match_heading(heading, patterns):
+    """Check if heading matches any pattern with fuzzy matching"""
+    heading_norm = normalize_text(heading.upper())
+    for pattern in patterns:
+        if re.search(pattern, heading_norm, re.IGNORECASE):
+            return True
+    return False
 
-                # assign label
-                if   col_mode is True:
-                    lbl = headers[ci] if ci<len(headers) else name
-                elif col_mode is False:
-                    lbl = headers[ri] if ri<len(headers) else name
+def get_table_context(tbl):
+    """Get comprehensive context information for table"""
+    heading = normalize_text(_prev_para_text(tbl))
+    headers = [normalize_text(c.text) for c in tbl.rows[0].cells if c.text.strip()]
+    col0 = [normalize_text(r.cells[0].text) for r in tbl.rows if r.cells[0].text.strip()]
+    first_cell = normalize_text(tbl.rows[0].cells[0].text) if tbl.rows else ""
+    all_cells = []
+    for row in tbl.rows:
+        for cell in row.cells:
+            text = normalize_text(cell.text)
+            if text:
+                all_cells.append(text)
+    return {
+        'heading': heading,
+        'headers': headers,
+        'col0': col0,
+        'first_cell': first_cell,
+        'all_cells': all_cells,
+        'num_rows': len(tbl.rows),
+        'num_cols': len(tbl.rows[0].cells) if tbl.rows else 0
+    }
+
+def calculate_schema_match_score(schema_name, spec, context):
+    """Calculate match score for a schema against table context"""
+    score = 0
+    reasons = []
+    if context['first_cell'] and context['first_cell'].upper() == schema_name.upper():
+        score += 100
+        reasons.append(f"Direct first cell match: '{context['first_cell']}'")
+    if spec.get("headings"):
+        for h in spec["headings"]:
+            if fuzzy_match_heading(context['heading'], [h["text"]]):
+                score += 50
+                reasons.append(f"Heading match: '{context['heading']}'")
+                break
+    if spec.get("orientation") == "left":
+        labels = [normalize_text(lbl) for lbl in spec["labels"]]
+        matches = 0
+        for lbl in labels:
+            if any(lbl.upper() in c.upper() or c.upper() in lbl.upper() for c in context['col0']):
+                matches += 1
+        if matches > 0:
+            score += (matches / len(labels)) * 30
+            reasons.append(f"Left orientation label matches: {matches}/{len(labels)}")
+    elif spec.get("orientation") == "row1":
+        labels = [normalize_text(lbl) for lbl in spec["labels"]]
+        matches = 0
+        for lbl in labels:
+            if any(lbl.upper() in h.upper() or h.upper() in lbl.upper() for h in context['headers']):
+                matches += 1
+        if matches > 0:
+            score += (matches / len(labels)) * 30
+            reasons.append(f"Row1 orientation header matches: {matches}/{len(labels)}")
+    if spec.get("columns"):
+        cols = [normalize_text(col) for col in spec["columns"]]
+        matches = 0
+        for col in cols:
+            if any(col.upper() in h.upper() for h in context['headers']):
+                matches += 1
+        if matches == len(cols):
+            score += 40
+            reasons.append(f"All column headers match: {cols}")
+    if schema_name == "Operator Declaration" and context['first_cell'].upper() == "PRINT NAME":
+        if "OPERATOR DECLARATION" in context['heading'].upper():
+            score += 80
+            reasons.append("Operator Declaration context match")
+        elif any("MANAGER" in cell.upper() for cell in context['all_cells']):
+            score += 60
+            reasons.append("Manager found in cells (likely Operator Declaration)")
+    if schema_name == "NHVAS Approved Auditor Declaration" and context['first_cell'].upper() == "PRINT NAME":
+        if any("MANAGER" in cell.upper() for cell in context['all_cells']):
+            score -= 50  # Penalty because auditors shouldn't be managers
+            reasons.append("Penalty: Manager found (not auditor)")
+    return score, reasons
+
+def match_table_schema(tbl):
+    """Improved table schema matching with scoring system"""
+    context = get_table_context(tbl)
+    best_match = None
+    best_score = 0
+    for name, spec in TABLE_SCHEMAS.items():
+        score, reasons = calculate_schema_match_score(name, spec, context)
+        if score > best_score:
+            best_score = score
+            best_match = name
+    if best_score >= 20:
+        return best_match
+    return None
+
+def check_multi_schema_table(tbl):
+    """Check if table contains multiple schemas and split appropriately"""
+    context = get_table_context(tbl)
+    operator_labels = ["Operator name (Legal entity)", "NHVAS Accreditation No.", "Registered trading name/s", 
+                      "Australian Company Number", "NHVAS Manual"]
+    contact_labels = ["Operator business address", "Operator Postal address", "Email address", "Operator Telephone Number"]
+    has_operator = any(any(op_lbl.upper() in cell.upper() for op_lbl in operator_labels) for cell in context['col0'])
+    has_contact = any(any(cont_lbl.upper() in cell.upper() for cont_lbl in contact_labels) for cell in context['col0'])
+    if has_operator and has_contact:
+        return ["Operator Information", "Operator contact details"]
+    return None
+
+def extract_multi_schema_table(tbl, schemas):
+    """Extract data from table with multiple schemas"""
+    result = {}
+    for schema_name in schemas:
+        if schema_name not in TABLE_SCHEMAS:
+            continue
+        spec = TABLE_SCHEMAS[schema_name]
+        schema_data = {}
+        for ri, row in enumerate(tbl.rows):
+            if ri == 0:
+                continue
+            row_label = normalize_text(row.cells[0].text)
+            belongs_to_schema = False
+            matched_label = None
+            for spec_label in spec["labels"]:
+                spec_norm = normalize_text(spec_label).upper()
+                row_norm = row_label.upper()
+                if spec_norm == row_norm or spec_norm in row_norm or row_norm in spec_norm:
+                    belongs_to_schema = True
+                    matched_label = spec_label
+                    break
+            if not belongs_to_schema:
+                continue
+            for ci, cell in enumerate(row.cells):
+                red_txt = "".join(run.text for p in cell.paragraphs for run in p.runs if is_red_font(run)).strip()
+                if red_txt:
+                    if matched_label not in schema_data:
+                        schema_data[matched_label] = []
+                    if red_txt not in schema_data[matched_label]:
+                        schema_data[matched_label].append(red_txt)
+        if schema_data:
+            result[schema_name] = schema_data
+    return result
+
+def extract_table_data(tbl, schema_name, spec):
+    """Extract red text data from table based on schema"""
+    labels = spec["labels"] + [schema_name]
+    collected = {lbl: [] for lbl in labels}
+    seen = {lbl: set() for lbl in labels}
+    by_col = (spec["orientation"] == "row1")
+    start_row = 1 if by_col else 0
+    rows = tbl.rows[start_row:]
+    for ri, row in enumerate(rows):
+        for ci, cell in enumerate(row.cells):
+            red_txt = "".join(run.text for p in cell.paragraphs for run in p.runs if is_red_font(run)).strip()
+            if not red_txt:
+                continue
+            if by_col:
+                if ci < len(spec["labels"]):
+                    lbl = spec["labels"][ci]
                 else:
-                    lbl = name
-
-                # dedupe & collect
-                if red not in seen[lbl]:
-                    seen[lbl].add(red)
-                    collected[lbl].append(red)
-
-        # only keep non-empty labels
-        filtered = {l:collected[l] for l in collected if collected[l]}
-        if filtered:
-            out[name] = filtered
-
-    # paragraphs
+                    lbl = schema_name
+            else:
+                raw_label = normalize_text(row.cells[0].text)
+                lbl = None
+                for spec_label in spec["labels"]:
+                    if normalize_text(spec_label).upper() == raw_label.upper():
+                        lbl = spec_label
+                        break
+                if not lbl:
+                    for spec_label in spec["labels"]:
+                        spec_norm = normalize_text(spec_label).upper()
+                        raw_norm = raw_label.upper()
+                        if spec_norm in raw_norm or raw_norm in spec_norm:
+                            lbl = spec_label
+                            break
+                if not lbl:
+                    lbl = schema_name
+            if red_txt not in seen[lbl]:
+                seen[lbl].add(red_txt)
+                collected[lbl].append(red_txt)
+    return {k: v for k, v in collected.items() if v}
+
+def extract_red_text(input_doc):
+    # input_doc: docx.Document object or file path
+    if isinstance(input_doc, str):
+        doc = Document(input_doc)
+    else:
+        doc = input_doc
+    out = {}
+    table_count = 0
+    for tbl in doc.tables:
+        table_count += 1
+        multi_schemas = check_multi_schema_table(tbl)
+        if multi_schemas:
+            multi_data = extract_multi_schema_table(tbl, multi_schemas)
+            for schema_name, schema_data in multi_data.items():
+                if schema_data:
+                    if schema_name in out:
+                        for k, v in schema_data.items():
+                            if k in out[schema_name]:
+                                out[schema_name][k].extend(v)
+                            else:
+                                out[schema_name][k] = v
+                    else:
+                        out[schema_name] = schema_data
+            continue
+        schema = match_table_schema(tbl)
+        if not schema:
+            continue
+        spec = TABLE_SCHEMAS[schema]
+        data = extract_table_data(tbl, schema, spec)
+        if data:
+            if schema in out:
+                for k, v in data.items():
+                    if k in out[schema]:
+                        out[schema][k].extend(v)
+                    else:
+                        out[schema][k] = v
+            else:
+                out[schema] = data
     paras = {}
-    for i,para in enumerate(doc.paragraphs):
-        red = "".join(r.text for r in para.runs if is_red_font(r)).strip()
-        if not red: continue
-        # find nearest non-red above
-        lab = None
-        for j in range(i-1,-1,-1):
-            if any(is_red_font(r) for r in doc.paragraphs[j].runs):
-                continue
-            txt = doc.paragraphs[j].text.strip()
+    for idx, para in enumerate(doc.paragraphs):
+        red_txt = "".join(r.text for r in para.runs if is_red_font(r)).strip()
+        if not red_txt:
+            continue
+        context = None
+        for j in range(idx-1, -1, -1):
+            txt = normalize_text(doc.paragraphs[j].text)
             if txt:
-                lab = txt; break
-        key = lab or "(para)"
-        paras.setdefault(key,[]).append(red)
-
+                all_patterns = HEADING_PATTERNS["main"] + HEADING_PATTERNS["sub"]
+                if any(re.search(p, txt, re.IGNORECASE) for p in all_patterns):
+                    context = txt
+                    break
+        if not context and re.fullmatch(PARAGRAPH_PATTERNS["date_line"], red_txt):
+            context = "Date"
+        if not context:
+            context = "(para)"
+        paras.setdefault(context, []).append(red_txt)
     if paras:
         out["paragraphs"] = paras
     return out
 
-if __name__=="__main__":
-    import sys, json
-    # Usage: python extract_red_text.py input.docx output.json
-    input_docx = sys.argv[1]
-    output_json = sys.argv[2]
-    word_data = extract_red_text(input_docx)
-    with open(output_json, 'w', encoding='utf-8') as f:
-        json.dump(word_data, f, indent=2, ensure_ascii=False)
-
-    # still print to console for immediate feedback
-    print(json.dumps(word_data, indent=2, ensure_ascii=False))
+def extract_red_text_filelike(input_file, output_file):
+    """
+    Accepts:
+      input_file: file-like object (BytesIO/File) or path
+      output_file: file-like object (opened for writing text) or path
+    """
+    if hasattr(input_file, "seek"):
+        input_file.seek(0)
+    doc = Document(input_file)
+    result = extract_red_text(doc)
+    if hasattr(output_file, "write"):
+        json.dump(result, output_file, indent=2, ensure_ascii=False)
+        output_file.flush()
+    else:
+        with open(output_file, "w", encoding="utf-8") as f:
+            json.dump(result, f, indent=2, ensure_ascii=False)
+    return result
+
+if __name__ == "__main__":
+    # Support both script and app/file-like usage
+    if len(sys.argv) == 3:
+        input_docx = sys.argv[1]
+        output_json = sys.argv[2]
+        doc = Document(input_docx)
+        word_data = extract_red_text(doc)
+        with open(output_json, 'w', encoding='utf-8') as f:
+            json.dump(word_data, f, indent=2, ensure_ascii=False)
+        print(json.dumps(word_data, indent=2, ensure_ascii=False))
+    else:
+        print("To use as a module: extract_red_text_filelike(input_file, output_file)")