aps-agent/kangni-analysis/extract.py

61 lines
1.9 KiB
Python

# -*- coding: utf-8 -*-
"""Extract text + tables from all docx files in this folder to .txt files."""
import os
from docx import Document
BASE = os.path.dirname(os.path.abspath(__file__))
def iter_block_items(doc):
from docx.oxml.ns import qn
body = doc.element.body
for child in body.iterchildren():
if child.tag == qn('w:p'):
yield ('p', child)
elif child.tag == qn('w:tbl'):
yield ('tbl', child)
def extract(path):
doc = Document(path)
out = []
for kind, el in iter_block_items(doc):
if kind == 'p':
from docx.text.paragraph import Paragraph
p = Paragraph(el, doc)
t = p.text.strip()
if t:
style = p.style.name if p.style else ''
prefix = '#' * 0
if style.startswith('Heading'):
try:
lvl = int(style.split()[-1])
prefix = '#' * lvl + ' '
except Exception:
prefix = '# '
out.append(prefix + t)
else:
from docx.table import Table
tbl = Table(el, doc)
out.append('[TABLE]')
for row in tbl.rows:
cells = [c.text.strip().replace('\n', ' / ') for c in row.cells]
# dedupe merged cells
dedup = []
prev = None
for c in cells:
dedup.append(c)
out.append(' | '.join(dedup))
out.append('[/TABLE]')
return '\n'.join(out)
for fn in os.listdir(BASE):
if fn.lower().endswith('.docx'):
src = os.path.join(BASE, fn)
dst = os.path.join(BASE, fn + '.txt')
try:
text = extract(src)
with open(dst, 'w', encoding='utf-8') as f:
f.write(text)
print(f'{fn}: {len(text)} chars')
except Exception as e:
print(f'{fn}: ERROR {e}')