Spaces:
Sleeping
Sleeping
File size: 5,690 Bytes
14c6c4f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | from docx import Document
from src.detectors.hybrid_detector import HybridDetector
from src.redactors.mapper import StatefulMapper
class DocxHandler:
"""
Handles reading, redacting, and saving DOCX files safely without destroying
document structure, tables, headers, or footers.
"""
def __init__(self, detector: HybridDetector, mapper: StatefulMapper):
self.detector = detector
self.mapper = mapper
def process_document(self, input_path: str, output_path: str):
"""
Loads a document, redacts PII across all sections, and saves it.
"""
doc = Document(input_path)
# 1. Process Core Paragraphs
self._process_paragraphs(doc.paragraphs)
# 2. Process Tables
for table in doc.tables:
for row_idx, row in enumerate(table.rows):
for cell_idx, cell in enumerate(row.cells):
is_address_cell = False
if row_idx > 0:
try:
above_cell = table.cell(row_idx - 1, cell_idx)
if above_cell.text.strip().upper() in ["REGISTERED OFFICE", "CORPORATE OFFICE"]:
is_address_cell = True
except Exception:
pass
self._process_paragraphs(cell.paragraphs, is_address_cell=is_address_cell)
# 3. Process Headers and Footers
for section in doc.sections:
self._process_paragraphs(section.header.paragraphs)
self._process_paragraphs(section.footer.paragraphs)
doc.save(output_path)
print(f"Redacted document successfully saved to {output_path}")
def _process_paragraphs(self, paragraphs, is_address_cell=False):
"""
Iterates over a list of paragraphs and redacts PII within them.
"""
for paragraph in paragraphs:
runs_text = "".join(r.text for r in paragraph.runs)
if not runs_text.strip():
continue
if is_address_cell:
entities = [{"entity_type": "ADDRESS", "start": 0, "end": len(runs_text), "text": runs_text}]
else:
# Detect entities in the run-reconstructed paragraph text
entities = self.detector.analyze(runs_text)
if not entities:
continue
# We process entities in reverse order (right to left) so that modifying
# the text of a run does not invalidate the start/end character indices
# of entities that appear earlier in the paragraph.
entities.sort(key=lambda x: x['start'], reverse=True)
self._redact_runs(paragraph, entities)
def _redact_runs(self, paragraph, entities):
"""
Replaces detected PII across Word 'Runs' safely.
A single paragraph is composed of multiple runs (sections of text with consistent formatting).
An entity might span across multiple runs, so we must map character indices to runs.
"""
for entity in entities:
e_start = entity['start']
e_end = entity['end']
e_type = entity['entity_type']
e_text = entity['text']
fake_text = self.mapper.get_replacement(e_type, e_text)
# Re-calculate run indices dynamically just to be safe,
# since we are modifying runs in-place.
# Because we process right-to-left, the indices to the left remain intact.
current_len = 0
runs_info = []
for idx, run in enumerate(paragraph.runs):
runs_info.append({
"index": idx,
"run": run,
"start": current_len,
"end": current_len + len(run.text),
"text": run.text
})
current_len += len(run.text)
# Find which runs the entity overlaps with
overlapping_runs = []
for r_info in runs_info:
# Overlap condition:
# Entity starts before run ends AND Entity ends after run starts
if e_start < r_info['end'] and e_end > r_info['start']:
overlapping_runs.append(r_info)
if not overlapping_runs:
continue
# If the entity is entirely contained in one run
if len(overlapping_runs) == 1:
r_info = overlapping_runs[0]
run = r_info['run']
local_start = e_start - r_info['start']
local_end = e_end - r_info['start']
# Replace the text inside the run
run.text = run.text[:local_start] + fake_text + run.text[local_end:]
# If the entity spans multiple runs
else:
first_r = overlapping_runs[0]
last_r = overlapping_runs[-1]
local_start = e_start - first_r['start']
local_end = e_end - last_r['start']
# Place the replacement text in the first run, keep the text before the entity
first_r['run'].text = first_r['run'].text[:local_start] + fake_text
# Clear all intermediate runs completely
for r_info in overlapping_runs[1:-1]:
r_info['run'].text = ""
# Keep the text after the entity in the last run
last_r['run'].text = last_r['run'].text[local_end:]
|