File size: 4,249 Bytes
abcd0c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
"""
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 .converter import ConversionError, ConversionResult, DocumentConverter, SUPPORTED_EXTENSIONS
from logger import get_logger

logger = get_logger(__name__)


@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 = min(8, (os.cpu_count() or 1) + 4),
    ) -> 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),
        )
        return BatchReport(
            total=total,
            succeeded=len(results),
            failed=len(errors),
            results=results,
            errors=errors,
            total_chars=sum(r.char_count for r in results),
            total_words=sum(r.word_count for r in results),
            total_duration_ms=sum(r.duration_ms for r in results),
        )

    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
        ]