Spaces:
Running on Zero
Running on Zero
| """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__) | |
| 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 | |