core/batch.py

"""
Batch processing utilities for the MarkItDown API.

Provides BatchProcessor for converting multiple files concurrently using
a shared DocumentConverter instance, and BatchReport for aggregating results.

These classes are used internally by the CLI and server batch endpoints.
"""

from __future__ import annotations

import concurrent.futures
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional, Sequence

from logger import get_logger

from .converter import (
    ConversionError,
    ConversionResult,
    DocumentConverter,
    SUPPORTED_EXTENSIONS,
)

logger = get_logger(__name__)

_DEFAULT_MAX_WORKERS = min(8, (os.cpu_count() or 1) + 4)


@dataclass
class BatchReport:
    """Aggregated results from a batch file conversion run."""

    total: int
    succeeded: int
    failed: int
    results: list[ConversionResult]
    errors: list[ConversionError]
    total_chars: int
    total_words: int
    total_duration_ms: float

    @property
    def success_rate(self) -> float:
        """Percentage of files converted successfully."""
        return (self.succeeded / self.total * 100) if self.total else 0.0


class BatchProcessor:
    """Convert multiple files concurrently using a thread pool.

    Parameters
    ----------
    converter:
        Shared DocumentConverter instance.
    max_workers:
        Number of threads in the pool. Defaults to min(8, cpu_count + 4).
    """

    def __init__(
        self,
        converter: DocumentConverter,
        max_workers: int = _DEFAULT_MAX_WORKERS,
    ) -> None:
        self._converter = converter
        self._max_workers = max_workers

    def process_files(
        self,
        paths: Sequence[str | Path],
        progress_callback: Optional[Callable[[int, int, str], None]] = None,
    ) -> BatchReport:
        """Convert all files in *paths* and return a BatchReport.

        Parameters
        ----------
        paths:
            Iterable of file paths to convert.
        progress_callback:
            Optional callable invoked after each file completes.
            Receives ``(completed_count, total_count, source_path)``.
        """
        results: list[ConversionResult] = []
        errors: list[ConversionError] = []
        total = len(paths)

        with concurrent.futures.ThreadPoolExecutor(max_workers=self._max_workers) as executor:
            future_to_path = {
                executor.submit(self._converter.convert_file, p): p for p in paths
            }
            completed = 0
            for future in concurrent.futures.as_completed(future_to_path):
                completed += 1
                outcome = future.result()
                source = str(future_to_path[future])
                if isinstance(outcome, ConversionResult):
                    results.append(outcome)
                else:
                    errors.append(outcome)
                if progress_callback:
                    progress_callback(completed, total, source)

        logger.info(
            "batch_processor | done | total=%d | succeeded=%d | failed=%d",
            total, len(results), len(errors),
        )
        total_chars = sum(r.char_count for r in results)
        total_words = sum(r.word_count for r in results)
        total_duration = sum(r.duration_ms for r in results)
        return BatchReport(
            total=total,
            succeeded=len(results),
            failed=len(errors),
            results=results,
            errors=errors,
            total_chars=total_chars,
            total_words=total_words,
            total_duration_ms=total_duration,
        )

    def discover_files(
        self,
        directory: str | Path,
        recursive: bool = True,
        extensions: Optional[set[str]] = None,
    ) -> list[Path]:
        """Return all convertible files under *directory*.

        Parameters
        ----------
        directory:
            Root directory to scan.
        recursive:
            When True, scan subdirectories as well.
        extensions:
            Set of extensions to include. Defaults to SUPPORTED_EXTENSIONS.
        """
        root = Path(directory).resolve()
        exts = extensions or SUPPORTED_EXTENSIONS
        glob_pattern = "**/*" if recursive else "*"
        return [
            p for p in root.glob(glob_pattern)
            if p.is_file() and p.suffix.lower() in exts
        ]