MrKAMELEON/ReferencesExtractor

1import re2from pathlib import Path3 4from PyPDF2 import PdfReader5 6 7def load_pdf_content(filepath: str | Path) -> str:8    """9    Load and extract text content from a PDF file.10 11    Args:12        filepath: Path to the PDF file, can be string or Path object13 14    Returns:15        Extracted text content from the PDF16 17    Raises:18        FileNotFoundError: If the PDF file doesn't exist19        ValueError: If the file is not a valid PDF20    """21    pdf_path = Path(filepath)22    if not pdf_path.exists():23        raise FileNotFoundError(f"PDF file not found: {pdf_path}")24 25    with pdf_path.open("rb") as file:26        reader = PdfReader(file)27        content = []28        for page in reader.pages:29            page_lines = page.extract_text().split("\n")30            page_lines = page_lines[1:]  # remove first line (page number & author)31            content.append("\n".join(page_lines))32    return "\n".join(content)33 34 35def save_references(output_path: Path, content: str) -> bool:36    """37    Save references content to a text file.38 39    Args:40        output_path: Path where to save the references41        content: References content to save42 43    Returns:44        True if save was successful, False otherwise45    """46    try:47        output_path.parent.mkdir(exist_ok=True)48 49        # Extract filename from path for header50        content_header = output_path.name.replace('.txt', '')51        save_content = f"{content_header}\n\n{content}"52 53        output_path.write_text(save_content, encoding="utf-8")54        return True55 56    except Exception as e:57        print(f"Error saving references to file: {str(e)}")58        return False59 60 61def fetch_files(input_directory: str | Path, file_type: str = ".pdf") -> list[str]:62    """63    Get list of files with specified extension from the input directory.64 65    Args:66        input_directory: Directory to search for files67        file_type: File extension to filter by (default: '.pdf')68 69    Returns:70        List of filenames matching the file_type71    """72    dir_path = Path(input_directory)73    dir_path.mkdir(exist_ok=True)74 75    return [76        f.name77        for f in dir_path.iterdir()78        if f.is_file() and f.suffix.lower() == file_type.lower()79    ]80 81 82def rfind_references_index(text: str, term_to_find: str = "references") -> int:83    """84    Finds the rightmost index of a given term (or its variations with whitespaces)85    in a text, similar to rfind.86 87    Args:88    - text (str): The input string to search within.89    - term_to_find (str): The term to search for (e.g., 'References', 'Bibliographie').90 91    Returns:92    - int: The starting index of the rightmost match, or -1 if not found.93    """94    # Build the regex pattern dynamically based on the term_to_find95    # Escape special characters in the term to avoid regex interpretation issues96    escaped_term = re.escape(term_to_find)97 98    # Insert '\s*' (zero or more whitespaces) between each character of the term99    pattern_string = r'[\s.,;:\'"\-_!?()]*'.join(list(escaped_term))100 101    # Compile the regex pattern with IGNORECASE flag102    pattern = re.compile(pattern_string, re.IGNORECASE)103 104    # Find all matches in the text105    matches = list(pattern.finditer(text))106 107    # If there are matches, return the start index of the last one108    if matches:109        return matches[-1].start()110    else:111        return -1112 113 114def find_references(content: str) -> str:115    """116    Find the references section in the last 40% of the text content.117 118    Args:119        content (str): The full text content to search in120 121    Returns:122        str: Cropped_content if found123    """124    references_index_relative = -1125    keyword_texts = [126        "References",127        "Bibliographie",128        "Referenzen",129        "Quellenangaben",130        "Bibliografia",131        "Literatura",132    ]133    for keyword in keyword_texts:134        index = rfind_references_index(content, keyword)135        references_index_relative = max(references_index_relative, index)136 137    if references_index_relative == -1:138        return "REFERENCES HEADER NOT FOUND " + content139 140    min_reference_skip_length = 10141    start_search_index = (142        references_index_relative + min_reference_skip_length143    )144    references_end_index = content.find(" ", start_search_index)145    cropped_content = content[references_end_index:]146    return cropped_content147 148