File size: 4,424 Bytes
abcd0c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c5638b0
 
 
 
 
 
 
abcd0c2
 
c5638b0
 
abcd0c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c5638b0
abcd0c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c5638b0
 
 
abcd0c2
 
 
 
 
 
c5638b0
 
 
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
137
138
139
140
141
142
143
144
145
146
147
"""
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
        ]