Spaces:
Running on Zero
Running on Zero
File size: 2,172 Bytes
3998cca | 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 | """Read the first pages of a file to learn its text and page count. Sync: run in a thread."""
import html
import io
import logging
import os
import re
import zipfile
from dataclasses import dataclass
from typing import Optional
log = logging.getLogger(__name__)
@dataclass
class DocInfo:
text: str = ""
pages: Optional[int] = None
def _pdf(data: bytes) -> DocInfo:
try:
from pypdf import PdfReader
except ImportError:
return DocInfo()
logging.getLogger("pypdf").setLevel(logging.ERROR)
reader = PdfReader(io.BytesIO(data))
pages = len(reader.pages)
chunks = []
for i in range(min(3, pages)):
chunks.append(reader.pages[i].extract_text() or "")
return DocInfo(text=" ".join(chunks), pages=pages)
def _pptx(data: bytes) -> DocInfo:
with zipfile.ZipFile(io.BytesIO(data)) as z:
slides = sorted(
(n for n in z.namelist() if re.fullmatch(r"ppt/slides/slide\d+\.xml", n)),
key=lambda n: int(re.findall(r"\d+", n)[0]),
)
chunks = []
for name in slides[:3]:
xml = z.read(name).decode("utf-8", "ignore")
chunks.append(" ".join(html.unescape(t) for t in re.findall(r"<a:t>([^<]*)</a:t>", xml)))
return DocInfo(text=" ".join(chunks), pages=len(slides))
def _docx(data: bytes) -> DocInfo:
with zipfile.ZipFile(io.BytesIO(data)) as z:
xml = z.read("word/document.xml").decode("utf-8", "ignore")
text = " ".join(html.unescape(t) for t in re.findall(r"<w:t[^>]*>([^<]*)</w:t>", xml))
return DocInfo(text=text)
def analyze(file_name: Optional[str], data: bytes) -> DocInfo:
ext = os.path.splitext(file_name or "")[1].lower()
try:
if ext == ".pdf":
info = _pdf(data)
elif ext == ".pptx":
info = _pptx(data)
elif ext == ".docx":
info = _docx(data)
else:
return DocInfo()
except Exception as exc: # encrypted / scanned / corrupt files must never break an upload
log.info("Could not read %s: %r", file_name, exc)
return DocInfo()
info.text = re.sub(r"\s+", " ", info.text)[:6000]
return info
|