pdf-to-markdown-cli 0.2.0__py3-none-any.whl

docs_to_md/pdf/splitter.py

@@ -0,0 +1,199 @@
+import logging
+import os
+import uuid
+from pathlib import Path
+from typing import List, Optional
+
+import pikepdf
+from pydantic import BaseModel
+from tqdm import tqdm
+
+from docs_to_md.utils.exceptions import PDFProcessingError
+from docs_to_md.utils.file_utils import ensure_directory
+from docs_to_md.utils.logging import ProgressTracker
+
+logger = logging.getLogger(__name__)
+
+
+def _sanitize_path(path: Path) -> Path:
+    """
+    Sanitize path to prevent path traversal attacks.
+    
+    Args:
+        path: Path to sanitize
+        
+    Returns:
+        Sanitized path
+        
+    Raises:
+        PDFProcessingError: If path contains suspicious patterns
+    """
+    try:
+        # Convert to absolute to resolve any .. or other path manipulations
+        path = path.absolute().resolve()
+        
+        # Get the filename only, discarding any directory components
+        safe_name = path.name
+        
+        # Check for suspicious patterns
+        if safe_name != path.name or ".." in str(path):
+            raise PDFProcessingError(f"Potentially unsafe path: {path}")
+            
+        return path
+    except Exception as e:
+        if isinstance(e, PDFProcessingError):
+            raise
+        raise PDFProcessingError(f"Failed to sanitize path {path}: {e}")
+
+
+class PDFChunkInfo(BaseModel):
+    """Information about a single PDF chunk."""
+    path: str
+    index: int
+    start_page: int
+    end_page: int
+
+
+class PDFChunks(BaseModel):
+    """Collection of PDF chunks."""
+    chunks: List[PDFChunkInfo]
+
+
+def _create_chunk(pdf: pikepdf.Pdf, chunks_dir: Path, chunk_num: int, num_chunks: int, start: int, end: int) -> str:
+    """
+    Create a single PDF chunk.
+    
+    Args:
+        pdf: Source PDF
+        chunks_dir: Directory to save chunks
+        chunk_num: Index of this chunk
+        num_chunks: Total number of chunks
+        start: Start page index (inclusive)
+        end: End page index (exclusive)
+        
+    Returns:
+        Path to created chunk file
+        
+    Raises:
+        PDFProcessingError: If chunk creation fails
+    """
+    chunk_pdf = pikepdf.Pdf.new()
+    try:
+        for i in range(start, end):
+            chunk_pdf.pages.append(pdf.pages[i])
+
+        # Ensure safe path creation
+        chunk_filename = f"{chunk_num+1:03d}of{num_chunks:03d}.pdf"
+        chunk_path = str(chunks_dir / chunk_filename)
+        
+        chunk_pdf.save(
+            chunk_path,
+            compress_streams=True,
+            object_stream_mode=pikepdf.ObjectStreamMode.generate
+        )
+        return chunk_path
+    except Exception as e:
+        raise PDFProcessingError(f"Failed to create PDF chunk {chunk_num+1}: {e}")
+    finally:
+        chunk_pdf.close()
+
+
+def _create_chunks(pdf: pikepdf.Pdf, path: Path, pages_per_chunk: int, tmp_dir: Path) -> PDFChunks:
+    """
+    Create multiple chunks from a PDF.
+    
+    Args:
+        pdf: Source PDF
+        path: Original PDF path (for naming)
+        pages_per_chunk: Number of pages per chunk
+        tmp_dir: Directory to save chunks
+        
+    Returns:
+        PDFChunks with information about created chunks
+        
+    Raises:
+        PDFProcessingError: If chunking fails
+    """
+    # Use provided temp directory
+    ensure_directory(tmp_dir)
+
+    num_chunks = (len(pdf.pages) + pages_per_chunk - 1) // pages_per_chunk
+    chunks: List[PDFChunkInfo] = []
+
+    progress = ProgressTracker(num_chunks, "Chunking PDF", "chunk")
+    
+    try:
+        for chunk_num in range(num_chunks):
+            start = chunk_num * pages_per_chunk
+            end = min(start + pages_per_chunk, len(pdf.pages))
+            chunk_path = _create_chunk(pdf, tmp_dir, chunk_num, num_chunks, start, end)
+            chunks.append(PDFChunkInfo(
+                path=chunk_path,
+                index=chunk_num,
+                start_page=start,
+                end_page=end-1
+            ))
+            progress.update()
+
+        if not chunks:
+            raise PDFProcessingError(f"Failed to create any chunks for {path}")
+
+        return PDFChunks(chunks=chunks)
+    
+    finally:
+        progress.close()
+
+
+def chunk_pdf_to_temp(pdf_path: str, pages_per_chunk: int = 10, tmp_dir: Optional[Path] = None) -> Optional[PDFChunks]:
+    """
+    Split a PDF into chunks of specified size and save to temp directory.
+
+    Args:
+        pdf_path: Path to the PDF file
+        pages_per_chunk: Number of pages per chunk (default: 10)
+        tmp_dir: Directory to save chunks in (default: creates one)
+
+    Returns:
+        PDFChunks containing information about each chunk, or None if no chunking needed
+
+    Raises:
+        PDFProcessingError: If the PDF is invalid or cannot be processed
+        ValueError: If pages_per_chunk < 1
+    """
+    if pages_per_chunk < 1:
+        raise ValueError("pages_per_chunk must be at least 1")
+
+    # Sanitize and validate path
+    path = _sanitize_path(Path(pdf_path))
+    
+    if not path.exists():
+        raise PDFProcessingError(f"File does not exist: {pdf_path}")
+
+    if path.suffix.lower() != '.pdf':
+        raise PDFProcessingError(f"File is not a PDF: {pdf_path}")
+
+    pdf = None
+    try:
+        pdf = pikepdf.Pdf.open(pdf_path)
+        if len(pdf.pages) == 0:
+            raise PDFProcessingError(f"PDF has no pages: {pdf_path}")
+
+        if len(pdf.pages) <= pages_per_chunk:
+            return None  # No chunking needed
+
+        # Create temp directory if not provided
+        if tmp_dir is None:
+            tmp_dir = Path("chunks") / f"{path.stem}_{uuid.uuid4().hex[:8]}"
+            ensure_directory(tmp_dir)
+
+        return _create_chunks(pdf, path, pages_per_chunk, tmp_dir)
+
+    except pikepdf.PdfError as e:
+        raise PDFProcessingError(f"Invalid PDF {pdf_path}: {e}")
+    except Exception as e:
+        if isinstance(e, PDFProcessingError):
+            raise
+        raise PDFProcessingError(f"Error processing PDF {pdf_path}: {e}")
+    finally:
+        if pdf is not None:
+            pdf.close() 

docs_to_md/storage/cache.py

@@ -0,0 +1,125 @@
+import logging
+import threading
+from pathlib import Path
+from typing import List, Optional
+
+from diskcache import Cache
+
+from docs_to_md.storage.models import ConversionRequest
+from docs_to_md.utils.exceptions import CacheError
+
+logger = logging.getLogger(__name__)
+
+
+class CacheManager:
+    """Handles persistence of conversion requests."""
+
+    def __init__(self, cache_dir: str):
+        """
+        Initialize cache manager with given directory.
+        
+        Args:
+            cache_dir: Directory for cache storage
+            
+        Raises:
+            CacheError: If cache initialization fails
+        """
+        try:
+            # Initialize with thread safety enabled
+            self.cache: Cache = Cache(cache_dir, statistics=True, timeout=60)
+            # Lock for thread-safe operations
+            self._lock = threading.RLock()
+        except Exception as e:
+            logger.error(f"Failed to initialize cache in {cache_dir}: {e}")
+            raise CacheError(f"Failed to initialize cache in {cache_dir}: {e}")
+
+    def save(self, request: ConversionRequest) -> None:
+        """
+        Save request to cache.
+        
+        Args:
+            request: Request to save
+            
+        Raises:
+            CacheError: If save operation fails
+        """
+        with self._lock:
+            try:
+                self.cache.set(request.request_id, request.model_dump())
+            except Exception as e:
+                logger.error(f"Failed to save request {request.request_id}: {e}")
+                raise CacheError(f"Failed to save request {request.request_id}: {e}")
+
+    def get(self, request_id: str) -> Optional[ConversionRequest]:
+        """
+        Get request from cache.
+        
+        Args:
+            request_id: ID of request to retrieve
+            
+        Returns:
+            ConversionRequest if found, None otherwise
+            
+        Raises:
+            CacheError: If retrieval fails
+        """
+        with self._lock:
+            try:
+                data = self.cache.get(request_id)
+                if data:
+                    return ConversionRequest.model_validate(data)
+                return None
+            except Exception as e:
+                logger.error(f"Failed to get request {request_id}: {e}")
+                return None
+
+    def delete(self, request_id: str) -> bool:
+        """
+        Delete request from cache.
+        
+        Args:
+            request_id: ID of request to delete
+            
+        Returns:
+            True if successful, False otherwise
+        """
+        with self._lock:
+            try:
+                return bool(self.cache.delete(request_id))
+            except Exception as e:
+                logger.error(f"Failed to delete request {request_id}: {e}")
+                return False
+
+    def get_all(self) -> List[ConversionRequest]:
+        """
+        Get all requests from cache.
+        
+        Returns:
+            List of all conversion requests
+        """
+        results = []
+        # Use a consistent view of the cache to avoid inconsistencies 
+        # during iteration
+        with self._lock:
+            keys = list(self.cache.iterkeys())
+            
+        for key in keys:
+            if request := self.get(str(key)):
+                results.append(request)
+        return results
+
+    def close(self) -> None:
+        """Close cache connection and free resources."""
+        with self._lock:
+            try:
+                self.cache.close()
+            except Exception as e:
+                logger.error(f"Error closing cache: {e}")
+            
+    def __enter__(self):
+        """Support for context manager."""
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Clean up resources when exiting context."""
+        self.close() 

docs_to_md/storage/models.py

@@ -0,0 +1,85 @@
+from enum import Enum
+from pathlib import Path
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class Status(str, Enum):
+    """Status of a request or chunk."""
+    PENDING = "pending"  # Waiting to be processed by the API
+    PROCESSING = "processing"  # Currently being processed by the API
+    COMPLETE = "complete"  # Successfully processed by the API
+    FAILED = "failed"  # Failed to process by the API
+
+
+class ChunkInfo(BaseModel):
+    """Information about a chunk or original file being processed."""
+    path: Path
+    index: int
+    request_id: Optional[str] = None
+    status: Status = Status.PENDING
+    error: Optional[str] = None
+
+    def mark_processing(self, request_id: str) -> None:
+        """Mark chunk as processing with given request ID."""
+        self.request_id = request_id
+        self.status = Status.PROCESSING
+
+    def mark_failed(self, error: str) -> None:
+        """Mark chunk as failed with error message."""
+        self.status = Status.FAILED
+        self.error = error
+
+    def mark_complete(self) -> None:
+        """Mark chunk as complete."""
+        self.status = Status.COMPLETE
+
+    def get_result_path(self, tmp_dir: Path) -> Path:
+        """Get the path where the result should be stored."""
+        return tmp_dir / f"{Path(self.path).name}.out"
+
+
+class ConversionRequest(BaseModel):
+    """Tracks a conversion request and its state."""
+    request_id: str
+    original_file: Path
+    target_file: Path
+    output_format: str = "markdown"
+    status: Status = Status.PENDING
+    error: Optional[str] = None
+    chunks: List[ChunkInfo] = []
+    chunk_size: int
+    tmp_dir: Optional[Path] = None  # Directory for temporary files for this conversion
+
+    def set_status(self, status: Status, error: Optional[str] = None) -> None:
+        """Set status and optional error message."""
+        self.status = status
+        if error:
+            self.error = error
+
+    def add_chunk(self, path: Path, index: int) -> ChunkInfo:
+        """Add a new chunk and return it."""
+        chunk = ChunkInfo(path=path, index=index)
+        self.chunks.append(chunk)
+        return chunk
+
+    @property
+    def pending_chunks(self) -> List[ChunkInfo]:
+        """Get all pending or processing chunks."""
+        return [c for c in self.chunks if c.status in (Status.PENDING, Status.PROCESSING)]
+
+    @property
+    def ordered_chunks(self) -> List[ChunkInfo]:
+        """Get chunks ordered by index."""
+        return sorted(self.chunks, key=lambda x: x.index)
+
+    @property
+    def has_failed(self) -> bool:
+        """Check if any chunks have failed."""
+        return any(c.status == Status.FAILED for c in self.chunks)
+
+    @property
+    def all_complete(self) -> bool:
+        """Check if all chunks are complete."""
+        return all(c.status == Status.COMPLETE for c in self.chunks) 

docs_to_md/utils/exceptions.py

@@ -0,0 +1,33 @@
+class DocsToMdError(Exception):
+    """Base exception for all docs-to-md errors."""
+    pass
+
+
+class ConfigurationError(DocsToMdError):
+    """Error related to configuration loading or validation."""
+    pass
+
+
+class FileError(DocsToMdError):
+    """Error related to file operations (reading, writing, discovery)."""
+    pass
+
+
+class APIError(DocsToMdError):
+    """Error related to the external conversion API communication."""
+    pass
+
+
+class PDFProcessingError(DocsToMdError):
+    """Error related to PDF manipulation (splitting, etc.)."""
+    pass
+
+
+class CacheError(DocsToMdError):
+    """Error related to cache operations."""
+    pass
+
+
+class ResultProcessingError(DocsToMdError):
+    """Error related to processing or combining results."""
+    pass 

docs_to_md/utils/file_utils.py

@@ -0,0 +1,189 @@
+import hashlib
+import logging
+import os
+import shutil
+import uuid
+from pathlib import Path
+from typing import List, Optional, Set
+
+import filetype
+
+from docs_to_md.utils.exceptions import FileError
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_directory(path: Path) -> None:
+    """Create directory if it doesn't exist."""
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def safe_delete(path: Path) -> None:
+    """Safely delete a file or directory."""
+    try:
+        path = Path(path)  # Convert to Path if string
+        if not path.exists():
+            return
+
+        if path.is_file():
+            path.unlink()
+        elif path.is_dir():
+            # Remove all files in directory recursively
+            shutil.rmtree(path, ignore_errors=True)
+    except Exception as e:
+        logger.error(f"Failed to delete {path}: {e}")
+
+
+def get_file_size(path: Path) -> int:
+    """Get file size in bytes."""
+    try:
+        return path.stat().st_size
+    except Exception as e:
+        logger.error(f"Error getting file size for {path}: {e}")
+        return 0
+
+
+def get_env_var(name: str, required: bool = True) -> Optional[str]:
+    """Get environment variable with optional requirement."""
+    value = os.getenv(name)
+    if required and not value:
+        raise FileError(f"Required environment variable {name} is not set")
+    return value
+
+
+class FileDiscovery:
+    """Handles finding and filtering files based on various criteria."""
+    
+    @staticmethod
+    def find_processable_files(input_path: Path, supported_types: Set[str], 
+                              supported_extensions: Optional[List[str]] = None) -> List[Path]:
+        """
+        Find all processable files from an input path.
+        
+        Args:
+            input_path: Directory or file path to search
+            supported_types: Set of supported MIME types
+            supported_extensions: List of supported file extensions (without dot)
+            
+        Returns:
+            List of processable file paths
+        """
+        if supported_extensions is None:
+            supported_extensions = ['.pdf', '.docx', '.doc', '.pptx', '.ppt', 
+                                   '.jpg', '.jpeg', '.png', '.gif', '.webp', '.tiff']
+        
+        files_to_process = []
+        
+        # Helper function to check file type without loading entire file
+        def check_file_type(file_path: Path) -> Optional[filetype.Type]:
+            try:
+                # Only read first 8KB for type detection rather than the whole file
+                with open(file_path, 'rb') as f:
+                    header = f.read(8192)  # First 8KB is enough for type detection
+                return filetype.guess(header)
+            except Exception as e:
+                logger.warning(f"Error reading file header for {file_path}: {e}")
+                return None
+        
+        # If input is a file, validate and return it
+        if input_path.is_file():
+            try:
+                kind = check_file_type(input_path)
+                if (kind and kind.mime in supported_types) or \
+                   (input_path.suffix.lower() in supported_extensions):
+                    files_to_process.append(input_path)
+                else:
+                    logger.warning(f"Unsupported file type: {input_path}")
+            except Exception as e:
+                logger.warning(f"Error checking file {input_path}: {e}")
+            
+            return files_to_process
+            
+        # If input is a directory, find all matching files
+        if not input_path.exists():
+            raise FileError(f"Input path does not exist: {input_path}")
+            
+        # Process directory
+        for p in input_path.glob("**/*"):
+            if p.is_file():
+                try:
+                    kind = check_file_type(p)
+                    if (kind and kind.mime in supported_types):
+                        files_to_process.append(p)
+                    elif p.suffix.lower() in supported_extensions:
+                        # Some file types might not be detected correctly by filetype
+                        files_to_process.append(p)
+                except Exception as e:
+                    logger.warning(f"Error checking file type for {p}: {e}")
+                    
+        return files_to_process
+
+
+class TemporaryDirectory:
+    """Context manager for temporary directories."""
+    
+    def __init__(self, base_path: Path, prefix: str):
+        """
+        Initialize a temporary directory.
+        
+        Args:
+            base_path: Base path where to create the temporary directory
+            prefix: Prefix for the directory name
+        """
+        self.path = base_path / f"{prefix}_{uuid.uuid4().hex[:8]}"
+        ensure_directory(self.path)
+        
+    def __enter__(self) -> Path:
+        """Enter the context and return the path to the temporary directory."""
+        return self.path
+        
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Clean up the temporary directory when exiting the context."""
+        safe_delete(self.path)
+
+
+class FileIO:
+    """Utility class for file operations."""
+    
+    @staticmethod
+    def read_file(path: Path) -> bytes:
+        """Read a file's content as bytes."""
+        try:
+            return path.read_bytes()
+        except Exception as e:
+            raise FileError(f"Failed to read file {path}: {e}")
+        
+    @staticmethod
+    def read_text(path: Path, encoding: str = 'utf-8') -> str:
+        """Read a file's content as text."""
+        try:
+            return path.read_text(encoding=encoding)
+        except Exception as e:
+            raise FileError(f"Failed to read file {path}: {e}")
+        
+    @staticmethod
+    def write_file(path: Path, content: str, encoding: str = 'utf-8') -> None:
+        """Write text content to a file."""
+        try:
+            ensure_directory(path.parent)
+            path.write_text(content, encoding=encoding)
+        except Exception as e:
+            raise FileError(f"Failed to write to file {path}: {e}")
+        
+    @staticmethod
+    def write_binary(path: Path, content: bytes) -> None:
+        """Write binary content to a file."""
+        try:
+            ensure_directory(path.parent)
+            path.write_bytes(content)
+        except Exception as e:
+            raise FileError(f"Failed to write binary data to {path}: {e}")
+            
+    @staticmethod
+    def copy_file(src: Path, dst: Path) -> None:
+        """Copy a file from source to destination."""
+        try:
+            ensure_directory(dst.parent)
+            shutil.copy2(src, dst)
+        except Exception as e:
+            raise FileError(f"Failed to copy file from {src} to {dst}: {e}")