| 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 |
|
|
