processing.py

from abc import ABC, abstractmethod
from pathlib import Path
from urllib.parse import urlparse
import requests
import fitz
import io
import re
import hashlib
import os


class FileTypeError(Exception):
    """Raised when the file type does not match the expected file type."""


class FileSchemeError(Exception):
    """Raised when the file scheme does not match the expected file scheme."""


class FileProcessor(ABC):
    type = None

    def __init__(self, path):
        self.path = path
        self.file_scheme = self._get_file_scheme()
        self.__class__._check_file_type(path)

    @abstractmethod
    def get_file_data(self):
        pass

    @abstractmethod
    def _get_file_metadata(self):
        pass

    @abstractmethod
    def _get_file_paragraphs(self):
        pass

    @classmethod
    def _check_file_type(cls, path):
        file_type = Path(path).suffix.lower()[1:]
        if file_type != cls.type:
            raise FileTypeError(
                f"Invalid file type. {cls.__name__} expects a {cls.type} file"
            )

    def _get_file_scheme(self):
        parsed_path = urlparse(self.path)
        if (
            not parsed_path.scheme
            or parsed_path.scheme.lower() == "file"
            or os.path.isfile(self.path)
        ):
            return "local"
        elif parsed_path.scheme.lower() in ["http", "https", "ftp"]:
            return "url"
        else:
            raise FileSchemeError("Unknown scheme")

    def _preprocess_text(self, text):
        text = text.replace("\n", " ")
        text = re.sub("\s+", " ", text)
        text = text.encode("utf-8", "ignore").decode("utf-8", "ignore")
        return text

    def _generate_hash(self, string):
        hash_object = hashlib.md5()
        hash_object.update(string.encode("utf-8", "ignore"))
        hex_dig = hash_object.hexdigest()

        return hex_dig

    def generate_paragraphs():
        raise NotImplementedError

    def generate_metadata():
        raise NotImplementedError


class PDFProcessor(FileProcessor):
    type = "pdf"

    def __init__(self, path):
        super().__init__(path)

    def get_file_data(self, merge_length=200):
        file = self._open_file()

        file_metadata = self._get_file_metadata(file)
        file_paragraphs = self._get_file_paragraphs(
            file, file_metadata, start_page=1, end_page=None, merge_length=merge_length
        )

        file.close()

        return file_metadata, file_paragraphs

    def _get_file_metadata(self, file):
        file_metadata = dict()

        metadata = file.metadata

        unique_string = str(Path(self.path).name) + metadata["title"]

        file_metadata["id"] = self._generate_hash(unique_string)
        file_metadata["title"] = metadata["title"]
        file_metadata["author"] = metadata["author"]
        file_metadata["subject"] = metadata["subject"]
        file_metadata["creation_date"] = metadata["creationDate"]
        file_metadata["modification_date"] = metadata["modDate"]
        file_metadata["n_pages"] = file.page_count
        if self.file_scheme == "local":
            file_metadata["url"] = str(Path(self.path).resolve())
        else:
            file_metadata["url"] = self.path
        file_metadata["file_name"] = Path(self.path).name
        file_metadata["short_name"] = Path(self.path).name
        file_metadata["release_date"] = ""
        file_metadata["report_type"] = ""
        file_metadata["source"] = ""

        return file_metadata

    def _get_file_paragraphs(
        self, file, file_metadata, start_page=1, end_page=None, merge_length=200
    ):
        if end_page is None:
            end_page = file_metadata["n_pages"]

        file_paragraphs = []

        for page_num in range(start_page - 1, end_page):
            page = file.load_page(page_num)
            blocks = page.get_text("blocks")

            for block in blocks:
                paragraph = self._process_block(
                    block, page, page_num + start_page, file_metadata["id"]
                )
                if paragraph is None:
                    continue

                first_char = paragraph["content"][0]
                if len(file_paragraphs) > 0:
                    if (
                        len(file_paragraphs[-1]["content"]) + len(paragraph["content"])
                        < merge_length
                    ) or (first_char.islower() and first_char.isalpha()):
                        file_paragraphs[-1]["content"] += " " + paragraph["content"]
                        file_paragraphs[-1]["length"] = len(
                            file_paragraphs[-1]["content"]
                        )
                    else:
                        file_paragraphs.append(paragraph)
                else:
                    file_paragraphs.append(paragraph)

        return file_paragraphs

    def _open_file(self):
        if self.file_scheme == "url":
            response = requests.get(self.path)
            file = fitz.open(stream=io.BytesIO(response.content), filetype="pdf")
        elif self.file_scheme == "local":
            file = fitz.open(self.path)
        return file

    def _process_block(self, block, page, page_number, file_id):
        x0, y0, x1, y1, content, block_no, block_type = block

        if content.isspace() or block_type == 1:
            return None

        content = self._preprocess_text(content)
        unique_content_string = "_".join(map(str, block))
        paragraph_id = self._generate_hash(unique_content_string)

        w, h = page.rect.width, page.rect.height
        paragraph = {
            "id": paragraph_id,
            "document_id": file_id,
            "content_type": "text" if block_type == 0 else "image",
            "content": content,
            "length": len(content),
            "idx_block": block_no,
            "page_number": page_number,
            "x0": x0 / h,
            "y0": y0 / w,
            "x1": x1 / h,
            "y1": y1 / w,
        }

        return paragraph


class HTMLProcessor(FileProcessor):
    type = "html"

    def __init__(self, path):
        super().__init__(path)

    def get_file_data(self):
        pass

    def _get_file_metadata(self):
        pass

    def _get_file_paragraphs(self):
        pass

    def _open_file(self):
        if self.file_scheme == "url":
            response = requests.get(self.path)
            file = response.text
        elif self.file_scheme == "local":
            file = open(self.path, "r").read()
        return file