diff --git a/README.rst b/README.rst index 90a0ba3..0ee03d9 100644 --- a/README.rst +++ b/README.rst @@ -119,6 +119,33 @@ Multiple bars if __name__ == '__main__': main() +Parallel execution +============================================================================== + +Run a function over a batch of items -- threads, processes, or asyncio -- +with a live progress bar, in one call: + +.. code:: python + + import progressbar + + results = progressbar.map(fetch, urls, workers=8) # threads + results = progressbar.map(crunch, files, pool='process') # processes + results = await progressbar.amap(fetch, urls) # asyncio + + # A progress-bar'd xargs -P: + progressbar.run('gzip -k {}', files, workers=4) + + # An overall bar plus one bar per in-flight task: + progressbar.map(crunch, files, workers=4, bar='multi') + +Results come back in input order; ``imap``/``imap_unordered`` stream +them instead, ``gather`` is a drop-in ``asyncio.gather`` with a bar, +and the bar keeps animating even while long tasks are running. See the +`parallel execution guide +`_ +for errors, timeouts, pools, and the decorator form. + Unknown length and animated bars ============================================================================== diff --git a/docs/howto/index.rst b/docs/howto/index.rst index 4366ccb..a77fbad 100644 --- a/docs/howto/index.rst +++ b/docs/howto/index.rst @@ -21,3 +21,4 @@ already read the :doc:`tutorial <../tutorial/index>`. unknown-length multibar multibar-line-offset + parallel-execution diff --git a/docs/howto/parallel-execution.rst b/docs/howto/parallel-execution.rst new file mode 100644 index 0000000..558f17b --- /dev/null +++ b/docs/howto/parallel-execution.rst @@ -0,0 +1,236 @@ +====================================== +Run a batch in parallel with progress +====================================== + +One call runs a function over a batch of items -- on threads, +processes, or asyncio -- and renders a live progress bar while it +happens: + +.. code-block:: python + + import progressbar + + results = progressbar.map(fetch, urls, workers=8) + +``progressbar.map`` mirrors the builtin ``map``: results come back in +input order, multiple iterables zip (``progressbar.map(pow, bases, +exps)``), and the bar counts completed items. The bar keeps animating +-- ETA, timers, spinners -- even while long tasks are running with +nothing finishing. + +Choosing where the work runs +============================ + +.. code-block:: python + + # Threads (default): I/O-bound work, no pickling requirements. + progressbar.map(fetch, urls, workers=8) + + # Processes: CPU-bound work; fn and items must be picklable and + # fn must be a module-level function. + progressbar.map(crunch, files, pool='process') + + # tqdm-style spellings of the same two calls: + progressbar.thread_map(fetch, urls) + progressbar.process_map(crunch, files) + + # An executor you already have (never shut down for you): + progressbar.map(fetch, urls, pool=my_executor) + +Process pools accept the executor construction keywords you would pass +to ``ProcessPoolExecutor`` -- ``initializer``/``initargs`` for +per-worker setup (database connections, loaded models), +``mp_context``, and ``max_tasks_per_child`` (Python 3.11+). Thread +pools accept ``thread_name_prefix``. On Python 3.14+, +``pool='interpreter'`` runs on an ``InterpreterPoolExecutor``. + +For many small items on a process pool, items are automatically +submitted in chunks to amortize the per-task overhead (the bar then +advances a chunk at a time); pass ``chunksize=`` to tune it. + +Streaming results as they arrive +================================ + +.. code-block:: python + + # Input order, lazily -- multiprocessing.Pool.imap semantics: + for result in progressbar.imap(fetch, urls, workers=8): + handle(result) + + # Completion order, as (item, result) pairs -- the pair restores + # the correspondence that completion order loses: + for url, result in progressbar.imap_unordered(fetch, urls): + print(f'{url} done') + +Breaking out of either loop cancels the not-yet-submitted work and +shuts the run down. For deterministic cleanup wrap the iterator in +``contextlib.closing`` (``contextlib.aclosing`` for the async +variants). + +Async code +========== + +.. code-block:: python + + # fn may be async -- or plain sync, which runs via asyncio.to_thread: + results = await progressbar.amap(fetch, urls, concurrency=8) + + # Lazy variants, mirroring the sync pair: + async for result in progressbar.aimap(fetch, urls, concurrency=8): + ... + async for url, result in progressbar.aimap_unordered(fetch, urls): + ... + + # A drop-in asyncio.gather with a bar (same signature, same + # ordering, same return_exceptions keyword): + results = await progressbar.gather(*coroutines) + +``concurrency=None`` (the default) creates every task up front, exactly +like ``asyncio.gather``; pass a limit for large batches so tasks are +created lazily in a window of that size. + +Shell commands: a progress-bar'd ``xargs -P`` +============================================= + +.. code-block:: python + + # {} (or {item}) is replaced per item; no placeholder appends the + # item as the last argument, like xargs: + procs = progressbar.run('gzip -k {}', files, workers=4) + + # List and callable templates for full control: + progressbar.run(['ffmpeg', '-i', '{}', '{}.mp4'], videos) + progressbar.run(lambda p: ['convert', p, p.with_suffix('.png')], images) + +Substitution replaces only the exact placeholder tokens -- never +``str.format`` -- so commands containing literal braces (``awk '{print +$1}'``) work, and an item containing spaces stays a single argument. +Each command's output is captured into its +:py:class:`subprocess.CompletedProcess` (so child output can't corrupt +the bar), and a non-zero exit raises +:py:class:`subprocess.CalledProcessError` by default (``check=False`` +to collect exit codes instead). ``shell=True`` is available for the +string form but substitutes items into a shell command line -- only +use it with items you trust. + +Sub-task bars with ``bar='multi'`` +================================== + +.. code-block:: python + + progressbar.map(crunch, files, workers=4, bar='multi') + +Instead of one aggregate bar, a ``MultiBar`` shows an overall bar plus +one line per in-flight task, labeled with the item. Inside a worker +(threads and asyncio), ``progressbar.current_task_bar()`` returns that +task's own bar for sub-progress: + +.. code-block:: python + + def crunch(path): + bar = progressbar.current_task_bar() + for i, block in enumerate(read_blocks(path)): + process(block) + if bar is not None: + bar.update(i) + +``'multi'`` suits modest worker counts -- the display occupies one +terminal row per in-flight task plus one for the total. Process +workers can't reach their bar this way yet +(``current_task_bar()`` returns ``None`` there): the bar object cannot +cross the process boundary in this release. + +Errors, timeouts and Ctrl-C +=========================== + +.. code-block:: python + + # Default: fail fast. First exception cancels the pending work and + # re-raises; running tasks finish first. + progressbar.map(crunch, files) + + # Collect instead: exceptions come back in place of results. + results = progressbar.map(crunch, files, on_error='return') + for item, result in zip(files, results): + if isinstance(result, Exception): + print(f'{item} failed: {result}') + + # An overall deadline; expiry cancels pending work and raises + # TimeoutError without waiting for stragglers. + progressbar.map(crunch, files, timeout=60) + +On ``KeyboardInterrupt`` the pending work is cancelled, the bar's +final state is left on its own line, and the interrupt re-raises -- +no garbled terminal. One honest limitation applies everywhere: +*running* tasks cannot be killed. A cancelled process task runs to +completion in the background, and a sync function inside ``amap`` is +abandoned in its thread, not interrupted. + +Reusing a pool across batches +============================= + +.. code-block:: python + + with progressbar.Pool(8) as pool: # Pool(8, 'process') for processes + thumbnails = pool.map(thumbnail, images) + uploads = pool.map(upload, thumbnails) + pool.run('touch {}', markers) + + async with progressbar.AsyncPool(8) as pool: + results = await pool.map(fetch, urls) + +Constructor keywords (``bar=``, ``on_error=``, ...) become per-call +defaults; the executor is created lazily on first use and +``pool.executor`` exposes it for direct ``submit()`` calls. + +The decorator spelling +====================== + +.. code-block:: python + + @progressbar.parallel(workers=4, pool='process') + def crunch(path): ... + + crunch(one_path) # still an ordinary call + crunch.map(paths) # parallel + bar, decorator config applied + await crunch.amap(paths) + +The decorator needs a plain module-level function (that's what +``pool='process'`` can pickle); it returns the same function object +with ``.map``/``.imap``/``.imap_unordered``/``.starmap``/``.amap``/ +``.aimap``/``.aimap_unordered`` attached. + +Progress for futures you already have +===================================== + +.. code-block:: python + + futures = [executor.submit(work, item) for item in items] + for future in progressbar.as_completed(futures): + handle(future.result()) + +A superset of :py:func:`concurrent.futures.as_completed` -- same +yield order and ``timeout`` behavior, plus the bar. + +Semantics worth knowing +======================= + +- **Bar options pass through.** Every verb forwards unknown keywords + to the bar: ``prefix=``/``desc=``, ``widgets=``, ``max_value=``, ... + A typo raises ``TypeError`` instead of being silently ignored. +- **Bar modes.** ``bar='plain'`` (default), ``bar='multi'``, + ``bar=False`` (run silently), or pass your own configured + ``ProgressBar``/``MultiBar`` instance to be driven. +- **One poll knob.** ``poll_interval`` (default 0.1s) is both how + often the coordinator wakes and how often the bar redraws with no + progress -- the keep-alive cadence. +- **Bounded memory.** Items are pulled from the input lazily and at + most ``buffersize`` tasks (default ``4 × workers``) are unfinished + at once, so million-item and generator inputs run in flat memory. + This deviates deliberately from ``Executor.map(buffersize=None)`` + on Python 3.14, where ``None`` means unbounded. +- **imap_unordered yields pairs.** ``multiprocessing.Pool``'s version + yields bare results; ours yields ``(item, result)`` because + completion order loses the correspondence. +- **Totals.** ``len()`` where available, ``operator.length_hint`` + otherwise; without either the bar runs in unknown-length mode. diff --git a/docs/reference/index.rst b/docs/reference/index.rst index 73b8d6e..ee9f106 100644 --- a/docs/reference/index.rst +++ b/docs/reference/index.rst @@ -9,5 +9,6 @@ Complete descriptions of the public API surface. progressbar multibar + parallel cli Full module autodoc <../progressbar> diff --git a/docs/reference/parallel.rst b/docs/reference/parallel.rst new file mode 100644 index 0000000..acda2b0 --- /dev/null +++ b/docs/reference/parallel.rst @@ -0,0 +1,118 @@ +================== +Parallel execution +================== + +The parallel verb family runs a callable over a batch of items -- on +threads, processes, or asyncio -- with a progress bar. The how-to +guide (:doc:`../howto/parallel-execution`) shows the idioms; this page +is the API reference. + +The shared keywords +=================== + +Every verb accepts (where applicable): + +.. list-table:: + :header-rows: 1 + :widths: 22 78 + + * - Keyword + - What it does + * - ``workers`` + - Pool size for the sync verbs. Default: the executor's own + default (``min(32, cpus + 4)`` threads, ``cpus`` processes). + Accepted as an alias for ``concurrency`` on the async verbs. + * - ``concurrency`` + - Async verbs: maximum in-flight tasks. ``None`` (default) + creates every task up front, like ``asyncio.gather``. + * - ``pool`` + - ``'thread'`` (default), ``'process'``, ``'interpreter'`` + (Python 3.14+), or an existing + :py:class:`concurrent.futures.Executor` instance (used as-is, + never shut down for you). + * - ``bar`` + - ``'plain'`` (one aggregate bar, default), ``'multi'`` (a + :py:class:`~progressbar.multi.MultiBar` with one sub-bar per + in-flight task), ``False`` (no output), or a configured + ``ProgressBar``/``MultiBar`` instance to drive. + * - ``on_error`` + - ``'raise'`` (default): first failure cancels pending work and + re-raises. ``'return'``: exceptions appear in place of their + results; ``KeyboardInterrupt``/``SystemExit`` still propagate. + * - ``chunksize`` + - Items per task. Default: 1 on threads; automatic on process + and interpreter pools (about 16 chunks per worker, capped at + 1000). The bar advances per chunk. + * - ``buffersize`` + - Maximum unfinished submitted tasks (sync verbs). Default + ``max(4 × workers, 16)``; keeps memory flat on huge or lazy + inputs. + * - ``timeout`` + - Overall deadline in seconds. Expiry cancels pending work and + raises ``TimeoutError`` without waiting for running tasks. + * - ``poll_interval`` + - Seconds between coordinator wakeups *and* no-progress bar + redraws (default 0.1) -- one knob for both. + * - ``initializer``, ``initargs``, ``mp_context``, + ``max_tasks_per_child``, ``thread_name_prefix`` + - Forwarded verbatim to the executor constructor + (``max_tasks_per_child`` needs Python 3.11+; each option is + validated against the pool kind). + * - ``**bar_kwargs`` + - Anything else goes to the bar: ``prefix=``/``desc=``, + ``suffix=``, ``widgets=``, ``max_value=``, ... Unknown names + raise ``TypeError``. + +Sync verbs +========== + +.. autofunction:: progressbar.map + :no-index: +.. autofunction:: progressbar.imap + :no-index: +.. autofunction:: progressbar.imap_unordered + :no-index: +.. autofunction:: progressbar.starmap + :no-index: +.. autofunction:: progressbar.thread_map + :no-index: +.. autofunction:: progressbar.process_map + :no-index: +.. autofunction:: progressbar.as_completed + :no-index: +.. autofunction:: progressbar.run + :no-index: + +Async verbs +=========== + +.. autofunction:: progressbar.amap + :no-index: +.. autofunction:: progressbar.aimap + :no-index: +.. autofunction:: progressbar.aimap_unordered + :no-index: +.. autofunction:: progressbar.gather + :no-index: + +Reusable layers +=============== + +.. autoclass:: progressbar.Pool + :members: + :member-order: bysource + :no-index: + +.. autoclass:: progressbar.AsyncPool + :members: + :member-order: bysource + :no-index: + +.. autofunction:: progressbar.parallel + :no-index: +.. autofunction:: progressbar.current_task_bar + :no-index: + +.. autoclass:: progressbar.ParallelFunction + :members: + :no-index: diff --git a/progressbar/__init__.py b/progressbar/__init__.py index 8111961..8958314 100644 --- a/progressbar/__init__.py +++ b/progressbar/__init__.py @@ -21,6 +21,28 @@ if typing.TYPE_CHECKING: # Eager imports for type checkers only, loaded lazily at runtime by # __getattr__ below. Names appear in __all__ so they read as re-exports. + # The redundant `name as name` aliases mark the namespace-only + # exports (kept out of ``__all__``, see ``_NAMESPACE_ONLY``) as + # deliberate re-exports for type checkers and linters. + from ._parallel import ( + AsyncPool, + ParallelFunction, + Pool, + aimap as aimap, + aimap_unordered as aimap_unordered, + amap as amap, + as_completed as as_completed, + current_task_bar, + gather as gather, + imap as imap, + imap_unordered as imap_unordered, + map as map, # noqa: A004 - namespaced use only, not in __all__ + parallel, + process_map, + run as run, + starmap, + thread_map, + ) from .algorithms import ( DoubleExponentialMovingAverage, ExponentialMovingAverage, @@ -94,6 +116,14 @@ 'UnknownLength': 'base', 'MultiBar': 'multi', 'SortKey': 'multi', + 'AsyncPool': '_parallel', + 'ParallelFunction': '_parallel', + 'Pool': '_parallel', + 'current_task_bar': '_parallel', + 'parallel': '_parallel', + 'process_map': '_parallel', + 'starmap': '_parallel', + 'thread_map': '_parallel', 'progressbar': 'shortcuts', 'LineOffsetStreamWrapper': 'terminal.stream', 'len_color': 'utils', @@ -130,6 +160,22 @@ 'VariableMixin': 'widgets', } +#: Exported name -> submodule, for names resolvable as +#: ``progressbar.`` (and by explicit import) but deliberately +#: kept out of ``__all__``: they would shadow builtins or stdlib names +#: under ``from progressbar import *`` (``map``, ``as_completed``, ...). +_NAMESPACE_ONLY: dict[str, str] = { + 'aimap': '_parallel', + 'aimap_unordered': '_parallel', + 'amap': '_parallel', + 'as_completed': '_parallel', + 'gather': '_parallel', + 'imap': '_parallel', + 'imap_unordered': '_parallel', + 'map': '_parallel', + 'run': '_parallel', +} + def __getattr__(name: str) -> typing.Any: """Lazily import submodules and exported names on first access.""" @@ -138,7 +184,7 @@ def __getattr__(name: str) -> typing.Any: globals()[name] = module # cache so __getattr__ runs only once return module - module_name = _NAME_TO_MODULE.get(name) + module_name = _NAME_TO_MODULE.get(name) or _NAMESPACE_ONLY.get(name) if module_name is None: raise AttributeError(f'module {__name__!r} has no attribute {name!r}') value = getattr(importlib.import_module(f'.{module_name}', __name__), name) @@ -161,6 +207,7 @@ def __dir__() -> list[str]: 'AdaptiveETA', 'AdaptiveTransferSpeed', 'AnimatedMarker', + 'AsyncPool', 'Bar', 'BouncingBar', 'Counter', @@ -182,8 +229,10 @@ def __dir__() -> list[str]: 'MultiProgressBar', 'MultiRangeBar', 'NullBar', + 'ParallelFunction', 'Percentage', 'PercentageLabelBar', + 'Pool', 'Postfix', 'ProgressBar', 'ReverseBar', @@ -199,7 +248,12 @@ def __dir__() -> list[str]: 'VariableMixin', '__author__', '__version__', + 'current_task_bar', 'len_color', + 'parallel', + 'process_map', 'progressbar', + 'starmap', 'streams', + 'thread_map', ] diff --git a/progressbar/_parallel/__init__.py b/progressbar/_parallel/__init__.py new file mode 100644 index 0000000..297cd72 --- /dev/null +++ b/progressbar/_parallel/__init__.py @@ -0,0 +1,57 @@ +"""Parallel execution with progress bars. + +Private implementation package behind the public ``progressbar.map`` / +``imap`` / ``amap`` / ``run`` / ``Pool`` family. The package is private +(underscored) because the public surface includes a ``parallel`` +*decorator* exported as ``progressbar.parallel`` -- a public +``progressbar/parallel.py`` module would clobber that attribute on +``import progressbar.parallel``. + +Public names are re-exported lazily from ``progressbar/__init__.py``; +star-import-unsafe ones (``map``, ``imap``, ``gather``, ...) resolve +through the namespace only and stay out of ``progressbar.__all__``. +""" + +from ._async import ( + AsyncPool, + aimap, + aimap_unordered, + amap, + gather, +) +from ._common import current_task_bar +from ._decorator import ( + ParallelFunction, + parallel, +) +from ._shell import run +from ._sync import ( + Pool, + as_completed, + imap, + imap_unordered, + map, # noqa: A004 - intentional builtin name, namespaced use only + process_map, + starmap, + thread_map, +) + +__all__ = [ + 'AsyncPool', + 'ParallelFunction', + 'Pool', + 'aimap', + 'aimap_unordered', + 'amap', + 'as_completed', + 'current_task_bar', + 'gather', + 'imap', + 'imap_unordered', + 'map', + 'parallel', + 'process_map', + 'run', + 'starmap', + 'thread_map', +] diff --git a/progressbar/_parallel/_async.py b/progressbar/_parallel/_async.py new file mode 100644 index 0000000..126ddd7 --- /dev/null +++ b/progressbar/_parallel/_async.py @@ -0,0 +1,452 @@ +"""The asyncio engine behind `amap`, `aimap` and `gather`. + +Mirrors the sync engine's coordination pattern -- windowed task +creation, a done-queue costing O(1) per completion, poll-timeout ticks +for the keep-alive guarantee -- with asyncio primitives. Sync callables +are welcome too: they run via `asyncio.to_thread`, so one async entry +point covers both worlds. +""" + +from __future__ import annotations + +import asyncio +import functools +import inspect +import time +import typing + +from . import ( + _common, + _display, +) + +#: One completion event: (item index, argument tuple, ok, result/error). +Completion = tuple[int, _common.ItemArgs, bool, typing.Any] + +#: Default seconds between coordinator wakeups; doubles as the bar's +#: redraw interval (one knob -- see the keep-alive contract). +DEFAULT_POLL_INTERVAL: float = 0.1 + + +def _call_strategy(fn: typing.Callable[..., typing.Any]) -> str: + """Classify `fn` as ``'async'`` or ``'sync'``. + + Unwraps `functools.partial` manually: on the 3.10 floor + `inspect.iscoroutinefunction` does not look through partials, and + `asyncio.iscoroutinefunction` (which does) is deprecated in 3.14. + """ + target: typing.Any = fn + while isinstance(target, functools.partial): + target = target.func + return 'async' if inspect.iscoroutinefunction(target) else 'sync' + + +async def _acall( + fn: typing.Callable[..., typing.Any], + args: _common.ItemArgs, + strategy: str, +) -> typing.Any: + """Await `fn(*args)` per the detected strategy. + + ``'async'`` awaits directly. ``'sync'`` runs in a thread via + `asyncio.to_thread` -- note cancellation *abandons* such a thread + rather than interrupting it -- and, if the call returned an + awaitable (a sync factory of coroutines), awaits that too. + """ + if strategy == 'async': + return await fn(*args) + value: typing.Any = await asyncio.to_thread(fn, *args) + if inspect.isawaitable(value): + value = await value + return value + + +async def _await_it(awaitable: typing.Awaitable[typing.Any]) -> typing.Any: + """Adapt a bare awaitable (the `gather` path) into a task coro.""" + return await awaitable + + +class _AsyncRun: + """State and coordination for one `execute_async` invocation.""" + + fn: typing.Callable[..., typing.Any] | None + strategy: str + awaitables: bool + total: typing.Any + on_error: str + single: bool + window: int | None + timeout: float | None + poll_interval: float + deadline: float | None + display: _display.Display + done: asyncio.Queue[asyncio.Task[typing.Any]] + in_flight: dict[ + asyncio.Task[typing.Any], tuple[int, _common.ItemArgs, int] + ] + item_source: typing.Iterator[tuple[int, _common.ItemArgs]] + seq: int + + def __init__( + self, + fn: typing.Callable[..., typing.Any] | None, + iterables: tuple[typing.Iterable[typing.Any], ...], + *, + concurrency: int | None, + bar: typing.Any, + on_error: str, + timeout: float | None, + poll_interval: float, + awaitables: bool, + bar_kwargs: dict[str, typing.Any], + ) -> None: + """Validate the configuration and set up the display.""" + if on_error not in ('raise', 'return'): + raise ValueError( + f"on_error={on_error!r} is not valid: expected 'raise' " + f"or 'return'" + ) + _common.validate_bar_kwargs(bar_kwargs) + self.fn = fn + self.strategy = '' if fn is None else _call_strategy(fn) + self.awaitables = awaitables + self.total = _common.detect_total(iterables) + self.on_error = on_error + self.single = len(iterables) == 1 + self.window = concurrency + self.timeout = timeout + self.poll_interval = poll_interval + self.deadline = None if timeout is None else time.monotonic() + timeout + self.display = _display.make_display( + bar, + total=self.total, + poll_interval=poll_interval, + bar_kwargs=bar_kwargs, + ) + self.done = asyncio.Queue() + self.in_flight = {} + self.item_source = enumerate(zip(*iterables, strict=False)) + self.seq = 0 + + async def completions(self) -> typing.AsyncIterator[Completion]: + """Drive the run, yielding per-item events in completion order.""" + self.display.start(self.total) + if self.window is None: + # gather semantics: everything in flight at once. + while self._launch_one(): + pass + else: + while len(self.in_flight) < self.window and self._launch_one(): + pass + while self.in_flight: + self._check_deadline() + task = await self._next_done() + if task is not None: + yield self._handle(task) + + def _launch_one(self) -> bool: + """Create the next task; `False` when the input is exhausted.""" + indexed: tuple[int, _common.ItemArgs] | None = next( + self.item_source, None + ) + if indexed is None: + return False + index, args = indexed + self.seq += 1 + label: str = str(_common.item_of(args, self.single)) + task_bar = self.display.task_started(self.seq, label) + coroutine: typing.Coroutine[typing.Any, typing.Any, typing.Any] + if self.awaitables: + coroutine = _await_it(args[0]) + else: + assert self.fn is not None + coroutine = _acall(self.fn, args, self.strategy) + if task_bar is None: + task: asyncio.Task[typing.Any] = asyncio.ensure_future(coroutine) + else: + # Task creation snapshots the current context, so binding + # the contextvar around it is what makes + # `current_task_bar()` work inside the task. + token = _common._task_bar_var.set(task_bar) # noqa: SLF001 + try: + task = asyncio.ensure_future(coroutine) + finally: + _common._task_bar_var.reset(token) # noqa: SLF001 + self.in_flight[task] = (index, args, self.seq) + task.add_done_callback(self.done.put_nowait) + return True + + async def _next_done(self) -> asyncio.Task[typing.Any] | None: + """Wait one poll for a completion; tick the display on none.""" + try: + return await asyncio.wait_for( + self.done.get(), timeout=self.poll_interval + ) + except asyncio.TimeoutError: + self.display.tick() + return None + + def _check_deadline(self) -> None: + """Raise once the overall `timeout` budget is spent.""" + if self.deadline is not None and time.monotonic() > self.deadline: + raise asyncio.TimeoutError( + f'parallel execution exceeded timeout={self.timeout}' + ) + + def _handle(self, task: asyncio.Task[typing.Any]) -> Completion: + """Turn one finished task into a completion event.""" + index, args, seq = self.in_flight.pop(task) + if task.cancelled(): + # Something outside this run cancelled the task; surface it + # rather than silently dropping the item. + self.display.task_finished(seq, ok=False) + raise asyncio.CancelledError + error: BaseException | None = task.exception() + if error is not None: + self.display.task_finished(seq, ok=False) + if self.on_error == 'raise' or isinstance( + error, (KeyboardInterrupt, SystemExit) + ): + raise error + self.display.advance() + self._launch_one() + return index, args, False, error + self.display.task_finished(seq, ok=True) + self.display.advance() + self._launch_one() + return index, args, True, task.result() + + async def close(self, *, success: bool) -> None: + """Cancel outstanding tasks, await them, release the display.""" + for task in self.in_flight: + task.cancel() + if self.in_flight: + # Awaiting the cancelled tasks prevents "Task exception was + # never retrieved"/"Task was destroyed" noise on teardown. + await asyncio.gather(*self.in_flight, return_exceptions=True) + self.display.finish(success=success) + + +async def execute_async( + fn: typing.Callable[..., typing.Any] | None, + iterables: tuple[typing.Iterable[typing.Any], ...], + *, + concurrency: int | None = None, + workers: int | None = None, + bar: typing.Any = 'plain', + on_error: str = 'raise', + timeout: float | None = None, + poll_interval: float = DEFAULT_POLL_INTERVAL, + awaitables: bool = False, + **bar_kwargs: typing.Any, +) -> typing.AsyncIterator[Completion]: + """Run `fn` over zipped `iterables` on the event loop. + + The async twin of the sync `execute`: yields ``(index, args, ok, + value)`` events in completion order. `workers` is accepted as an + alias for `concurrency` (same concept, sync spelling). + ``concurrency=None`` creates every task up front (`asyncio.gather` + semantics -- pass a limit for large batches); with a limit, tasks + are created lazily in a window of that size. + + With ``awaitables=True`` (the `gather` path) the single iterable + contains awaitables to schedule directly and `fn` is ignored. + + Raises: + ValueError: Invalid `on_error`. + TypeError: Unknown bar keyword. + asyncio.TimeoutError: The overall `timeout` expired; outstanding + tasks are cancelled and awaited first. + """ + if concurrency is None: + concurrency = workers + run: _AsyncRun = _AsyncRun( + fn, + iterables, + concurrency=concurrency, + bar=bar, + on_error=on_error, + timeout=timeout, + poll_interval=poll_interval, + awaitables=awaitables, + bar_kwargs=bar_kwargs, + ) + success: bool = False + try: + async for event in run.completions(): + yield event + success = True + finally: + await run.close(success=success) + + +async def aimap( + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> typing.AsyncIterator[typing.Any]: + """Lazily apply `fn` on the event loop, yielding in input order. + + The async counterpart of `imap`: results-only, ordered, with + out-of-order completions held back until their turn. Use + `contextlib.aclosing` for deterministic cleanup on early exit. + See `execute_async` for keywords. + """ + held: dict[int, typing.Any] = {} + next_index: int = 0 + async for index, _args, _ok, value in execute_async( + fn, iterables, **kwargs + ): + held[index] = value + while next_index in held: + yield held.pop(next_index) + next_index += 1 + + +async def aimap_unordered( + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> typing.AsyncIterator[tuple[typing.Any, typing.Any]]: + """Lazily apply `fn` on the event loop, yielding as tasks finish. + + The async counterpart of `imap_unordered`: ``(item, result)`` pairs + in completion order (the pair shape restores the correspondence + completion order loses). See `execute_async` for keywords. + """ + single: bool = len(iterables) == 1 + async for _index, args, _ok, value in execute_async( + fn, iterables, **kwargs + ): + yield _common.item_of(args, single), value + + +async def gather( + *awaitables: typing.Awaitable[typing.Any], + return_exceptions: bool = False, + bar: typing.Any = 'plain', + poll_interval: float = DEFAULT_POLL_INTERVAL, + timeout: float | None = None, + **bar_kwargs: typing.Any, +) -> list[typing.Any]: + """`asyncio.gather` with a progress bar. + + A drop-in replacement: results in argument order, no arguments + yields ``[]``, and `return_exceptions` keeps asyncio's exact + keyword (mapped to ``on_error='return'`` internally). Unlike + `amap` there is no concurrency limiting -- the awaitables already + exist, matching `asyncio.gather` semantics. + """ + if not awaitables: + return [] + results: dict[int, typing.Any] = { + index: value + async for index, _args, _ok, value in execute_async( + None, + (awaitables,), + on_error='return' if return_exceptions else 'raise', + bar=bar, + poll_interval=poll_interval, + timeout=timeout, + awaitables=True, + **bar_kwargs, + ) + } + return [results[index] for index in range(len(results))] + + +class AsyncPool: + """Shared concurrency limit plus per-call defaults for async verbs. + + The async sibling of `Pool`. There is no executor to manage -- + tasks run on the caller's event loop -- so this is configuration + reuse: a concurrency bound and default keywords applied to every + call, overridable per call:: + + async with progressbar.AsyncPool(8) as pool: + first = await pool.map(fetch, urls) + async for item, result in pool.imap_unordered(fetch, more): + ... + """ + + _concurrency: int | None + _defaults: dict[str, typing.Any] + + def __init__( + self, concurrency: int | None = None, **defaults: typing.Any + ) -> None: + """Store the concurrency bound and per-call defaults.""" + self._concurrency = concurrency + self._defaults = defaults + + def _merged(self, kwargs: dict[str, typing.Any]) -> dict[str, typing.Any]: + """Per-call keywords override the pool's defaults.""" + return { + 'concurrency': self._concurrency, + **self._defaults, + **kwargs, + } + + def map( # noqa: A003 - mirrors the module verb + self, + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, + ) -> typing.Coroutine[typing.Any, typing.Any, list[typing.Any]]: + """`amap` with this pool's limit and defaults; awaitable.""" + return amap(fn, *iterables, **self._merged(kwargs)) + + def imap( + self, + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, + ) -> typing.AsyncIterator[typing.Any]: + """`aimap` with this pool's limit and defaults.""" + return aimap(fn, *iterables, **self._merged(kwargs)) + + def imap_unordered( + self, + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, + ) -> typing.AsyncIterator[tuple[typing.Any, typing.Any]]: + """`aimap_unordered` with this pool's limit and defaults.""" + return aimap_unordered(fn, *iterables, **self._merged(kwargs)) + + async def __aenter__(self) -> AsyncPool: + """Return the pool (no resource to acquire; symmetry with Pool).""" + return self + + async def __aexit__(self, *exc_info: typing.Any) -> None: + """Nothing to release; tasks belong to the caller's loop.""" + + +async def amap( + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> list[typing.Any]: + """Apply `fn` to every zipped item on the event loop; ordered. + + The async counterpart of `progressbar.map`. `fn` may be an async + *or* a plain sync callable -- sync callables run in a thread via + `asyncio.to_thread`. Results come back in input order:: + + results = await progressbar.amap(fetch, urls, concurrency=8) + + See `execute_async` for the keyword reference. + """ + results: dict[int, typing.Any] = { + index: value + async for index, _args, _ok, value in execute_async( + fn, iterables, **kwargs + ) + } + return [results[index] for index in range(len(results))] diff --git a/progressbar/_parallel/_common.py b/progressbar/_parallel/_common.py new file mode 100644 index 0000000..2a1cb47 --- /dev/null +++ b/progressbar/_parallel/_common.py @@ -0,0 +1,212 @@ +"""Shared plumbing for the parallel execution verbs. + +Everything here is engine-agnostic: argument validation, total +detection, chunking, worker/window defaults, and the context variable +that gives workers access to their own sub-bar under ``bar='multi'``. +""" + +from __future__ import annotations + +import contextvars +import functools +import inspect +import itertools +import operator +import os +import typing + +from .. import ( + bar as bar_module, + base, +) + +#: One zipped argument tuple, i.e. one call's positional arguments. +ItemArgs = tuple[typing.Any, ...] + +T = typing.TypeVar('T') + +#: Chunk sizing targets ~16 chunks per worker so completion events stay +#: frequent enough for a lively bar while amortizing per-task overhead. +_CHUNKS_PER_WORKER: int = 16 +#: Hard cap so gigantic inputs still produce regular progress updates. +_MAX_AUTO_CHUNKSIZE: int = 1_000 +#: Submission window per worker; the floor keeps tiny pools busy. +_WINDOWS_PER_WORKER: int = 4 +_MIN_BUFFERSIZE: int = 16 + +#: The bar owned by the currently executing task, set by `with_task_bar` +#: around each worker invocation under ``bar='multi'``. Workers read it +#: through `current_task_bar`. +_task_bar_var: contextvars.ContextVar[bar_module.ProgressBar | None] = ( + contextvars.ContextVar('current_task_bar', default=None) +) + + +def current_task_bar() -> bar_module.ProgressBar | None: + """Return the calling task's own progress bar, if it has one. + + Inside a function executed by `progressbar.map`/`amap` with + ``bar='multi'`` this returns the per-task bar so the worker can + report sub-progress (``current_task_bar().update(i)``). Anywhere + else -- including process-pool workers, which cannot share a bar + object with the parent in v1 -- it returns `None`. + """ + return _task_bar_var.get() + + +def with_task_bar( + task_bar: bar_module.ProgressBar, + inner: typing.Callable[[], T], +) -> typing.Callable[[], T]: + """Wrap `inner` so `current_task_bar` returns `task_bar` inside it.""" + + def _bound() -> T: + token: contextvars.Token[bar_module.ProgressBar | None] = ( + _task_bar_var.set(task_bar) + ) + try: + return inner() + finally: + _task_bar_var.reset(token) + + return _bound + + +def detect_total( + iterables: tuple[typing.Iterable[typing.Any], ...], +) -> int | typing.Any: + """Return the number of items `zip(*iterables)` will yield. + + Uses `len` where available, falling back to `operator.length_hint`; + any iterable without either makes the total `base.UnknownLength`. + Multiple iterables zip, so the total is their minimum. + """ + totals: list[int] = [] + for iterable in iterables: + total: int = _total_of(iterable) + if total < 0: + return base.UnknownLength + totals.append(total) + return min(totals) if totals else 0 + + +def _total_of(iterable: typing.Iterable[typing.Any]) -> int: + """Return `len`/`length_hint` for one iterable, -1 when unknown.""" + try: + return len(iterable) # type: ignore[arg-type] + except TypeError: + return operator.length_hint(iterable, -1) + + +@functools.cache +def known_bar_kwargs(cls: type) -> frozenset[str]: + """Collect every keyword parameter accepted along `cls`'s MRO.""" + names: set[str] = set() + for klass in cls.__mro__: + init: typing.Any = klass.__dict__.get('__init__') + if init is None: + continue + for parameter in inspect.signature(init).parameters.values(): + if parameter.kind in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ): + names.add(parameter.name) + names.discard('self') + return frozenset(names) + + +def validate_bar_kwargs(bar_kwargs: dict[str, typing.Any]) -> None: + """Reject unknown bar keyword arguments loudly. + + `ProgressBarMixinBase.__init__` swallows unknown ``**kwargs`` + silently, so a typo like ``worker=8`` would otherwise run with + defaults and no error -- exactly the failure this guard exists for. + """ + # FastProgressBar subclasses ProgressBar, so its MRO covers both. + allowed: frozenset[str] = known_bar_kwargs(bar_module.ProgressBar) + unknown: set[str] = set(bar_kwargs) - allowed + if unknown: + raise TypeError( + f'unknown progress bar argument(s): {sorted(unknown)!r}. ' + f'Not a bar option and not a parallel option either.' + ) + + +def resolve_workers(workers: int | None, kind: str) -> int: + """Return the effective pool size, mirroring the executor defaults.""" + if workers is not None: + return workers + cpu_count: int = os.cpu_count() or 1 + if kind == 'thread': + # ThreadPoolExecutor's documented default. + return min(32, cpu_count + 4) + return cpu_count + + +def default_buffersize(workers: int) -> int: + """Return the default submission window (unfinished futures).""" + return max(_WINDOWS_PER_WORKER * workers, _MIN_BUFFERSIZE) + + +def auto_chunksize(total: int | typing.Any, workers: int) -> int: + """Pick a chunk size for process pools from the batch size. + + Targets `_CHUNKS_PER_WORKER` chunks per worker, capped at + `_MAX_AUTO_CHUNKSIZE` so progress updates stay regular. Streaming + inputs (unknown total) get 1: correctness first, tuning explicit. + """ + if total is base.UnknownLength: + return 1 + return max( + 1, min(total // (workers * _CHUNKS_PER_WORKER), _MAX_AUTO_CHUNKSIZE) + ) + + +def iter_chunks( + iterables: tuple[typing.Iterable[typing.Any], ...], + chunksize: int, +) -> typing.Iterator[list[ItemArgs]]: + """Lazily zip `iterables` and batch the argument tuples.""" + zipped: typing.Iterator[ItemArgs] = zip(*iterables, strict=False) + while chunk := list(itertools.islice(zipped, chunksize)): + yield chunk + + +def item_of(args: ItemArgs, single: bool) -> typing.Any: + """Return the user-facing item: bare for one iterable, tuple else.""" + return args[0] if single else args + + +def run_chunk( + fn: typing.Callable[..., typing.Any], + chunk: list[ItemArgs], + catch: bool, +) -> list[tuple[bool, typing.Any]]: + """Run `fn` over a chunk of argument tuples in one task. + + Top-level and closed over nothing so process pools can pickle it. + + Args: + fn: The callable applied per argument tuple. + chunk: The argument tuples for this task. + catch: Under ``on_error='return'`` each item's `Exception` is + captured as a ``(False, exc)`` outcome so one failure loses + no other results. `KeyboardInterrupt`/`SystemExit` always + escape -- errors may be *returned*, never swallowed. With + ``catch=False`` the first exception escapes, aborting the + chunk's remainder (documented fail-fast semantics). + + Returns: + One ``(ok, result_or_exception)`` pair per completed item. + """ + outcomes: list[tuple[bool, typing.Any]] = [] + for args in chunk: + if catch: + try: + outcomes.append((True, fn(*args))) + except Exception as exc: # noqa: BLE001 - returned, not silenced + outcomes.append((False, exc)) + else: + outcomes.append((True, fn(*args))) + return outcomes diff --git a/progressbar/_parallel/_decorator.py b/progressbar/_parallel/_decorator.py new file mode 100644 index 0000000..72437dd --- /dev/null +++ b/progressbar/_parallel/_decorator.py @@ -0,0 +1,138 @@ +"""`@parallel`: attach the batch verbs to a plain function. + +Sugar over the module verbs:: + + @progressbar.parallel(workers=4, pool='process') + def crunch(path: str) -> Result: ... + + + crunch(one_path) # unchanged direct call + crunch.map(paths) # parallel + bar, decorator config applied + +The decorator returns the *original function object* with bound verbs +attached as attributes. That shape is deliberate: pickling by +qualified name stays intact (``pool='process'`` needs it), and spawn's +re-import just re-runs the decoration. +""" + +from __future__ import annotations + +import functools +import inspect +import typing + +from . import ( + _async, + _sync, +) + + +class ParallelFunction(typing.Protocol): + """A function enriched with the parallel batch verbs.""" + + def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Any: + """The original, undecorated call.""" + ... + + def map( # noqa: A003 - mirrors the module verb + self, *iterables: typing.Any, **kwargs: typing.Any + ) -> list[typing.Any]: + """Parallel ordered map over the iterables; see the module verb.""" + ... + + def imap( + self, *iterables: typing.Any, **kwargs: typing.Any + ) -> typing.Generator[typing.Any, None, None]: + """Lazy ordered results; see the module verb.""" + ... + + def imap_unordered( + self, *iterables: typing.Any, **kwargs: typing.Any + ) -> typing.Generator[tuple[typing.Any, typing.Any], None, None]: + """Completion-order pairs; see the module verb.""" + ... + + def starmap( + self, iterable: typing.Any, **kwargs: typing.Any + ) -> list[typing.Any]: + """Parallel map over pre-tupled arguments; see the module verb.""" + ... + + def amap( + self, *iterables: typing.Any, **kwargs: typing.Any + ) -> typing.Coroutine[typing.Any, typing.Any, list[typing.Any]]: + """Async ordered map; see the module verb.""" + ... + + def aimap( + self, *iterables: typing.Any, **kwargs: typing.Any + ) -> typing.AsyncIterator[typing.Any]: + """Async lazy ordered results; see the module verb.""" + ... + + def aimap_unordered( + self, *iterables: typing.Any, **kwargs: typing.Any + ) -> typing.AsyncIterator[tuple[typing.Any, typing.Any]]: + """Async completion-order pairs; see the module verb.""" + ... + + +#: Verb name -> module implementation bound by the decorator. +_VERBS: dict[str, typing.Callable[..., typing.Any]] = { + 'map': _sync.map, + 'imap': _sync.imap, + 'imap_unordered': _sync.imap_unordered, + 'starmap': _sync.starmap, + 'amap': _async.amap, + 'aimap': _async.aimap, + 'aimap_unordered': _async.aimap_unordered, +} + + +def _bind( + verb: typing.Callable[..., typing.Any], + fn: typing.Callable[..., typing.Any], + config: dict[str, typing.Any], +) -> typing.Callable[..., typing.Any]: + """Close `verb` over `fn` and the decorator config.""" + + @functools.wraps(verb, assigned=('__doc__',), updated=()) + def bound(*iterables: typing.Any, **kwargs: typing.Any) -> typing.Any: + return verb(fn, *iterables, **{**config, **kwargs}) + + return bound + + +def parallel( + **config: typing.Any, +) -> typing.Callable[[typing.Callable[..., typing.Any]], ParallelFunction]: + """Attach the parallel batch verbs to a function. + + Args: + **config: Default keywords for every attached verb (`workers`, + `pool`, `bar`, `on_error`, ...); individual calls override + them. + + Returns: + A decorator returning the same function object with `.map`, + `.imap`, `.imap_unordered`, `.starmap`, `.amap`, `.aimap` and + `.aimap_unordered` attached. + + Raises: + TypeError: The target is not a plain named function. Lambdas + and bound methods are rejected because ``pool='process'`` + pickles the function by qualified name. + """ + + def decorate(fn: typing.Callable[..., typing.Any]) -> ParallelFunction: + if not inspect.isfunction(fn) or fn.__name__ == '': + raise TypeError( + f'parallel() requires a plain named function, got ' + f'{fn!r} (needed so pool="process" can pickle it by ' + f'qualified name)' + ) + for name, verb in _VERBS.items(): + setattr(fn, name, _bind(verb, fn, config)) + return typing.cast(ParallelFunction, fn) + + return decorate diff --git a/progressbar/_parallel/_display.py b/progressbar/_parallel/_display.py new file mode 100644 index 0000000..d481a31 --- /dev/null +++ b/progressbar/_parallel/_display.py @@ -0,0 +1,345 @@ +"""Display backends for the parallel verbs. + +One small protocol so both execution engines can drive any of four +rendering modes -- ``'plain'`` (one aggregate bar), ``'multi'`` +(a MultiBar with per-task sub-bars), ``False`` (silent), or a +caller-configured bar instance -- through the same six calls. + +The keep-alive contract lives here: `PlainDisplay` constructs its bar +with ``poll_interval`` set, because `ProgressBar.update()` without a +new value only redraws timers/animations when the bar's own +``poll_interval`` says a redraw is due. A default-configured bar +no-ops, which would freeze ETA/spinners during long tasks. +""" + +from __future__ import annotations + +import os +import sys +import typing + +from .. import ( + bar as bar_module, + base, + fast as fast_module, +) +from . import _common + +if typing.TYPE_CHECKING: + from .. import multi as multi_module + + +@typing.runtime_checkable +class Display(typing.Protocol): + """What the execution engines require from a rendering backend.""" + + def start(self, total: typing.Any) -> None: + """Begin rendering a run of `total` items (or `UnknownLength`).""" + ... + + def task_started( + self, seq: int, label: str + ) -> bar_module.ProgressBar | None: + """Register in-flight task `seq`; return its own bar, if any.""" + ... + + def task_finished(self, seq: int, ok: bool) -> None: + """Retire task `seq` from the in-flight set.""" + ... + + def advance(self, n: int = 1) -> None: + """Count `n` more items as completed.""" + ... + + def tick(self) -> None: + """Keep time widgets moving when nothing completed this poll.""" + ... + + def finish(self, *, success: bool = True) -> None: + """Stop rendering; a failed run must not jump the bar to 100%.""" + ... + + +class NullDisplay: + """The ``bar=False`` backend: the machinery runs, nothing renders.""" + + def start(self, total: typing.Any) -> None: + """Ignore the run start.""" + + def task_started( + self, seq: int, label: str + ) -> bar_module.ProgressBar | None: + """Report no per-task bar.""" + return None + + def task_finished(self, seq: int, ok: bool) -> None: + """Ignore the task end.""" + + def advance(self, n: int = 1) -> None: + """Ignore progress.""" + + def tick(self) -> None: + """Ignore the poll.""" + + def finish(self, *, success: bool = True) -> None: + """Ignore the run end.""" + + +def _select_bar_class( + bar_kwargs: dict[str, typing.Any], +) -> type[bar_module.ProgressBar]: + """Pick the lean fast bar unless a kwarg needs the widget machinery. + + Mirrors the dispatch rule of `progressbar.shortcuts.progressbar`: + anything widget-shaped (custom widgets, variables, units, postfix) + forces the full bar; the plain percentage/ETA case takes the cheap + renderer. + """ + needs_full: bool = bool( + bar_kwargs.get('widgets') + or bar_kwargs.get('variables') + or bar_kwargs.get('unit_scale') + or bar_kwargs.get('postfix') + or bar_kwargs.get('unit', 'it') != 'it' + or os.environ.get('PROGRESSBAR_DISABLE_FASTPATH') + ) + return ( + bar_module.ProgressBar if needs_full else fast_module.FastProgressBar + ) + + +class PlainDisplay: + """One aggregate bar counting completed items.""" + + _bar: bar_module.ProgressBar + _owned: bool + _started_by_us: bool + _value: int + + def __init__( + self, + *, + total: typing.Any, + poll_interval: float, + bar_kwargs: dict[str, typing.Any], + instance: bar_module.ProgressBar | None = None, + ) -> None: + """Create the bar (or adopt `instance` without reconfiguring it).""" + self._owned = instance is None + self._started_by_us = False + self._value = 0 + if instance is not None: + self._bar = instance + else: + # poll_interval is what makes no-progress ticks redraw; a + # caller-supplied value in bar_kwargs wins (same knob). + kwargs: dict[str, typing.Any] = dict(bar_kwargs) + kwargs.setdefault('poll_interval', poll_interval) + kwargs.setdefault('max_value', total) + bar_class: type[bar_module.ProgressBar] = _select_bar_class(kwargs) + self._bar = bar_class(**kwargs) + + def start(self, total: typing.Any) -> None: + """Start the bar unless the caller already started it.""" + if not self._bar.started(): + self._bar.start() + self._started_by_us = True + + def task_started( + self, seq: int, label: str + ) -> bar_module.ProgressBar | None: + """Report no per-task bar (plain mode has only the aggregate).""" + return None + + def task_finished(self, seq: int, ok: bool) -> None: + """Nothing tracked per task in plain mode.""" + + def advance(self, n: int = 1) -> None: + """Add `n` completions and redraw.""" + self._value += n + self._bar.update(self._value) + + def tick(self) -> None: + """Redraw with no new value so time widgets stay alive.""" + self._bar.update() + + def finish(self, *, success: bool = True) -> None: + """Finish the bar; only if this display started it.""" + if self._started_by_us: + # dirty=True on failure: keep the last real value on screen + # instead of forcing the bar to max. + self._bar.finish(dirty=not success) + + +class MultiDisplay: + """A MultiBar: one overall bar plus one bar per in-flight task. + + Per-task bars are keyed ``'{seq}: {label}'`` -- ``seq`` is unique + per run, so two items with the same ``str()`` can never collide on + a key (a plain-label key would silently merge them and corrupt the + display). Finished task bars are deleted immediately: MultiBar's + own default keeps them around for an hour, which would stack + thousands of dead lines on a big batch. + + Rendering is done by MultiBar's daemon thread at + ``update_interval=poll_interval``, so `tick` needs no work here. + Best suited to modest worker counts: the block occupies one + terminal row per in-flight task plus one for the total. + """ + + #: Overall bar's key in the multibar (also its visible label). + _TOTAL_KEY: typing.ClassVar[str] = 'Total' + #: Bounded wait for the render thread; it wakes every + #: `update_interval` seconds, so this is generous. + _STOP_TIMEOUT: typing.ClassVar[float] = 5.0 + + multibar: multi_module.MultiBar + _keys: dict[int, str] + _owned: bool + _started_thread: bool + _value: int + _poll_interval: float + + def __init__( + self, + *, + total: typing.Any, + poll_interval: float, + bar_kwargs: dict[str, typing.Any], + instance: multi_module.MultiBar | None = None, + ) -> None: + """Create (or adopt) the multibar and its overall bar.""" + # Deferred import: keeps `import progressbar` + plain-mode use + # from paying for the multibar/widget machinery. + from .. import multi as multi_module + + self._keys = {} + self._value = 0 + self._started_thread = False + self._poll_interval = poll_interval + self._owned = instance is None + kwargs: dict[str, typing.Any] = dict(bar_kwargs) + if instance is not None: + self.multibar = instance + else: + self.multibar = multi_module.MultiBar( + fd=kwargs.pop('fd', sys.stderr), + update_interval=poll_interval, + show_finished=False, + remove_finished=0, + sort_reverse=False, + ) + # Remaining bar kwargs style the overall bar. + kwargs.setdefault('poll_interval', poll_interval) + kwargs.setdefault('max_value', total) + self._overall: bar_module.ProgressBar = bar_module.ProgressBar( + **kwargs + ) + + def start(self, total: typing.Any) -> None: + """Add and start the overall bar, then the render thread.""" + self.multibar[self._TOTAL_KEY] = self._overall + self._overall.start() + # `_thread` is the only handle MultiBar exposes for "already + # running"; an adopted, already-started instance must not be + # started twice (MultiBar asserts on that). + if self.multibar._thread is None: # noqa: SLF001 + self.multibar.start() + self._started_thread = True + + def task_started( + self, seq: int, label: str + ) -> bar_module.ProgressBar | None: + """Add a per-task bar and hand it out for `current_task_bar`.""" + key: str = f'{seq}: {label}' + self._keys[seq] = key + task_bar: bar_module.ProgressBar = bar_module.ProgressBar( + max_value=base.UnknownLength, + ) + self.multibar[key] = task_bar + task_bar.start() + return task_bar + + def task_finished(self, seq: int, ok: bool) -> None: + """Drop the task's bar -- finished bars must not pile up.""" + key: str | None = self._keys.pop(seq, None) + if key is not None and key in self.multibar: + del self.multibar[key] + + def advance(self, n: int = 1) -> None: + """Count completions on the overall bar. + + The overall bar is `paused` (MultiBar contract), so `update` + only records the value; the render thread draws it. + """ + self._value += n + self._overall.update(self._value) + + def tick(self) -> None: + """No-op: the render thread redraws every `poll_interval`.""" + + def finish(self, *, success: bool = True) -> None: + """Finish the overall bar and wind down the render thread.""" + for seq in list(self._keys): + self.task_finished(seq, ok=success) + self._overall.finish(dirty=not success) + if self._started_thread: + # One last frame so the final state is on screen even if + # the render thread never woke between finish and stop. + self.multibar.render(force=True) + self.multibar.stop(timeout=self._STOP_TIMEOUT) + + +def make_display( + bar_mode: typing.Any, + *, + total: typing.Any, + poll_interval: float, + bar_kwargs: dict[str, typing.Any], +) -> Display: + """Build the display backend for one run. + + Args: + bar_mode: ``'plain'`` | ``'multi'`` | ``False`` | a + `ProgressBar` or `MultiBar` instance to drive. + total: Item count or `base.UnknownLength`. + poll_interval: Redraw cadence; also the engines' wake interval. + bar_kwargs: Validated passthrough for the constructed bar. + + Raises: + TypeError: For unknown or unvalidated bar kwargs, or an + unrecognized `bar_mode`. + """ + _common.validate_bar_kwargs(bar_kwargs) + if bar_mode is False or bar_mode is None: + return NullDisplay() + if bar_mode == 'plain': + return PlainDisplay( + total=total, poll_interval=poll_interval, bar_kwargs=bar_kwargs + ) + if bar_mode == 'multi': + return MultiDisplay( + total=total, poll_interval=poll_interval, bar_kwargs=bar_kwargs + ) + if isinstance(bar_mode, bar_module.ProgressBar): + return PlainDisplay( + total=total, + poll_interval=poll_interval, + bar_kwargs=bar_kwargs, + instance=bar_mode, + ) + # Deferred import mirrors MultiDisplay's: only pay for the multibar + # machinery when a MultiBar is actually in play. + from .. import multi as multi_module + + if isinstance(bar_mode, multi_module.MultiBar): + return MultiDisplay( + total=total, + poll_interval=poll_interval, + bar_kwargs=bar_kwargs, + instance=bar_mode, + ) + raise TypeError( + f'bar={bar_mode!r} is not a valid mode: expected "plain", ' + f'"multi", False, a ProgressBar or a MultiBar' + ) diff --git a/progressbar/_parallel/_shell.py b/progressbar/_parallel/_shell.py new file mode 100644 index 0000000..bc74778 --- /dev/null +++ b/progressbar/_parallel/_shell.py @@ -0,0 +1,189 @@ +"""`run()`: execute a shell command per item, in parallel, with a bar. + +A progress-bar'd ``xargs -P`` in Python. Templates never go through +`str.format` -- only the exact placeholder tokens ``{}`` and ``{item}`` +are substituted -- so commands containing literal braces (``awk +'{print $1}'``) pass through untouched. +""" + +from __future__ import annotations + +import functools +import os +import shlex +import subprocess +import typing + +from . import _sync + +#: The two placeholder spellings recognized in command templates. +_PLACEHOLDERS: tuple[str, str] = ('{}', '{item}') + +#: A command template: a string, an argv list, or a callable that +#: builds the argv for one item. +CommandT = ( + str + | typing.Sequence[str] + | typing.Callable[[typing.Any], typing.Sequence[str]] +) + + +def _substitute(token: str, item_text: str) -> str: + """Replace the placeholder spellings inside one token.""" + for placeholder in _PLACEHOLDERS: + token = token.replace(placeholder, item_text) + return token + + +def _has_placeholder(tokens: typing.Iterable[str]) -> bool: + """Return whether any token contains a placeholder.""" + return any( + placeholder in token + for token in tokens + for placeholder in _PLACEHOLDERS + ) + + +def build_argv( + command: CommandT, item: typing.Any, *, shell: bool +) -> list[str] | str: + """Build the command for one item from a template. + + Args: + command: A str template (split with `shlex.split`, non-POSIX + mode on Windows so backslash paths survive), an argv list + template, or a callable returning the argv. In the str and + list forms every ``{}``/``{item}`` inside a token is + replaced by ``str(item)`` -- an item containing spaces + stays a single argv element. Without any placeholder the + item is appended as the final argument (the ``xargs`` + convention). + item: The batch item; substituted as ``str(item)``. + shell: With the str form, substitute into (and return) the + whole command string for ``subprocess.run(shell=True)``; + an appended item is `shlex.quote`-escaped. The caller must + trust its items -- see `run`. + + Returns: + The argv list, or the command string when ``shell=True``. + """ + item_text: str = str(item) + if callable(command): + return [str(part) for part in command(item)] + if isinstance(command, str): + if shell: + if _has_placeholder((command,)): + return _substitute(command, item_text) + return f'{command} {shlex.quote(item_text)}' + tokens: list[str] = shlex.split(command, posix=os.name != 'nt') + else: + tokens = list(command) + if _has_placeholder(tokens): + return [_substitute(token, item_text) for token in tokens] + return [*tokens, item_text] + + +def _run_one( + command: CommandT, + item: typing.Any, + *, + check: bool, + capture_output: bool, + text: bool, + shell: bool, + cwd: typing.Any, + env: typing.Any, +) -> subprocess.CompletedProcess[typing.Any]: + """Execute the command for one item (thread-pool worker).""" + argv: list[str] | str = build_argv(command, item, shell=shell) + # The argv is assembled from the caller's own template and items; + # shell=True is opt-in and documented as trusting both. + return subprocess.run( # noqa: S603, PLW1510 + argv, + check=check, + capture_output=capture_output, + text=text, + shell=shell, # noqa: S602 + cwd=cwd, + env=env, + ) + + +def make_runner( + command: CommandT, + *, + check: bool = True, + capture_output: bool = True, + text: bool = True, + shell: bool = False, + cwd: typing.Any = None, + env: typing.Any = None, +) -> typing.Callable[[typing.Any], subprocess.CompletedProcess[typing.Any]]: + """Bind a command template into a per-item callable for `map`.""" + return functools.partial( + _run_one, + command, + check=check, + capture_output=capture_output, + text=text, + shell=shell, + cwd=cwd, + env=env, + ) + + +def run( + command: CommandT, + items: typing.Iterable[typing.Any], + /, + *, + check: bool = True, + capture_output: bool = True, + text: bool = True, + shell: bool = False, + cwd: typing.Any = None, + env: typing.Any = None, + **kwargs: typing.Any, +) -> list[subprocess.CompletedProcess[typing.Any]]: + """Run a shell command for every item in parallel, with a bar. + + ``progressbar.run('gzip -k {}', files, workers=4)`` is a + progress-bar'd ``xargs -P``. Subprocesses release the GIL, so this + always runs on threads (`Pool.run` reuses a pool's executor). + + Args: + command: Template -- see `build_argv` for the three forms and + the placeholder rules. + items: The batch; each becomes one subprocess. + check: Raise `subprocess.CalledProcessError` on a non-zero + exit (feeding `on_error` like any other worker error). + capture_output: Capture stdout/stderr into the results -- + the default, so child output cannot corrupt the bar. + text: Decode captured output as text. + shell: Run through the shell (str form only). The items are + substituted into the command line: only use with trusted + items, this is the documented injection risk. + cwd: Working directory for the subprocesses. + env: Environment for the subprocesses. + **kwargs: The shared execution keywords (`workers`, `bar`, + `on_error`, `timeout`, ...); see `_sync.execute`. + + Returns: + One `subprocess.CompletedProcess` per item, in input order + (exceptions in place under ``on_error='return'``). + """ + if 'pool' in kwargs: + raise TypeError( + 'run() always uses threads (subprocesses release the GIL); ' + 'use Pool.run() to reuse an existing pool' + ) + runner = make_runner( + command, + check=check, + capture_output=capture_output, + text=text, + shell=shell, + cwd=cwd, + env=env, + ) + return _sync.map(runner, items, pool='thread', **kwargs) diff --git a/progressbar/_parallel/_sync.py b/progressbar/_parallel/_sync.py new file mode 100644 index 0000000..8fc308a --- /dev/null +++ b/progressbar/_parallel/_sync.py @@ -0,0 +1,825 @@ +"""The `concurrent.futures` engine behind the sync parallel verbs. + +One generator -- `execute` -- owns the whole coordination pattern: +windowed submission, a done-queue that costs O(1) per completion, and +poll-timeout ticks that keep the bar animating while nothing finishes. +Every public sync verb (`map`, and its siblings) is a thin consumer of +`execute`'s completion stream. +""" + +from __future__ import annotations + +import concurrent.futures +import functools +import inspect +import queue +import sys +import time +import typing + +from . import ( + _common, + _display, +) + +#: One completion event: (item index, argument tuple, ok, result/error). +Completion = tuple[int, _common.ItemArgs, bool, typing.Any] + +#: `pool=` values that name an executor kind rather than an instance. +_POOL_KINDS: frozenset[str] = frozenset({'thread', 'process', 'interpreter'}) + +#: Default seconds between coordinator wakeups; doubles as the bar's +#: redraw interval (one knob -- see the keep-alive contract). +DEFAULT_POLL_INTERVAL: float = 0.1 + + +def _pool_kind(pool: typing.Any) -> str: + """Map a `pool=` argument to 'thread'/'process'/'interpreter'.""" + if isinstance(pool, str): + return pool + if isinstance(pool, concurrent.futures.ProcessPoolExecutor): + return 'process' + return 'thread' + + +def _adopt_executor( + pool: concurrent.futures.Executor, + workers: int | None, + constructor_kwargs: dict[str, typing.Any], +) -> tuple[concurrent.futures.Executor, bool, int]: + """Adopt a caller-owned executor; reject construction kwargs.""" + configured: list[str] = [ + name for name, value in constructor_kwargs.items() if value + ] + if configured: + raise ValueError( + f'{configured!r} configure a new executor and cannot be ' + f'combined with an existing executor instance' + ) + return pool, False, _common.resolve_workers(workers, _pool_kind(pool)) + + +def _thread_executor( + workers: int, + initializer: typing.Callable[..., None] | None, + initargs: tuple[typing.Any, ...], + thread_name_prefix: str, +) -> concurrent.futures.Executor: + """Build the owned thread pool.""" + return concurrent.futures.ThreadPoolExecutor( + max_workers=workers, + thread_name_prefix=thread_name_prefix, + initializer=initializer, + initargs=initargs, + ) + + +def _process_executor( + workers: int, + initializer: typing.Callable[..., None] | None, + initargs: tuple[typing.Any, ...], + mp_context: typing.Any, + max_tasks_per_child: int | None, +) -> concurrent.futures.Executor: + """Build the owned process pool.""" + process_kwargs: dict[str, typing.Any] = { + 'max_workers': workers, + 'mp_context': mp_context, + 'initializer': initializer, + 'initargs': initargs, + } + if max_tasks_per_child is not None: + # Assignment before the version gate so the line is reachable + # (and measured) on every Python; on 3.10 the raise below + # discards the dict anyway. + process_kwargs['max_tasks_per_child'] = max_tasks_per_child + if sys.version_info < (3, 11): # pragma: no cover - version gate + raise ValueError('max_tasks_per_child requires Python 3.11+') + return concurrent.futures.ProcessPoolExecutor(**process_kwargs) + + +def _interpreter_executor( + workers: int, + initializer: typing.Callable[..., None] | None, + initargs: tuple[typing.Any, ...], +) -> concurrent.futures.Executor: + """Build the owned interpreter pool (Python 3.14+).""" + try: + # typed Any: the class only exists on 3.14+, so pyright has no + # signature for it on the 3.10 floor. + interpreter_pool: typing.Any = ( + concurrent.futures.InterpreterPoolExecutor # type: ignore[attr-defined] + ) + except AttributeError: # pragma: no cover - version gate (<=3.13) + raise ValueError('pool="interpreter" requires Python 3.14+') from None + return interpreter_pool( # pragma: no cover - reachable on 3.14+ only + max_workers=workers, + initializer=initializer, + initargs=initargs, + ) + + +def resolve_executor( + pool: typing.Any, + workers: int | None, + *, + initializer: typing.Callable[..., None] | None, + initargs: tuple[typing.Any, ...], + mp_context: typing.Any, + max_tasks_per_child: int | None, + thread_name_prefix: str, +) -> tuple[concurrent.futures.Executor, bool, int]: + """Create (or adopt) the executor for one run. + + Args: + pool: ``'thread'`` | ``'process'`` | ``'interpreter'`` (3.14+) + or an existing `concurrent.futures.Executor` instance. + workers: Pool size; `None` uses the executor defaults. + initializer: Per-worker setup callable, forwarded verbatim. + initargs: Arguments for `initializer`. + mp_context: `multiprocessing` context for process pools. + max_tasks_per_child: Worker recycling limit (3.11+). + thread_name_prefix: Thread pool naming, forwarded verbatim. + + Returns: + ``(executor, owned, effective_workers)`` -- `owned` is whether + this run created (and must shut down) the executor. + + Raises: + ValueError: Unknown `pool` string, construction kwargs combined + with an executor instance or the wrong pool kind, or a + version-gated option on an unsupported Python. + """ + if isinstance(pool, concurrent.futures.Executor): + return _adopt_executor( + pool, + workers, + { + 'initializer': initializer, + 'initargs': initargs, + 'mp_context': mp_context, + 'max_tasks_per_child': max_tasks_per_child, + 'thread_name_prefix': thread_name_prefix, + }, + ) + if pool not in _POOL_KINDS: + raise ValueError( + f'pool={pool!r} is not valid: expected "thread", "process", ' + f'"interpreter" or a concurrent.futures.Executor instance' + ) + if pool != 'process' and ( + mp_context is not None or max_tasks_per_child is not None + ): + raise ValueError( + 'mp_context/max_tasks_per_child only apply to process pools' + ) + if pool != 'thread' and thread_name_prefix: + raise ValueError('thread_name_prefix only applies to thread pools') + + effective_workers: int = _common.resolve_workers(workers, pool) + executor: concurrent.futures.Executor + if pool == 'thread': + executor = _thread_executor( + effective_workers, initializer, initargs, thread_name_prefix + ) + elif pool == 'interpreter': + executor = _interpreter_executor( + effective_workers, initializer, initargs + ) + else: + executor = _process_executor( + effective_workers, + initializer, + initargs, + mp_context, + max_tasks_per_child, + ) + return executor, True, effective_workers + + +def _indexed_chunks( + iterables: tuple[typing.Iterable[typing.Any], ...], + chunksize: int, +) -> typing.Iterator[tuple[int, list[_common.ItemArgs]]]: + """Yield ``(first item index, chunk)`` pairs, consuming lazily.""" + index: int = 0 + for chunk in _common.iter_chunks(iterables, chunksize): + yield index, chunk + index += len(chunk) + + +class _Run: + """State and coordination for one `execute` invocation. + + Split from `execute` so each concern -- submission, completion + handling, deadline, shutdown -- stays a small method; `execute` + itself only owns the generator's try/finally lifecycle. + """ + + fn: typing.Callable[..., typing.Any] + kind: str + total: typing.Any + on_error: str + catch: bool + single: bool + window: int + timeout: float | None + poll_interval: float + deadline: float | None + executor: concurrent.futures.Executor + owned: bool + display: _display.Display + done: queue.SimpleQueue[concurrent.futures.Future[typing.Any]] + in_flight: dict[ + concurrent.futures.Future[typing.Any], + tuple[int, list[_common.ItemArgs], int], + ] + chunk_source: typing.Iterator[tuple[int, list[_common.ItemArgs]]] + seq: int + + def __init__( + self, + fn: typing.Callable[..., typing.Any], + iterables: tuple[typing.Iterable[typing.Any], ...], + *, + workers: int | None, + pool: typing.Any, + bar: typing.Any, + on_error: str, + chunksize: int | None, + buffersize: int | None, + timeout: float | None, + poll_interval: float, + initializer: typing.Callable[..., None] | None, + initargs: tuple[typing.Any, ...], + mp_context: typing.Any, + max_tasks_per_child: int | None, + thread_name_prefix: str, + bar_kwargs: dict[str, typing.Any], + ) -> None: + """Validate the configuration and set up executor and display.""" + if inspect.iscoroutinefunction(fn): + raise TypeError( + f'{fn!r} is a coroutine function; use progressbar.amap() ' + f'-- the sync verbs cannot await it' + ) + if on_error not in ('raise', 'return'): + raise ValueError( + f"on_error={on_error!r} is not valid: expected 'raise' " + f"or 'return'" + ) + _common.validate_bar_kwargs(bar_kwargs) + + self.fn = fn + self.kind = _pool_kind(pool) + self.total = _common.detect_total(iterables) + self.on_error = on_error + self.catch = on_error == 'return' + self.single = len(iterables) == 1 + self.timeout = timeout + self.poll_interval = poll_interval + self.deadline = None if timeout is None else time.monotonic() + timeout + self.executor, self.owned, effective_workers = resolve_executor( + pool, + workers, + initializer=initializer, + initargs=initargs, + mp_context=mp_context, + max_tasks_per_child=max_tasks_per_child, + thread_name_prefix=thread_name_prefix, + ) + if chunksize is None: + chunksize = ( + _common.auto_chunksize(self.total, effective_workers) + if self.kind in ('process', 'interpreter') + else 1 + ) + self.window = ( + buffersize + if buffersize is not None + else _common.default_buffersize(effective_workers) + ) + self.display = _display.make_display( + bar, + total=self.total, + poll_interval=poll_interval, + bar_kwargs=bar_kwargs, + ) + self.done = queue.SimpleQueue() + self.in_flight = {} + self.chunk_source = _indexed_chunks(iterables, chunksize) + self.seq = 0 + + def completions(self) -> typing.Iterator[Completion]: + """Drive the run, yielding per-item events in completion order.""" + self.display.start(self.total) + while len(self.in_flight) < self.window and self._submit_one(): + pass + while self.in_flight: + self._check_deadline() + future = self._next_done() + if future is not None: + yield from self._handle(future) + + def _submit_one(self) -> bool: + """Submit the next chunk; `False` when the input is exhausted.""" + indexed: tuple[int, list[_common.ItemArgs]] | None = next( + self.chunk_source, None + ) + if indexed is None: + return False + start_index, chunk = indexed + self.seq += 1 + label: str = str(_common.item_of(chunk[0], self.single)) + task_bar = self.display.task_started(self.seq, label) + inner: typing.Callable[[], list[tuple[bool, typing.Any]]] = ( + functools.partial(_common.run_chunk, self.fn, chunk, self.catch) + ) + if task_bar is not None and self.kind == 'thread': + # Threads share our address space, so the worker can update + # its sub-bar through `current_task_bar()`. Process (and + # interpreter) workers cannot -- the bar object does not + # survive pickling; see the phase-2 note in the docs. + inner = _common.with_task_bar(task_bar, inner) + future: concurrent.futures.Future[typing.Any] = self.executor.submit( + inner + ) + self.in_flight[future] = (start_index, chunk, self.seq) + future.add_done_callback(self.done.put) + return True + + def _next_done( + self, + ) -> concurrent.futures.Future[typing.Any] | None: + """Wait one poll for a completion; tick the display on none.""" + try: + return self.done.get(timeout=self.poll_interval) + except queue.Empty: + self.display.tick() + return None + + def _check_deadline(self) -> None: + """Raise once the overall `timeout` budget is spent.""" + if self.deadline is not None and time.monotonic() > self.deadline: + raise concurrent.futures.TimeoutError( + f'parallel execution exceeded timeout={self.timeout}' + ) + + def _handle( + self, future: concurrent.futures.Future[typing.Any] + ) -> typing.Iterator[Completion]: + """Turn one finished future into per-item completion events.""" + start_index, chunk, chunk_seq = self.in_flight.pop(future) + error: BaseException | None = future.exception() + if error is not None: + # Fail-fast fn errors (catch=False), machinery errors (e.g. + # BrokenProcessPool, unpicklable results) and + # KeyboardInterrupt/SystemExit escaping run_chunk all land + # here: never silently, always raised. + self.display.task_finished(chunk_seq, ok=False) + raise error + outcomes: list[tuple[bool, typing.Any]] = future.result() + self.display.task_finished(chunk_seq, ok=all(ok for ok, _ in outcomes)) + self.display.advance(len(chunk)) + self._submit_one() + for offset, (ok, value) in enumerate(outcomes): + yield start_index + offset, chunk[offset], ok, value + + def close(self, *, interrupted: bool, success: bool) -> None: + """Cancel leftovers and release executor and display.""" + for future in self.in_flight: + future.cancel() + if self.owned: + self.executor.shutdown(wait=not interrupted, cancel_futures=True) + self.display.finish(success=success) + + +def execute( + fn: typing.Callable[..., typing.Any], + iterables: tuple[typing.Iterable[typing.Any], ...], + *, + workers: int | None = None, + pool: typing.Any = 'thread', + bar: typing.Any = 'plain', + on_error: str = 'raise', + chunksize: int | None = None, + buffersize: int | None = None, + timeout: float | None = None, + poll_interval: float = DEFAULT_POLL_INTERVAL, + initializer: typing.Callable[..., None] | None = None, + initargs: tuple[typing.Any, ...] = (), + mp_context: typing.Any = None, + max_tasks_per_child: int | None = None, + thread_name_prefix: str = '', + **bar_kwargs: typing.Any, +) -> typing.Iterator[Completion]: + """Run `fn` over zipped `iterables`, yielding completion events. + + The single sync coordination loop. Yields one ``(index, args, ok, + value)`` tuple per item in *completion order*; consumers impose + their own ordering (`map` collects by index, `imap` holds back, + `imap_unordered` passes through). + + Cleanup is the generator's ``finally``: closing this generator (an + early ``break`` in a consumer) cancels unsubmitted work and shuts + down an owned executor. Running tasks cannot be interrupted -- they + finish in the background of the shutdown; on `KeyboardInterrupt` + the shutdown does not wait for them. + + Raises: + TypeError: `fn` is a coroutine function (belongs to `amap`), or + an unknown bar keyword was passed. + ValueError: `on_error` is not ``'raise'``/``'return'``, or the + executor configuration is invalid. + concurrent.futures.TimeoutError: The overall `timeout` expired; + pending work is cancelled first. + """ + run: _Run = _Run( + fn, + iterables, + workers=workers, + pool=pool, + bar=bar, + on_error=on_error, + chunksize=chunksize, + buffersize=buffersize, + timeout=timeout, + poll_interval=poll_interval, + initializer=initializer, + initargs=initargs, + mp_context=mp_context, + max_tasks_per_child=max_tasks_per_child, + thread_name_prefix=thread_name_prefix, + bar_kwargs=bar_kwargs, + ) + interrupted: bool = False + success: bool = False + try: + yield from run.completions() + success = True + except BaseException as error: + # Ctrl-C means "stop now" and a timeout is a promise about + # wall-clock time -- neither may block on running tasks. Plain + # fail-fast still waits so results/streams aren't torn down + # under a worker mid-write. + interrupted = isinstance( + error, + (KeyboardInterrupt, TimeoutError, concurrent.futures.TimeoutError), + ) + raise + finally: + run.close(interrupted=interrupted, success=success) + + +#: Keywords that configure a `Pool`'s lazily created executor rather +#: than an individual run. +_EXECUTOR_KWARG_NAMES: frozenset[str] = frozenset( + { + 'initializer', + 'initargs', + 'mp_context', + 'max_tasks_per_child', + 'thread_name_prefix', + } +) + + +class Pool: + """A reusable executor plus per-call defaults for the sync verbs. + + The flat verbs create and destroy an executor per call; a `Pool` + keeps one alive across calls:: + + with progressbar.Pool(8) as pool: + first = pool.map(fetch, urls) + second = pool.map(fetch, more_urls) + + Positional shorthand: ``Pool(8)`` is eight threads, ``Pool(8, + 'process')`` eight processes. ``Pool(executor=existing)`` adopts a + caller-owned executor (never shut down here). Every other keyword + becomes a per-call default that individual calls can override. + + The executor is created lazily on first use, so an unused + ``Pool(kind='process')`` spawns nothing. + """ + + _workers: int | None + _kind: str + _external: concurrent.futures.Executor | None + _executor: concurrent.futures.Executor | None + _executor_kwargs: dict[str, typing.Any] + _defaults: dict[str, typing.Any] + + def __init__( + self, + workers: int | None = None, + kind: str = 'thread', + *, + executor: concurrent.futures.Executor | None = None, + **defaults: typing.Any, + ) -> None: + """Validate eagerly (fail fast); create nothing yet.""" + executor_kwargs: dict[str, typing.Any] = { + name: defaults.pop(name) + for name in list(defaults) + if name in _EXECUTOR_KWARG_NAMES + } + if executor is not None: + if workers is not None or kind != 'thread' or executor_kwargs: + raise ValueError( + 'workers/kind/executor-construction options cannot ' + 'be combined with an existing executor instance' + ) + elif kind not in _POOL_KINDS: + raise ValueError( + f'kind={kind!r} is not valid: expected "thread", ' + f'"process" or "interpreter"' + ) + self._workers = workers + self._kind = kind + self._external = executor + self._executor = None + self._executor_kwargs = executor_kwargs + self._defaults = defaults + + @property + def executor(self) -> concurrent.futures.Executor: + """The underlying executor, created on first access.""" + if self._external is not None: + return self._external + if self._executor is None: + self._executor, _owned, _workers = resolve_executor( + self._kind, + self._workers, + initializer=self._executor_kwargs.get('initializer'), + initargs=self._executor_kwargs.get('initargs', ()), + mp_context=self._executor_kwargs.get('mp_context'), + max_tasks_per_child=self._executor_kwargs.get( + 'max_tasks_per_child' + ), + thread_name_prefix=self._executor_kwargs.get( + 'thread_name_prefix', '' + ), + ) + return self._executor + + def _merged(self, kwargs: dict[str, typing.Any]) -> dict[str, typing.Any]: + """Per-call keywords override the pool's defaults.""" + return {**self._defaults, **kwargs, 'pool': self.executor} + + def map( # noqa: A003 - mirrors the module verb + self, + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, + ) -> list[typing.Any]: + """`map` on this pool's executor; see the module `map`.""" + return map(fn, *iterables, **self._merged(kwargs)) + + def imap( + self, + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, + ) -> typing.Generator[typing.Any, None, None]: + """`imap` on this pool's executor; see the module `imap`.""" + return imap(fn, *iterables, **self._merged(kwargs)) + + def imap_unordered( + self, + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, + ) -> typing.Generator[tuple[typing.Any, typing.Any], None, None]: + """`imap_unordered` on this pool's executor; see the module verb.""" + return imap_unordered(fn, *iterables, **self._merged(kwargs)) + + def starmap( + self, + fn: typing.Callable[..., typing.Any], + iterable: typing.Iterable[typing.Any], + /, + **kwargs: typing.Any, + ) -> list[typing.Any]: + """`starmap` on this pool's executor; see the module `starmap`.""" + return starmap(fn, iterable, **self._merged(kwargs)) + + def run( + self, + command: typing.Any, + items: typing.Iterable[typing.Any], + /, + **kwargs: typing.Any, + ) -> list[typing.Any]: + """`run` a shell command per item on this pool's executor.""" + # Deferred import: _shell imports this module. + from . import _shell + + runner = _shell.make_runner( + command, + **{ + name: kwargs.pop(name) + for name in ( + 'check', + 'capture_output', + 'text', + 'shell', + 'cwd', + 'env', + ) + if name in kwargs + }, + ) + return self.map(runner, items, **kwargs) + + def shutdown( + self, wait: bool = True, *, cancel_futures: bool = False + ) -> None: + """Shut down the owned executor; adopted executors are spared.""" + if self._executor is not None: + self._executor.shutdown(wait=wait, cancel_futures=cancel_futures) + + def __enter__(self) -> Pool: + """Return the pool; the executor still waits for first use.""" + return self + + def __exit__(self, *exc_info: typing.Any) -> None: + """Shut down the owned executor, waiting for running work.""" + self.shutdown(wait=True) + + +def _star_call( + fn: typing.Callable[..., typing.Any], args: typing.Any +) -> typing.Any: + """Unpack one `starmap` argument tuple into a call (picklable).""" + return fn(*args) + + +def starmap( + fn: typing.Callable[..., typing.Any], + iterable: typing.Iterable[typing.Any], + /, + **kwargs: typing.Any, +) -> list[typing.Any]: + """`map` over pre-tupled arguments (``multiprocessing.Pool.starmap``). + + ``starmap(fn, [(1, 2), (3, 4)])`` calls ``fn(1, 2)`` and + ``fn(3, 4)`` in parallel. See `execute` for keywords. + """ + return map(functools.partial(_star_call, fn), iterable, **kwargs) + + +def thread_map( + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> list[typing.Any]: + """`map` pinned to a thread pool (tqdm-compatible spelling).""" + if 'pool' in kwargs: + raise TypeError( + "thread_map() already sets pool='thread'; use map() to pick " + 'a pool explicitly' + ) + return map(fn, *iterables, pool='thread', **kwargs) + + +def process_map( + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> list[typing.Any]: + """`map` pinned to a process pool (tqdm-compatible spelling).""" + if 'pool' in kwargs: + raise TypeError( + "process_map() already sets pool='process'; use map() to " + 'pick a pool explicitly' + ) + return map(fn, *iterables, pool='process', **kwargs) + + +def as_completed( + futures: typing.Iterable[concurrent.futures.Future[typing.Any]], + timeout: float | None = None, + *, + bar: typing.Any = 'plain', + poll_interval: float = DEFAULT_POLL_INTERVAL, + **bar_kwargs: typing.Any, +) -> typing.Generator[concurrent.futures.Future[typing.Any], None, None]: + """`concurrent.futures.as_completed` with a progress bar. + + A superset of the stdlib function: same yield order and `timeout` + semantics, plus a bar counting completions (total inferred from the + futures). The caller owns the futures -- an early ``break`` or a + timeout never cancels them. + """ + futures_list: list[concurrent.futures.Future[typing.Any]] = list(futures) + display: _display.Display = _display.make_display( + bar, + total=len(futures_list), + poll_interval=poll_interval, + bar_kwargs=bar_kwargs, + ) + done: queue.SimpleQueue[concurrent.futures.Future[typing.Any]] = ( + queue.SimpleQueue() + ) + pending: set[concurrent.futures.Future[typing.Any]] = set(futures_list) + deadline: float | None = ( + None if timeout is None else time.monotonic() + timeout + ) + success: bool = False + for future in pending: + future.add_done_callback(done.put) + try: + display.start(len(futures_list)) + while pending: + if deadline is not None and time.monotonic() > deadline: + raise concurrent.futures.TimeoutError( + f'{len(pending)} (of {len(futures_list)}) futures ' + f'unfinished within timeout={timeout}' + ) + try: + future = done.get(timeout=poll_interval) + except queue.Empty: + display.tick() + continue + pending.discard(future) + display.advance() + yield future + success = True + finally: + display.finish(success=success) + + +def imap( + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> typing.Generator[typing.Any, None, None]: + """Lazily apply `fn` in parallel, yielding results in input order. + + The parallel counterpart of ``multiprocessing.Pool.imap``: same + ordering, same laziness, same results-only element shape. Results + completed out of order are held back until their turn; the held + set stays bounded by the submission window (`buffersize`). + + Closing the generator early (``break``) cancels unsubmitted work + and shuts down the run's executor; wrap in `contextlib.closing` + for deterministic cleanup. See `execute` for keywords. + """ + held: dict[int, typing.Any] = {} + next_index: int = 0 + for index, _args, _ok, value in execute(fn, iterables, **kwargs): + held[index] = value + while next_index in held: + yield held.pop(next_index) + next_index += 1 + + +def imap_unordered( + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> typing.Generator[tuple[typing.Any, typing.Any], None, None]: + """Lazily apply `fn` in parallel, yielding as tasks complete. + + Yields ``(item, result)`` pairs in *completion* order -- the pair + shape exists because completion order loses the input + correspondence (a deliberate deviation from + ``multiprocessing.Pool.imap_unordered``, which yields bare + results). With multiple iterables, ``item`` is the argument tuple. + + Closing the generator early cancels unsubmitted work; see `imap`. + """ + single: bool = len(iterables) == 1 + for _index, args, _ok, value in execute(fn, iterables, **kwargs): + yield _common.item_of(args, single), value + + +def map( # noqa: A001 - intentional builtin name, namespaced use only + fn: typing.Callable[..., typing.Any], + /, + *iterables: typing.Iterable[typing.Any], + **kwargs: typing.Any, +) -> list[typing.Any]: + """Apply `fn` to every zipped item in parallel; results in order. + + The parallel counterpart of the builtin ``map``: + ``progressbar.map(fn, items, workers=8)`` runs on a thread pool by + default, renders a progress bar, and returns the results in input + order once the batch completes. ``pool='process'`` switches to + processes, ``bar='multi'`` shows per-task sub-bars, and + ``on_error='return'`` swaps fail-fast for exceptions-in-place. See + `execute` for the full keyword reference. + """ + results: dict[int, typing.Any] = { + index: value + for index, _args, _ok, value in execute(fn, iterables, **kwargs) + } + return [results[index] for index in range(len(results))] diff --git a/tests/api_surface_snapshot.json b/tests/api_surface_snapshot.json index 86820c5..b254d0e 100644 --- a/tests/api_surface_snapshot.json +++ b/tests/api_surface_snapshot.json @@ -4,6 +4,7 @@ "AdaptiveETA": "class(exponential_smoothing=?, exponential_smoothing_factor=?, **kwargs)", "AdaptiveTransferSpeed": "class(**kwargs)", "AnimatedMarker": "class(markers=?, default=?, fill=?, marker_wrap=?, fill_wrap=?, **kwargs)", + "AsyncPool": "class(concurrency=?, **defaults)", "Bar": "class(marker=?, left=?, right=?, fill=?, fill_left=?, marker_wrap=?, **kwargs)", "BouncingBar": "class(marker=?, left=?, right=?, fill=?, fill_left=?, marker_wrap=?, **kwargs)", "Counter": "class(format=?, **kwargs)", @@ -26,8 +27,10 @@ "MultiProgressBar": "class(name, markers=?, **kwargs)", "MultiRangeBar": "class(name, markers, **kwargs)", "NullBar": "class(min_value=?, max_value=?, widgets=?, left_justify=?, initial_value=?, poll_interval=?, widget_kwargs=?, custom_len=?, max_error=?, prefix=?, suffix=?, variables=?, min_poll_interval=?, desc=?, total=?, unit=?, unit_scale=?, postfix=?, **kwargs)", + "ParallelFunction": "type-alias", "Percentage": "class(format=?, na=?, **kwargs)", "PercentageLabelBar": "class(format=?, na=?, **kwargs)", + "Pool": "class(workers=?, kind=?, *, executor=?, **defaults)", "Postfix": "class(name=?, prefix=?, separator=?, **kwargs)", "ProgressBar": "class(min_value=?, max_value=?, widgets=?, left_justify=?, initial_value=?, poll_interval=?, widget_kwargs=?, custom_len=?, max_error=?, prefix=?, suffix=?, variables=?, min_poll_interval=?, desc=?, total=?, unit=?, unit_scale=?, postfix=?, **kwargs)", "ReverseBar": "class(marker=?, left=?, right=?, fill=?, fill_left=?, **kwargs)", @@ -43,9 +46,14 @@ "VariableMixin": "class(name, **kwargs)", "__author__": "str", "__version__": "str", + "current_task_bar": "callable()", "len_color": "callable(value)", + "parallel": "callable(**config)", + "process_map": "callable(fn, *iterables, **kwargs)", "progressbar": "callable(iterator, min_value=?, max_value=?, widgets=?, prefix=?, suffix=?, fast=?, desc=?, total=?, unit=?, unit_scale=?, postfix=?, **kwargs)", - "streams": "StreamWrapper" + "starmap": "callable(fn, iterable, **kwargs)", + "streams": "StreamWrapper", + "thread_map": "callable(fn, *iterables, **kwargs)" }, "progressbar.algorithms": { "DoubleExponentialMovingAverage": "class(alpha=?)", @@ -66,11 +74,13 @@ "ResizableMixin": "class(term_width=?, **kwargs)", "StdRedirectMixin": "class(redirect_stderr=?, redirect_stdout=?, redirect_blank_line=?, **kwargs)", "T": "type-alias", + "TracebackType": "re-export", "ValueT": "type-alias", "annotations": "_Feature", "datetime": "re-export", "deepcopy": "callable(x, memo=?, _nil=?)", - "logger": "Logger" + "logger": "Logger", + "timedelta": "re-export" }, "progressbar.base": { "FalseMeta": "classsignature-unavailable", @@ -85,6 +95,7 @@ "COLOR_SUPPORT": "ColorSupport", "ColorSupport": "enum(NONE,XTERM,XTERM_256,XTERM_TRUECOLOR,WINDOWS)", "JUPYTER": "bool", + "TRUECOLOR_TERMS": "frozenset", "annotations": "_Feature", "env_flag": "callable(name, default=?)", "is_ansi_terminal": "callable(fd, is_terminal=?)", diff --git a/tests/test_init_exports.py b/tests/test_init_exports.py index 0f09096..6ad07b3 100644 --- a/tests/test_init_exports.py +++ b/tests/test_init_exports.py @@ -1,8 +1,11 @@ -"""Guard the three hand-synced export lists in ``progressbar/__init__.py``. - -``_NAME_TO_MODULE`` is the single source of truth for the lazily re-exported -public names. ``__all__`` and the ``TYPE_CHECKING`` import block must stay in -sync with it; these tests fail loudly if any of the three drift apart. +"""Guard the hand-synced export lists in ``progressbar/__init__.py``. + +``_NAME_TO_MODULE`` is the single source of truth for the lazily +re-exported public names; ``_NAMESPACE_ONLY`` holds the names that +resolve via the namespace but stay out of ``__all__`` (they would +shadow builtins or stdlib names under a star-import). ``__all__`` and +the ``TYPE_CHECKING`` import block must stay in sync with them; these +tests fail loudly if any of the lists drift apart. """ from __future__ import annotations @@ -15,6 +18,7 @@ # Alias (not a `from` import) so CodeQL doesn't flag `progressbar` as imported # with both `import` and `import from`. _NAME_TO_MODULE = progressbar._NAME_TO_MODULE +_NAMESPACE_ONLY = progressbar._NAMESPACE_ONLY #: Dunders that are eagerly imported (not part of ``_NAME_TO_MODULE``) but are #: still part of the public ``__all__``. @@ -28,18 +32,33 @@ def test_every_mapping_name_resolves() -> None: assert getattr(progressbar, name) is not None, name +def test_every_namespace_only_name_resolves() -> None: + for name in _NAMESPACE_ONLY: + assert getattr(progressbar, name) is not None, name + + def test_all_matches_mapping_plus_dunders() -> None: # ``__all__`` must contain exactly the mapping names plus the eager - # dunders. The concrete ordering is delegated to ruff's RUF022, so this - # compares contents rather than the exact list order. + # dunders -- and none of the namespace-only names, which would + # shadow builtins/stdlib under ``from progressbar import *``. The + # concrete ordering is delegated to ruff's RUF022, so this compares + # contents rather than the exact list order. assert set(progressbar.__all__) == set(_NAME_TO_MODULE) | _EAGER_DUNDERS +def test_namespace_only_stays_out_of_all() -> None: + assert not set(_NAMESPACE_ONLY) & set(progressbar.__all__) + + +def test_mappings_do_not_overlap() -> None: + assert not set(_NAMESPACE_ONLY) & set(_NAME_TO_MODULE) + + def test_all_has_no_duplicates() -> None: assert len(progressbar.__all__) == len(set(progressbar.__all__)) -def test_type_checking_block_imports_exactly_the_mapping() -> None: +def test_type_checking_block_imports_exactly_the_mappings() -> None: source: str = pathlib.Path(progressbar.__file__).read_text() tree: ast.Module = ast.parse(source) @@ -60,4 +79,4 @@ def test_type_checking_block_imports_exactly_the_mapping() -> None: for alias in stmt.names: imported.add(alias.asname or alias.name) - assert imported == set(_NAME_TO_MODULE) + assert imported == set(_NAME_TO_MODULE) | set(_NAMESPACE_ONLY) diff --git a/tests/test_parallel_aliases.py b/tests/test_parallel_aliases.py new file mode 100644 index 0000000..fc84bb2 --- /dev/null +++ b/tests/test_parallel_aliases.py @@ -0,0 +1,110 @@ +"""`starmap`, the tqdm-style aliases and progress-aware `as_completed`.""" + +from __future__ import annotations + +import concurrent.futures +import io +import operator +import time + +import pytest + +from progressbar._parallel import _sync + + +def _add(left: int, right: int) -> int: + return left + right + + +def _double(value: int) -> int: + return value * 2 + + +class TestStarmap: + def test_unpacks_argument_tuples(self) -> None: + assert _sync.starmap(operator.add, [(1, 2), (3, 4)], bar=False) == [ + 3, + 7, + ] + + def test_process_pool(self) -> None: + assert _sync.starmap( + _add, [(1, 2), (3, 4)], pool='process', workers=2, bar=False + ) == [3, 7] + + def test_empty(self) -> None: + assert _sync.starmap(operator.add, [], bar=False) == [] + + +class TestTqdmStyleAliases: + def test_thread_map(self) -> None: + assert _sync.thread_map(_double, range(5), bar=False) == [ + 0, + 2, + 4, + 6, + 8, + ] + + def test_process_map(self) -> None: + assert _sync.process_map(_double, range(5), workers=2, bar=False) == [ + 0, + 2, + 4, + 6, + 8, + ] + + def test_pool_kwarg_rejected(self) -> None: + with pytest.raises(TypeError, match='pool'): + _sync.thread_map(_double, range(3), pool='process', bar=False) + with pytest.raises(TypeError, match='pool'): + _sync.process_map(_double, range(3), pool='thread', bar=False) + + +class TestAsCompleted: + def _futures( + self, executor: concurrent.futures.ThreadPoolExecutor + ) -> list[concurrent.futures.Future[int]]: + return [executor.submit(_double, value) for value in range(5)] + + def test_yields_every_future_with_a_bar(self) -> None: + stream = io.StringIO() + with concurrent.futures.ThreadPoolExecutor(2) as executor: + futures = self._futures(executor) + seen = list(_sync.as_completed(futures, fd=stream)) + assert sorted(fut.result() for fut in seen) == [0, 2, 4, 6, 8] + assert '5' in stream.getvalue() + assert stream.getvalue().endswith('\n') + + @pytest.mark.no_freezegun + def test_timeout_raises(self) -> None: + with concurrent.futures.ThreadPoolExecutor(1) as executor: + blocker = executor.submit(time.sleep, 3) + with pytest.raises(concurrent.futures.TimeoutError): + list( + _sync.as_completed( + [blocker], + timeout=0.2, + poll_interval=0.05, + bar=False, + ) + ) + blocker.cancel() + + def test_early_break_leaves_futures_untouched(self) -> None: + with concurrent.futures.ThreadPoolExecutor(2) as executor: + futures = self._futures(executor) + for _future in _sync.as_completed(futures, bar=False): + break + # Not cancelled: the caller owns these futures. + assert sorted(fut.result() for fut in futures) == [ + 0, + 2, + 4, + 6, + 8, + ] + + def test_empty(self) -> None: + assert list(_sync.as_completed([], bar=False)) == [] diff --git a/tests/test_parallel_async.py b/tests/test_parallel_async.py new file mode 100644 index 0000000..32118ca --- /dev/null +++ b/tests/test_parallel_async.py @@ -0,0 +1,402 @@ +"""The asyncio engine: `amap` and its call-strategy handling.""" + +from __future__ import annotations + +import asyncio +import io +import operator +import typing + +import pytest + +from progressbar._parallel import _async + + +async def _async_double(value: int) -> int: + return value * 2 + + +def _sync_double(value: int) -> int: + return value * 2 + + +def _returns_awaitable(value: int) -> typing.Awaitable[int]: + return _async_double(value) + + +async def _boom_on_two(value: int) -> int: + if value == 2: + raise ValueError('boom') + return value + + +class TestAmap: + def test_async_fn_ordered(self) -> None: + async def _run() -> list[int]: + return await _async.amap(_async_double, range(10), bar=False) + + assert asyncio.run(_run()) == [value * 2 for value in range(10)] + + @pytest.mark.no_freezegun + def test_ordered_despite_scrambled_completion(self) -> None: + async def _staggered(value: int) -> int: + await asyncio.sleep((5 - value) * 0.02) + return value + + async def _run() -> list[int]: + return await _async.amap(_staggered, range(5), bar=False) + + assert asyncio.run(_run()) == list(range(5)) + + @pytest.mark.no_freezegun + def test_sync_fn_wrapped_in_thread(self) -> None: + async def _run() -> list[int]: + return await _async.amap(_sync_double, range(5), bar=False) + + assert asyncio.run(_run()) == [0, 2, 4, 6, 8] + + @pytest.mark.no_freezegun + def test_sync_fn_returning_awaitable_is_awaited(self) -> None: + async def _run() -> list[int]: + return await _async.amap(_returns_awaitable, range(4), bar=False) + + assert asyncio.run(_run()) == [0, 2, 4, 6] + + def test_multiple_iterables_zip(self) -> None: + async def _add(left: int, right: int) -> int: + return left + right + + async def _run() -> list[int]: + return await _async.amap(_add, [1, 2], [10, 20], bar=False) + + assert asyncio.run(_run()) == [11, 22] + + def test_empty(self) -> None: + async def _run() -> list[int]: + return await _async.amap(_async_double, [], bar=False) + + assert asyncio.run(_run()) == [] + + @pytest.mark.no_freezegun + def test_concurrency_capped(self) -> None: + running: list[int] = [0] + seen_max: list[int] = [0] + + async def _tracked(value: int) -> int: + running[0] += 1 + seen_max[0] = max(seen_max[0], running[0]) + await asyncio.sleep(0.02) + running[0] -= 1 + return value + + async def _run() -> list[int]: + return await _async.amap( + _tracked, range(10), concurrency=2, bar=False + ) + + assert asyncio.run(_run()) == list(range(10)) + assert seen_max[0] <= 2 + + def test_workers_alias(self) -> None: + async def _run() -> list[int]: + return await _async.amap( + _async_double, range(4), workers=2, bar=False + ) + + assert asyncio.run(_run()) == [0, 2, 4, 6] + + +class TestAmapErrors: + def test_fail_fast(self) -> None: + async def _run() -> list[int]: + return await _async.amap( + _boom_on_two, range(10), concurrency=1, bar=False + ) + + with pytest.raises(ValueError, match='boom'): + asyncio.run(_run()) + + def test_on_error_return(self) -> None: + async def _run() -> list[typing.Any]: + return await _async.amap( + _boom_on_two, range(5), on_error='return', bar=False + ) + + results: list[typing.Any] = asyncio.run(_run()) + assert results[1] == 1 + assert isinstance(results[2], ValueError) + assert results[4] == 4 + + @pytest.mark.no_freezegun + def test_timeout_cancels_cleanly(self) -> None: + async def _slow(value: int) -> int: + await asyncio.sleep(30) + return value # pragma: no cover - always cancelled + + async def _run() -> list[int]: + return await _async.amap( + _slow, + range(4), + timeout=0.2, + poll_interval=0.05, + bar=False, + ) + + with pytest.raises(asyncio.TimeoutError): + asyncio.run(_run()) + # asyncio.run closing the loop without warnings proves the + # outstanding tasks were cancelled and awaited. + + def test_invalid_on_error(self) -> None: + async def _run() -> list[int]: + return await _async.amap( + _async_double, range(3), on_error='ignore', bar=False + ) + + with pytest.raises(ValueError, match='on_error'): + asyncio.run(_run()) + + +class TestKeepAlive: + @pytest.mark.no_freezegun + def test_bar_ticks_during_long_task(self) -> None: + stream = io.StringIO() + + async def _slow(value: int) -> int: + await asyncio.sleep(0.4) + return value + + async def _run() -> list[int]: + return await _async.amap( + _slow, range(2), poll_interval=0.05, fd=stream + ) + + assert asyncio.run(_run()) == [0, 1] + # Multiple renders happened while the tasks slept: the output + # contains far more than the start + finish frames. Count + # rendered frames by their timer text -- the line separator + # depends on stream/tty detection. + assert stream.getvalue().count('Elapsed Time') > 3 + + +class TestAimap: + @pytest.mark.no_freezegun + def test_ordered_despite_scrambled_completion(self) -> None: + async def _staggered(value: int) -> int: + await asyncio.sleep((5 - value) * 0.02) + return value + + async def _run() -> list[int]: + return [ + value + async for value in _async.aimap( + _staggered, range(5), bar=False + ) + ] + + assert asyncio.run(_run()) == list(range(5)) + + @pytest.mark.no_freezegun + def test_early_break_with_aclosing(self) -> None: + import contextlib + + async def _run() -> list[int]: + collected: list[int] = [] + async with contextlib.aclosing( + _async.aimap( + _async_double, range(10), concurrency=2, bar=False + ) + ) as iterator: + async for value in iterator: + collected.append(value) + if len(collected) == 2: + break + return collected + + assert asyncio.run(_run()) == [0, 2] + + +class TestAimapUnordered: + @pytest.mark.no_freezegun + def test_yields_pairs_in_completion_order(self) -> None: + async def _staggered(value: int) -> int: + await asyncio.sleep((5 - value) * 0.02) + return value + + async def _run() -> list[tuple[int, int]]: + return [ + pair + async for pair in _async.aimap_unordered( + _staggered, range(5), bar=False + ) + ] + + pairs: list[tuple[int, int]] = asyncio.run(_run()) + assert sorted(pairs) == [(value, value) for value in range(5)] + assert pairs[0] == (4, 4) + + def test_multi_iterable_pairs_use_args_tuple(self) -> None: + async def _add(left: int, right: int) -> int: + return left + right + + async def _run() -> list[tuple[typing.Any, int]]: + return [ + pair + async for pair in _async.aimap_unordered( + _add, [1, 2], [10, 20], bar=False + ) + ] + + assert sorted(asyncio.run(_run())) == [ + ((1, 10), 11), + ((2, 20), 22), + ] + + +class TestGather: + def test_ordered_results(self) -> None: + async def _run() -> list[int]: + return await _async.gather( + _async_double(1), + _async_double(2), + _async_double(3), + bar=False, + ) + + assert asyncio.run(_run()) == [2, 4, 6] + + def test_empty_returns_empty_list(self) -> None: + async def _run() -> list[typing.Any]: + return await _async.gather() + + assert asyncio.run(_run()) == [] + + def test_return_exceptions(self) -> None: + async def _run() -> list[typing.Any]: + return await _async.gather( + _async_double(1), + _boom_on_two(2), + _async_double(3), + return_exceptions=True, + bar=False, + ) + + results: list[typing.Any] = asyncio.run(_run()) + assert results[0] == 2 + assert isinstance(results[1], ValueError) + assert results[2] == 6 + + def test_fail_fast_by_default(self) -> None: + async def _run() -> list[typing.Any]: + return await _async.gather( + _async_double(1), _boom_on_two(2), bar=False + ) + + with pytest.raises(ValueError, match='boom'): + asyncio.run(_run()) + + +class TestAsyncPool: + @pytest.mark.no_freezegun + def test_bounds_concurrency(self) -> None: + running: list[int] = [0] + seen_max: list[int] = [0] + + async def _tracked(value: int) -> int: + running[0] += 1 + seen_max[0] = max(seen_max[0], running[0]) + await asyncio.sleep(0.02) + running[0] -= 1 + return value + + async def _run() -> list[int]: + async with _async.AsyncPool(2, bar=False) as pool: + return await pool.map(_tracked, range(8)) + + assert asyncio.run(_run()) == list(range(8)) + assert seen_max[0] <= 2 + + def test_defaults_merge_and_override(self) -> None: + async def _run() -> list[typing.Any]: + async with _async.AsyncPool(2, bar=False) as pool: + return await pool.map( + _boom_on_two, range(4), on_error='return' + ) + + results: list[typing.Any] = asyncio.run(_run()) + assert isinstance(results[2], ValueError) + + def test_imap_methods(self) -> None: + async def _run() -> tuple[list[int], list[tuple[int, int]]]: + async with _async.AsyncPool(2, bar=False) as pool: + ordered: list[int] = [ + value async for value in pool.imap(_async_double, range(3)) + ] + pairs: list[tuple[int, int]] = sorted( + [ + pair + async for pair in pool.imap_unordered( + _async_double, range(3) + ) + ] + ) + return ordered, pairs + + ordered, pairs = asyncio.run(_run()) + assert ordered == [0, 2, 4] + assert pairs == [(0, 0), (1, 2), (2, 4)] + + +class TestMultiBarMode: + def test_async_workers_see_their_task_bar(self) -> None: + from progressbar._parallel import _common + + seen: list[bool] = [] + + async def _check(value: int) -> int: + seen.append(_common.current_task_bar() is not None) + return value + + async def _run() -> list[int]: + return await _async.amap( + _check, range(3), bar='multi', fd=io.StringIO() + ) + + assert asyncio.run(_run()) == [0, 1, 2] + assert seen == [True, True, True] + + +class TestExternalCancellation: + def test_self_cancelling_task_surfaces(self) -> None: + async def _self_cancel(value: int) -> int: + if value == 1: + task = asyncio.current_task() + assert task is not None + task.cancel() + await asyncio.sleep(1) + return value + + async def _run() -> list[int]: + return await _async.amap( + _self_cancel, range(3), concurrency=1, bar=False + ) + + # A cancellation this run did not initiate must surface, never + # silently drop the item. + with pytest.raises(asyncio.CancelledError): + asyncio.run(_run()) + + +class TestCallStrategy: + def test_detects_coroutine_function(self) -> None: + assert _async._call_strategy(_async_double) == 'async' + + def test_detects_partial_of_coroutine_function(self) -> None: + import functools + + partial = functools.partial(_async_double) + assert _async._call_strategy(partial) == 'async' + + def test_sync_fallback(self) -> None: + assert _async._call_strategy(_sync_double) == 'sync' + assert _async._call_strategy(operator.add) == 'sync' diff --git a/tests/test_parallel_common.py b/tests/test_parallel_common.py new file mode 100644 index 0000000..e0d1fb8 --- /dev/null +++ b/tests/test_parallel_common.py @@ -0,0 +1,156 @@ +"""Tests for the shared plumbing behind the parallel execution verbs.""" + +from __future__ import annotations + +import typing + +import pytest + +import progressbar +from progressbar import base +from progressbar._parallel import _common + + +def _boom_on_two(value: int) -> int: + """Module-level worker: raises on input 2, else returns the input.""" + if value == 2: + raise ValueError('boom') + return value + + +def _unsized() -> typing.Iterator[int]: + yield from (1, 2, 3) + + +class TestDetectTotal: + def test_sized(self) -> None: + assert _common.detect_total(([1, 2, 3],)) == 3 + + def test_length_hint(self) -> None: + # A list_iterator has no __len__ but supports __length_hint__. + assert _common.detect_total((iter([1, 2]),)) == 2 + + def test_unsized(self) -> None: + assert _common.detect_total((_unsized(),)) is base.UnknownLength + + def test_multiple_iterables_take_min(self) -> None: + assert _common.detect_total(([1, 2, 3], [1, 2])) == 2 + + def test_mixed_sized_and_unsized(self) -> None: + assert _common.detect_total(([1, 2], _unsized())) is base.UnknownLength + + def test_no_iterables(self) -> None: + assert _common.detect_total(()) == 0 + + +class TestValidateBarKwargs: + def test_known_keys_pass(self) -> None: + _common.validate_bar_kwargs({'prefix': 'x', 'max_value': 10}) + + def test_typo_raises(self) -> None: + with pytest.raises(TypeError, match='worker'): + _common.validate_bar_kwargs({'worker': 8}) + + def test_known_names_present(self) -> None: + names: frozenset[str] = _common.known_bar_kwargs( + progressbar.ProgressBar + ) + assert {'poll_interval', 'max_value', 'widgets'} <= names + assert 'self' not in names + + +class TestResolveWorkers: + def test_explicit(self) -> None: + assert _common.resolve_workers(7, 'thread') == 7 + + def test_thread_default_capped(self) -> None: + assert 1 <= _common.resolve_workers(None, 'thread') <= 32 + + def test_process_default(self) -> None: + assert _common.resolve_workers(None, 'process') >= 1 + + +class TestBufferAndChunkDefaults: + def test_default_buffersize(self) -> None: + assert _common.default_buffersize(8) == 32 + assert _common.default_buffersize(1) == 16 + + def test_auto_chunksize_unknown_total(self) -> None: + assert _common.auto_chunksize(base.UnknownLength, 8) == 1 + + def test_auto_chunksize_scales(self) -> None: + assert _common.auto_chunksize(100_000, 8) == 781 + + def test_auto_chunksize_small_batch(self) -> None: + assert _common.auto_chunksize(10, 8) == 1 + + def test_auto_chunksize_capped(self) -> None: + assert _common.auto_chunksize(10_000_000, 1) == 1_000 + + +class TestIterChunks: + def test_single_iterable(self) -> None: + chunks: list[list[tuple[int, ...]]] = list( + _common.iter_chunks(([1, 2, 3, 4, 5],), 2) + ) + assert chunks == [[(1,), (2,)], [(3,), (4,)], [(5,)]] + + def test_zips_multiple_iterables(self) -> None: + chunks = list(_common.iter_chunks(([1, 2], ['a', 'b']), 10)) + assert chunks == [[(1, 'a'), (2, 'b')]] + + def test_lazy(self) -> None: + # Consuming one chunk must not consume the whole source. + source: typing.Iterator[int] = iter(range(100)) + first: list[tuple[int, ...]] = next(_common.iter_chunks((source,), 3)) + assert first == [(0,), (1,), (2,)] + assert next(source) < 10 + + +class TestItemOf: + def test_single(self) -> None: + assert _common.item_of((42,), single=True) == 42 + + def test_multiple(self) -> None: + assert _common.item_of((1, 2), single=False) == (1, 2) + + +class TestRunChunk: + def test_catch_returns_outcomes(self) -> None: + outcomes: list[tuple[bool, typing.Any]] = _common.run_chunk( + _boom_on_two, [(1,), (2,), (3,)], catch=True + ) + assert [ok for ok, _ in outcomes] == [True, False, True] + assert outcomes[0] == (True, 1) + assert isinstance(outcomes[1][1], ValueError) + + def test_no_catch_aborts_chunk(self) -> None: + with pytest.raises(ValueError, match='boom'): + _common.run_chunk(_boom_on_two, [(1,), (2,), (3,)], catch=False) + + def test_catch_lets_keyboard_interrupt_escape(self) -> None: + def _interrupt(_: int) -> None: + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + _common.run_chunk(_interrupt, [(1,)], catch=True) + + +class TestCurrentTaskBar: + def test_default_is_none(self) -> None: + assert _common.current_task_bar() is None + + def test_with_task_bar_binds_and_restores(self) -> None: + marker: progressbar.ProgressBar = progressbar.ProgressBar(max_value=1) + seen: list[progressbar.ProgressBar | None] = [] + + def _inner() -> str: + seen.append(_common.current_task_bar()) + return 'done' + + wrapped: typing.Callable[[], str] = _common.with_task_bar( + marker, _inner + ) + assert wrapped() == 'done' + assert seen == [marker] + assert _common.current_task_bar() is None diff --git a/tests/test_parallel_decorator.py b/tests/test_parallel_decorator.py new file mode 100644 index 0000000..fc5027f --- /dev/null +++ b/tests/test_parallel_decorator.py @@ -0,0 +1,86 @@ +"""The `@parallel` decorator: batch verbs attached to plain functions.""" + +from __future__ import annotations + +import asyncio +import pickle + +import pytest + +from progressbar._parallel import _decorator + + +@_decorator.parallel(workers=2, bar=False) +def _double(value: int) -> int: + return value * 2 + + +@_decorator.parallel(workers=2, bar=False) +def _square(value: int) -> int: + return value * value + + +class TestDecoratedFunction: + def test_direct_call_unchanged(self) -> None: + assert _double(21) == 42 + + def test_map_with_config_defaults(self) -> None: + assert _double.map(range(5)) == [0, 2, 4, 6, 8] + + def test_per_call_override_beats_config(self) -> None: + # The config sets bar=False; overriding on_error per call works + # alongside it. + assert _double.map(range(3), workers=1) == [0, 2, 4] + + def test_imap_variants(self) -> None: + assert list(_double.imap(range(3))) == [0, 2, 4] + pairs = sorted(_double.imap_unordered(range(3))) + assert pairs == [(0, 0), (1, 2), (2, 4)] + + def test_starmap(self) -> None: + assert _double.starmap([(1,), (2,)]) == [2, 4] + + def test_amap(self) -> None: + async def _run() -> list[int]: + return await _double.amap(range(3)) + + assert asyncio.run(_run()) == [0, 2, 4] + + def test_async_iterators(self) -> None: + async def _run() -> tuple[list[int], list[tuple[int, int]]]: + ordered = [value async for value in _double.aimap(range(3))] + pairs = sorted( + [pair async for pair in _double.aimap_unordered(range(3))] + ) + return ordered, pairs + + ordered, pairs = asyncio.run(_run()) + assert ordered == [0, 2, 4] + assert pairs == [(0, 0), (1, 2), (2, 4)] + + +class TestPicklability: + def test_pickle_round_trip_preserves_identity(self) -> None: + # Attributes attach to the original function object, so pickle + # by qualified name still works -- required for pool='process'. + assert pickle.loads(pickle.dumps(_square)) is _square + + def test_process_pool_map(self) -> None: + assert _square.map(range(4), pool='process') == [0, 1, 4, 9] + + +class TestDecorationErrors: + def test_lambda_rejected(self) -> None: + with pytest.raises(TypeError, match='named function'): + _decorator.parallel()(lambda value: value) + + def test_bound_method_rejected(self) -> None: + class _Thing: + def method(self) -> None: ... + + with pytest.raises(TypeError, match='named function'): + _decorator.parallel()(_Thing().method) + + def test_non_callable_rejected(self) -> None: + with pytest.raises(TypeError, match='named function'): + _decorator.parallel()(42) # type: ignore[arg-type] diff --git a/tests/test_parallel_display.py b/tests/test_parallel_display.py new file mode 100644 index 0000000..e51ccfa --- /dev/null +++ b/tests/test_parallel_display.py @@ -0,0 +1,194 @@ +"""Tests for the parallel display layer (plain/null/instance modes).""" + +from __future__ import annotations + +import io +import time +import typing + +import pytest + +import progressbar +from progressbar._parallel import _display + + +def _make( + mode: typing.Any, + total: typing.Any = 3, + poll_interval: float = 0.1, + **bar_kwargs: typing.Any, +) -> _display.Display: + return _display.make_display( + mode, total=total, poll_interval=poll_interval, bar_kwargs=bar_kwargs + ) + + +class TestPlainDisplay: + def test_advances_to_completion(self) -> None: + stream = io.StringIO() + display: _display.Display = _make('plain', fd=stream) + display.start(3) + display.advance() + display.advance(2) + display.finish() + assert '3' in stream.getvalue() + assert stream.getvalue().endswith('\n') + + def test_tick_redraws_without_progress(self) -> None: + # Keep-alive: with zero completions, a tick after poll_interval + # must produce a new render (pins the rev-2 spec blocker). + stream = io.StringIO() + display = _make('plain', fd=stream) + display.start(3) + before: int = len(stream.getvalue()) + time.sleep(0.5) # advances the frozen clock via autouse fixture + display.tick() + assert len(stream.getvalue()) > before + display.finish() + + def test_task_hooks_are_noops(self) -> None: + display = _make('plain', fd=io.StringIO()) + display.start(1) + assert display.task_started(1, 'label') is None + display.task_finished(1, ok=True) + display.finish() + + def test_failure_finish_does_not_force_max(self) -> None: + stream = io.StringIO() + display = _make('plain', fd=stream) + display.start(3) + time.sleep(0.01) # let the frozen clock pass the update gate + display.advance() + display.finish(success=False) + # A dirty finish keeps the value at 1 instead of jumping to 3. + last_line: str = stream.getvalue().rstrip('\n').rsplit('\r', 1)[-1] + assert '1 of 3' in last_line + assert stream.getvalue().endswith('\n') + + +class TestNullDisplay: + def test_everything_is_a_noop(self) -> None: + display = _make(False) + display.start(1) + display.advance() + display.tick() + assert display.task_started(1, 'x') is None + display.task_finished(1, ok=False) + display.finish(success=False) + + +class TestInstanceDisplay: + def test_wraps_and_drives_unstarted_bar(self) -> None: + stream = io.StringIO() + bar: progressbar.ProgressBar = progressbar.ProgressBar( + max_value=2, fd=stream, poll_interval=0.1 + ) + display = _make(bar) + display.start(2) + display.advance(2) + display.finish() + assert bar.finished() + + def test_prestarted_bar_is_not_finished_by_us(self) -> None: + stream = io.StringIO() + bar = progressbar.ProgressBar(max_value=2, fd=stream) + bar.start() + display = _make(bar) + display.start(2) + display.advance() + display.finish() + assert not bar.finished() + + +class TestMultiDisplay: + def _multi(self, total: int = 3) -> _display.MultiDisplay: + display = _display.make_display( + 'multi', + total=total, + poll_interval=0.05, + bar_kwargs={'fd': io.StringIO()}, + ) + assert isinstance(display, _display.MultiDisplay) + return display + + def test_satisfies_protocol(self) -> None: + display = self._multi() + assert isinstance(display, _display.Display) + display.finish() + + def test_per_task_bars_appear_and_disappear(self) -> None: + display = self._multi() + display.start(3) + task_bar = display.task_started(1, 'item-a') + assert task_bar is not None + assert '1: item-a' in display.multibar + display.task_finished(1, ok=True) + assert '1: item-a' not in display.multibar + display.finish() + + def test_duplicate_labels_get_distinct_keys(self) -> None: + display = self._multi() + display.start(3) + display.task_started(1, 'same') + display.task_started(2, 'same') + assert '1: same' in display.multibar + assert '2: same' in display.multibar + display.task_finished(1, ok=True) + display.task_finished(2, ok=False) + display.finish() + + def test_overall_bar_counts_completions(self) -> None: + display = self._multi() + display.start(3) + display.advance() + display.advance(2) + assert display.multibar['Total'].value == 3 + display.finish() + + def test_render_thread_stopped_after_finish(self) -> None: + display = self._multi() + display.start(3) + display.advance(3) + display.finish() + assert display.multibar._thread is None # noqa: SLF001 + + def test_task_finished_with_unknown_seq_is_a_noop(self) -> None: + display = self._multi() + display.start(1) + display.task_finished(99, ok=True) + display.finish() + + def test_finish_removes_live_task_bars(self) -> None: + display = self._multi() + display.start(3) + display.task_started(1, 'still-running') + display.task_started(2, 'also-running') + display.finish(success=False) + assert '1: still-running' not in display.multibar + assert '2: also-running' not in display.multibar + + def test_adopts_existing_multibar_without_stopping_it(self) -> None: + multibar = progressbar.MultiBar(fd=io.StringIO()) + multibar.start() + display = _display.make_display( + multibar, total=2, poll_interval=0.05, bar_kwargs={} + ) + display.start(2) + display.advance(2) + display.finish() + assert multibar._thread is not None # noqa: SLF001 + multibar.stop(timeout=5) + + +class TestMakeDisplay: + def test_unknown_mode_raises(self) -> None: + with pytest.raises(TypeError, match='bogus'): + _make('bogus') + + def test_bar_kwargs_reach_the_bar(self) -> None: + stream = io.StringIO() + display = _make('plain', fd=stream, prefix='PFX ') + display.start(3) + display.advance() + display.finish() + assert 'PFX' in stream.getvalue() diff --git a/tests/test_parallel_errors.py b/tests/test_parallel_errors.py new file mode 100644 index 0000000..6e93ff1 --- /dev/null +++ b/tests/test_parallel_errors.py @@ -0,0 +1,126 @@ +"""Pin the sync engine's error, timeout and interrupt contracts.""" + +from __future__ import annotations + +import concurrent.futures +import io +import threading +import time + +import pytest + +from progressbar._parallel import _sync + +_executed: set[int] = set() +_executed_lock: threading.Lock = threading.Lock() + + +def _boom(value: int) -> int: + if value == 3: + raise ValueError('boom') + return value * 2 + + +def _record_and_boom(value: int) -> int: + with _executed_lock: + _executed.add(value) + if value == 0: + raise ValueError('early boom') + return value + + +def _raise_interrupt(value: int) -> int: + if value == 1: + raise KeyboardInterrupt + return value + + +def _sleep_long(value: int) -> int: + # Long enough to trip the 0.3s deadline, short enough that the two + # straggler worker threads drain quickly in the background. + time.sleep(3) + return value + + +class TestFailFast: + def test_raises_original_exception(self) -> None: + with pytest.raises(ValueError, match='boom'): + _sync.map(_boom, range(10), workers=2, bar=False) + + def test_cancels_pending_work(self) -> None: + _executed.clear() + with pytest.raises(ValueError, match='early boom'): + _sync.map( + _record_and_boom, + range(50), + workers=1, + buffersize=2, + bar=False, + ) + # workers=1 runs items sequentially; item 0 fails, so at most + # the already-submitted window (2 chunks) ever executed. + assert len(_executed) <= 3 + + def test_keyboard_interrupt_propagates(self) -> None: + with pytest.raises(KeyboardInterrupt): + _sync.map(_raise_interrupt, range(10), workers=1, bar=False) + + def test_keyboard_interrupt_propagates_with_on_error_return( + self, + ) -> None: + # `on_error='return'` must never swallow an interrupt. + with pytest.raises(KeyboardInterrupt): + _sync.map( + _raise_interrupt, + range(10), + workers=1, + on_error='return', + bar=False, + ) + + +class TestOnErrorReturn: + def test_exceptions_in_place(self) -> None: + results = _sync.map(_boom, range(5), on_error='return', bar=False) + assert results[0] == 0 + assert results[2] == 4 + assert isinstance(results[3], ValueError) + assert results[4] == 8 + + def test_invalid_on_error_rejected(self) -> None: + with pytest.raises(ValueError, match='on_error'): + _sync.map(_boom, range(3), on_error='ignore', bar=False) + + +class TestTimeout: + @pytest.mark.no_freezegun + def test_timeout_raises_and_cancels(self) -> None: + start: float = time.monotonic() + with pytest.raises(concurrent.futures.TimeoutError, match='timeout'): + _sync.map( + _sleep_long, + range(4), + workers=2, + timeout=0.3, + poll_interval=0.05, + bar=False, + ) + # The engine must give up at the deadline instead of waiting + # for the 3-second workers: running tasks are documented as + # uncancellable but the shutdown must not block on them. + assert time.monotonic() - start < 2 + + +class TestBarFinalState: + def test_error_finishes_bar_on_own_line(self) -> None: + stream = io.StringIO() + with pytest.raises(ValueError, match='boom'): + _sync.map(_boom, range(10), workers=1, fd=stream) + assert stream.getvalue().endswith('\n') + + def test_error_does_not_jump_to_full(self) -> None: + stream = io.StringIO() + with pytest.raises(ValueError, match='boom'): + _sync.map(_boom, range(10), workers=1, buffersize=1, fd=stream) + final_line: str = stream.getvalue().rstrip('\n').rsplit('\r', 1)[-1] + assert '10 of 10' not in final_line diff --git a/tests/test_parallel_imap.py b/tests/test_parallel_imap.py new file mode 100644 index 0000000..c020e37 --- /dev/null +++ b/tests/test_parallel_imap.py @@ -0,0 +1,121 @@ +"""Lazy ordered `imap` and completion-order `imap_unordered`.""" + +from __future__ import annotations + +import contextlib +import operator +import threading +import time +import typing + +import pytest + +from progressbar._parallel import _sync + +_executed: set[int] = set() +_executed_lock: threading.Lock = threading.Lock() + + +def _double(value: int) -> int: + return value * 2 + + +def _record(value: int) -> int: + with _executed_lock: + _executed.add(value) + return value + + +def _sleep_inverse(value: int) -> int: + # Later items finish first, scrambling completion order. + time.sleep((5 - value) * 0.03) + return value + + +def _boom_on_two(value: int) -> int: + if value == 2: + raise ValueError('boom') + return value + + +class TestImap: + def test_yields_results_in_input_order(self) -> None: + assert list(_sync.imap(_double, range(10), bar=False)) == [ + value * 2 for value in range(10) + ] + + @pytest.mark.no_freezegun + def test_ordered_despite_scrambled_completion(self) -> None: + assert list( + _sync.imap(_sleep_inverse, range(5), workers=5, bar=False) + ) == list(range(5)) + + def test_lazy(self) -> None: + iterator: typing.Generator[typing.Any, None, None] = _sync.imap( + _double, range(10), workers=1, bar=False + ) + assert next(iterator) == 0 + iterator.close() + + def test_on_error_return_yields_exceptions_in_place(self) -> None: + results: list[typing.Any] = list( + _sync.imap(_boom_on_two, range(4), on_error='return', bar=False) + ) + assert results[0] == 0 + assert isinstance(results[2], ValueError) + assert results[3] == 3 + + def test_on_error_raise_raises_at_iteration(self) -> None: + iterator = _sync.imap(_boom_on_two, range(4), workers=1, bar=False) + with pytest.raises(ValueError, match='boom'): + list(iterator) + + +class TestImapUnordered: + @pytest.mark.no_freezegun + def test_yields_pairs_in_completion_order(self) -> None: + pairs: list[tuple[int, int]] = list( + _sync.imap_unordered( + _sleep_inverse, range(5), workers=5, bar=False + ) + ) + assert sorted(pairs) == [(value, value) for value in range(5)] + # The fastest item (highest input) completes and is seen first. + assert pairs[0] == (4, 4) + + def test_multi_iterable_pairs_use_args_tuple(self) -> None: + pairs = list( + _sync.imap_unordered( + operator.add, [1, 2], [10, 20], workers=1, bar=False + ) + ) + assert sorted(pairs) == [((1, 10), 11), ((2, 20), 22)] + + def test_on_error_return_pairs_exceptions(self) -> None: + pairs = dict( + _sync.imap_unordered( + _boom_on_two, range(4), on_error='return', bar=False + ) + ) + assert isinstance(pairs[2], ValueError) + assert pairs[3] == 3 + + +class TestGeneratorCleanup: + def test_early_break_stops_submission(self) -> None: + _executed.clear() + for item, _result in _sync.imap_unordered( + _record, range(100), workers=1, buffersize=2, bar=False + ): + if item >= 1: + break + # workers=1, window=2: only the in-window items ever ran; the + # other ~97 must have been cancelled by the generator close. + time.sleep(0.2) # let any straggler drain before asserting + assert len(_executed) <= 5 + + def test_contextlib_closing_recipe(self) -> None: + with contextlib.closing( + _sync.imap(_double, range(10), workers=1, bar=False) + ) as iterator: + assert next(iterator) == 0 diff --git a/tests/test_parallel_map.py b/tests/test_parallel_map.py new file mode 100644 index 0000000..51381e7 --- /dev/null +++ b/tests/test_parallel_map.py @@ -0,0 +1,157 @@ +"""Tests for the sync engine's ordered `map` over threads.""" + +from __future__ import annotations + +import io +import operator +import threading +import time + +import pytest + +from progressbar._parallel import _sync + + +def _double(value: int) -> int: + return value * 2 + + +def _sleep_inverse(value: int) -> int: + # Later items finish first, scrambling completion order. + time.sleep((5 - value) * 0.02) + return value + + +async def _async_double(value: int) -> int: # pragma: no cover - never runs + return value * 2 + + +class TestMap: + def test_ordered_results(self) -> None: + assert _sync.map(_double, range(10), bar=False) == [ + value * 2 for value in range(10) + ] + + def test_multiple_iterables_zip(self) -> None: + assert _sync.map(operator.add, [1, 2], [10, 20], bar=False) == [ + 11, + 22, + ] + + def test_empty_input(self) -> None: + assert _sync.map(_double, [], bar=False) == [] + + @pytest.mark.no_freezegun + def test_order_preserved_under_scrambled_completion(self) -> None: + assert _sync.map( + _sleep_inverse, range(5), workers=5, bar=False + ) == list(range(5)) + + def test_single_worker(self) -> None: + assert _sync.map(_double, range(5), workers=1, bar=False) == [ + 0, + 2, + 4, + 6, + 8, + ] + + def test_small_buffersize_completes(self) -> None: + assert _sync.map( + _double, range(20), workers=2, buffersize=2, bar=False + ) == [value * 2 for value in range(20)] + + def test_generator_input(self) -> None: + assert _sync.map( + _double, (value for value in range(5)), bar=False + ) == [0, 2, 4, 6, 8] + + def test_coroutine_function_rejected(self) -> None: + with pytest.raises(TypeError, match='amap'): + _sync.map(_async_double, range(3), bar=False) + + def test_bar_false_produces_no_output( + self, capsys: pytest.CaptureFixture[str] + ) -> None: + _sync.map(_double, range(3), bar=False) + captured = capsys.readouterr() + assert captured.out == '' + assert captured.err == '' + + def test_typo_kwarg_raises(self) -> None: + with pytest.raises(TypeError, match='worker'): + _sync.map(_double, range(3), worker=8) + + def test_bar_renders_progress(self) -> None: + stream = io.StringIO() + _sync.map(_double, range(3), fd=stream) + assert '3' in stream.getvalue() + assert stream.getvalue().endswith('\n') + + def test_runs_in_worker_threads(self) -> None: + main_thread: threading.Thread = threading.current_thread() + seen: set[str] = set() + + def _record(value: int) -> int: + seen.add(threading.current_thread().name) + return value + + _sync.map(_record, range(10), workers=2, bar=False) + assert main_thread.name not in seen + + +class TestMultiBarMode: + def test_workers_see_their_task_bar(self) -> None: + from progressbar._parallel import _common + + seen: list[bool] = [] + + def _check(value: int) -> int: + seen.append(_common.current_task_bar() is not None) + return value + + _sync.map(_check, range(4), workers=2, bar='multi', fd=io.StringIO()) + assert seen == [True, True, True, True] + + def test_plain_mode_has_no_task_bar(self) -> None: + from progressbar._parallel import _common + + seen: list[bool] = [] + + def _check(value: int) -> int: + seen.append(_common.current_task_bar() is None) + return value + + _sync.map(_check, range(2), workers=2, bar=False) + assert seen == [True, True] + + +class TestResolveExecutor: + def test_thread_pool_created_and_owned(self) -> None: + executor, owned, workers = _sync.resolve_executor( + 'thread', + 3, + initializer=None, + initargs=(), + mp_context=None, + max_tasks_per_child=None, + thread_name_prefix='', + ) + try: + assert owned is True + assert workers == 3 + assert executor.submit(_double, 2).result() == 4 + finally: + executor.shutdown() + + def test_unknown_pool_raises(self) -> None: + with pytest.raises(ValueError, match='bogus'): + _sync.resolve_executor( + 'bogus', + None, + initializer=None, + initargs=(), + mp_context=None, + max_tasks_per_child=None, + thread_name_prefix='', + ) diff --git a/tests/test_parallel_pool.py b/tests/test_parallel_pool.py new file mode 100644 index 0000000..bcc2d8f --- /dev/null +++ b/tests/test_parallel_pool.py @@ -0,0 +1,113 @@ +"""The reusable `Pool` layer over the sync verbs.""" + +from __future__ import annotations + +import concurrent.futures + +import pytest + +from progressbar._parallel import _sync + + +def _double(value: int) -> int: + return value * 2 + + +def _boom(value: int) -> int: + raise ValueError('boom') + + +class TestPoolLifecycle: + def test_lazy_executor(self) -> None: + pool = _sync.Pool(2) + assert pool._executor is None # noqa: SLF001 + pool.shutdown() + + def test_executor_reused_across_calls(self) -> None: + with _sync.Pool(2) as pool: + first = pool.executor + pool.map(_double, range(3), bar=False) + pool.map(_double, range(3), bar=False) + assert pool.executor is first + + def test_context_manager_shuts_down(self) -> None: + with _sync.Pool(2) as pool: + pool.map(_double, range(3), bar=False) + executor = pool.executor + with pytest.raises(RuntimeError): + executor.submit(_double, 1) + + def test_adopted_executor_not_shut_down(self) -> None: + with concurrent.futures.ThreadPoolExecutor(2) as executor: + with _sync.Pool(executor=executor) as pool: + assert pool.map(_double, range(3), bar=False) == [0, 2, 4] + # Leaving the Pool context must not kill the adopted + # executor -- the caller owns it. + assert executor.submit(_double, 2).result() == 4 + + def test_invalid_kind_rejected_eagerly(self) -> None: + with pytest.raises(ValueError, match='bogus'): + _sync.Pool(2, 'bogus') + + def test_workers_with_executor_rejected(self) -> None: + with ( + concurrent.futures.ThreadPoolExecutor(2) as executor, + pytest.raises(ValueError, match='executor'), + ): + _sync.Pool(2, executor=executor) + + +class TestPoolVerbs: + def test_map(self) -> None: + with _sync.Pool(2) as pool: + assert pool.map(_double, range(5), bar=False) == [ + 0, + 2, + 4, + 6, + 8, + ] + + def test_imap(self) -> None: + with _sync.Pool(2) as pool: + assert list(pool.imap(_double, range(5), bar=False)) == [ + 0, + 2, + 4, + 6, + 8, + ] + + def test_imap_unordered(self) -> None: + with _sync.Pool(2) as pool: + pairs = sorted(pool.imap_unordered(_double, range(3), bar=False)) + assert pairs == [(0, 0), (1, 2), (2, 4)] + + def test_starmap(self) -> None: + with _sync.Pool(2) as pool: + assert pool.starmap(_double_args, [(1,), (2,)], bar=False) == [ + 2, + 4, + ] + + +def _double_args(value: int) -> int: + return value * 2 + + +class TestPoolDefaults: + def test_constructor_defaults_apply(self) -> None: + with _sync.Pool(2, bar=False, on_error='return') as pool: + results = pool.map(_boom, range(2)) + assert all(isinstance(result, ValueError) for result in results) + + def test_per_call_override_beats_default(self) -> None: + with ( + _sync.Pool(2, bar=False, on_error='return') as pool, + pytest.raises(ValueError, match='boom'), + ): + pool.map(_boom, range(2), on_error='raise') + + def test_process_kind(self) -> None: + with _sync.Pool(2, 'process') as pool: + assert pool.map(_double, range(4), bar=False) == [0, 2, 4, 6] diff --git a/tests/test_parallel_process.py b/tests/test_parallel_process.py new file mode 100644 index 0000000..cabfc87 --- /dev/null +++ b/tests/test_parallel_process.py @@ -0,0 +1,229 @@ +"""Process/interpreter pool support, chunking and executor passthrough.""" + +from __future__ import annotations + +import concurrent.futures +import multiprocessing +import sys +import typing + +import pytest + +from progressbar._parallel import ( + _common, + _sync, +) + +# Process tests pay real spawn cost; keep batches small. +_INIT_VALUE: int = 0 + + +def _square(value: int) -> int: + return value * value + + +def _boom_on_two(value: int) -> int: + if value == 2: + raise ValueError('boom') + return value + + +def _init_worker(value: int) -> None: + global _INIT_VALUE # noqa: PLW0603 - the per-worker setup contract + _INIT_VALUE = value + + +def _read_init(_: int) -> int: + return _INIT_VALUE + + +class TestProcessPool: + def test_ordered_results(self) -> None: + assert _sync.map(_square, range(12), pool='process', bar=False) == [ + value * value for value in range(12) + ] + + def test_explicit_chunksize(self) -> None: + assert _sync.map( + _square, range(10), pool='process', chunksize=3, bar=False + ) == [value * value for value in range(10)] + + def test_auto_chunksize_engaged( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + calls: list[tuple[typing.Any, int]] = [] + original: typing.Callable[[typing.Any, int], int] = ( + _common.auto_chunksize + ) + + def _spy(total: typing.Any, workers: int) -> int: + calls.append((total, workers)) + return original(total, workers) + + monkeypatch.setattr(_sync._common, 'auto_chunksize', _spy) + _sync.map(_square, range(4), pool='process', workers=2, bar=False) + assert calls == [(4, 2)] + + def test_auto_chunksize_not_used_for_threads( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def _fail(total: typing.Any, workers: int) -> int: + raise AssertionError('auto_chunksize must not run for threads') + + monkeypatch.setattr(_sync._common, 'auto_chunksize', _fail) + _sync.map(_square, range(4), pool='thread', bar=False) + + def test_initializer_reaches_workers(self) -> None: + results = _sync.map( + _read_init, + range(4), + pool='process', + workers=2, + initializer=_init_worker, + initargs=(42,), + bar=False, + ) + assert results == [42, 42, 42, 42] + + def test_mp_context(self) -> None: + context = multiprocessing.get_context('spawn') + assert _sync.map( + _square, + range(4), + pool='process', + workers=2, + mp_context=context, + bar=False, + ) == [0, 1, 4, 9] + + @pytest.mark.skipif( + sys.version_info >= (3, 11), + reason='the ValueError applies before Python 3.11', + ) + def test_max_tasks_per_child_rejected_on_310( + self, + ) -> None: # pragma: no cover + with pytest.raises(ValueError, match=r'3\.11'): + _sync.map( + _square, + range(2), + pool='process', + max_tasks_per_child=2, + bar=False, + ) + + @pytest.mark.skipif( + sys.version_info < (3, 11), + reason='max_tasks_per_child needs Python 3.11+', + ) + def test_max_tasks_per_child(self) -> None: + assert _sync.map( + _square, + range(4), + pool='process', + workers=2, + max_tasks_per_child=2, + bar=False, + ) == [0, 1, 4, 9] + + def test_chunked_on_error_return_keeps_partial_chunk(self) -> None: + results = _sync.map( + _boom_on_two, + range(6), + pool='process', + chunksize=3, + on_error='return', + bar=False, + ) + # Item 2 fails inside the first chunk; 0, 1 and the whole + # second chunk survive (per-item catch, no data loss). + assert results[0] == 0 + assert results[1] == 1 + assert isinstance(results[2], ValueError) + assert results[3:] == [3, 4, 5] + + def test_chunked_on_error_raise(self) -> None: + with pytest.raises(ValueError, match='boom'): + _sync.map( + _boom_on_two, + range(6), + pool='process', + chunksize=3, + bar=False, + ) + + +class TestInterpreterPool: + @pytest.mark.skipif( + sys.version_info < (3, 14), + reason='InterpreterPoolExecutor needs Python 3.14+', + ) + def test_interpreter_pool_runs(self) -> None: # pragma: no cover + # Subinterpreter workers cannot import test modules, so the + # worker callable must come from an importable module -- the + # builtin `abs` qualifies everywhere. + assert _sync.map( + abs, [-1, -2, -3], pool='interpreter', workers=2, bar=False + ) == [1, 2, 3] + + @pytest.mark.skipif( + sys.version_info >= (3, 14), + reason='the ValueError applies before Python 3.14', + ) + def test_interpreter_pool_rejected_before_314(self) -> None: + with pytest.raises(ValueError, match=r'3\.14'): + _sync.map(_square, range(4), pool='interpreter', bar=False) + + +class TestExecutorInstance: + def test_used_but_not_shut_down(self) -> None: + with concurrent.futures.ThreadPoolExecutor(2) as executor: + assert _sync.map(_square, range(5), pool=executor, bar=False) == [ + 0, + 1, + 4, + 9, + 16, + ] + # Still usable afterwards: the engine must not shut it down. + assert executor.submit(_square, 3).result() == 9 + + def test_construction_kwargs_rejected(self) -> None: + with ( + concurrent.futures.ThreadPoolExecutor(2) as executor, + pytest.raises(ValueError, match='initializer'), + ): + _sync.map( + _square, + range(3), + pool=executor, + initializer=_init_worker, + initargs=(1,), + bar=False, + ) + + +class TestPoolValidation: + def test_unknown_pool_string(self) -> None: + with pytest.raises(ValueError, match='bogus'): + _sync.map(_square, range(3), pool='bogus', bar=False) + + def test_thread_pool_rejects_process_options(self) -> None: + with pytest.raises(ValueError, match='process pools'): + _sync.map( + _square, + range(3), + pool='thread', + mp_context=multiprocessing.get_context('spawn'), + bar=False, + ) + + def test_process_pool_rejects_thread_options(self) -> None: + with pytest.raises(ValueError, match='thread pools'): + _sync.map( + _square, + range(3), + pool='process', + thread_name_prefix='x', + bar=False, + ) diff --git a/tests/test_parallel_shell.py b/tests/test_parallel_shell.py new file mode 100644 index 0000000..100f662 --- /dev/null +++ b/tests/test_parallel_shell.py @@ -0,0 +1,141 @@ +"""The `run()` shell helper and its brace-safe template engine.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import typing + +import pytest + +from progressbar._parallel import ( + _shell, + _sync, +) + +#: A tiny portable command: exit with the given code. +_EXIT: list[str] = [sys.executable, '-c', 'import sys; sys.exit(0)'] + + +class TestBuildArgv: + def test_str_template_placeholder(self) -> None: + assert _shell.build_argv('gzip -k {}', 'a.txt', shell=False) == [ + 'gzip', + '-k', + 'a.txt', + ] + + def test_item_with_spaces_stays_one_argument(self) -> None: + assert _shell.build_argv('gzip -k {}', 'a file.txt', shell=False) == [ + 'gzip', + '-k', + 'a file.txt', + ] + + def test_literal_braces_survive(self) -> None: + # str.format would blow up on awk's braces; replacement of the + # exact placeholder token must not. + argv = _shell.build_argv( + "awk '{print $1}' {}", 'data.csv', shell=False + ) + if os.name == 'nt': + # Windows uses non-POSIX splitting (so backslash paths + # survive), which also preserves quote characters. + assert argv == ['awk', "'{print $1}'", 'data.csv'] + else: + assert argv == ['awk', '{print $1}', 'data.csv'] + + def test_item_placeholder_synonym(self) -> None: + assert _shell.build_argv( + 'convert {item} out-{item}.png', 'x', shell=False + ) == ['convert', 'x', 'out-x.png'] + + def test_no_placeholder_appends_item(self) -> None: + assert _shell.build_argv('gzip -k', 'a.txt', shell=False) == [ + 'gzip', + '-k', + 'a.txt', + ] + + def test_list_template(self) -> None: + assert _shell.build_argv( + ['ffmpeg', '-i', '{}', '{}.mp4'], 'in.avi', shell=False + ) == ['ffmpeg', '-i', 'in.avi', 'in.avi.mp4'] + + def test_list_without_placeholder_appends(self) -> None: + assert _shell.build_argv(['echo'], 'hi', shell=False) == [ + 'echo', + 'hi', + ] + + def test_callable_template(self) -> None: + assert _shell.build_argv( + lambda item: ['echo', str(item).upper()], 'hi', shell=False + ) == ['echo', 'HI'] + + def test_shell_string_replacement(self) -> None: + command = _shell.build_argv( + 'gzip -k {} > /dev/null', 'a.txt', shell=True + ) + assert command == 'gzip -k a.txt > /dev/null' + + def test_shell_string_appends_quoted(self) -> None: + command = _shell.build_argv('gzip -k', 'a file.txt', shell=True) + assert command == "gzip -k 'a file.txt'" + + +class TestRun: + @pytest.mark.no_freezegun + def test_runs_commands_and_returns_completed_processes(self) -> None: + results: list[subprocess.CompletedProcess[str]] = _shell.run( + [sys.executable, '-c', 'print({})'], + [1, 2, 3], + workers=2, + bar=False, + ) + assert [proc.stdout.strip() for proc in results] == ['1', '2', '3'] + assert all(proc.returncode == 0 for proc in results) + + @pytest.mark.no_freezegun + def test_check_raises_called_process_error(self) -> None: + with pytest.raises(subprocess.CalledProcessError): + _shell.run( + [sys.executable, '-c', 'import sys; sys.exit({})'], + [0, 1], + workers=1, + bar=False, + ) + + @pytest.mark.no_freezegun + def test_check_false_returns_failures(self) -> None: + results = _shell.run( + [sys.executable, '-c', 'import sys; sys.exit({})'], + [0, 1], + check=False, + workers=1, + bar=False, + ) + assert [proc.returncode for proc in results] == [0, 1] + + @pytest.mark.no_freezegun + def test_on_error_return_embeds_the_error(self) -> None: + results: list[typing.Any] = _shell.run( + [sys.executable, '-c', 'import sys; sys.exit({})'], + [0, 1], + on_error='return', + workers=1, + bar=False, + ) + assert results[0].returncode == 0 + assert isinstance(results[1], subprocess.CalledProcessError) + + def test_pool_kwarg_rejected(self) -> None: + with pytest.raises(TypeError, match=r'Pool\.run'): + _shell.run(_EXIT, [1], pool='process', bar=False) + + @pytest.mark.no_freezegun + def test_pool_run_method(self) -> None: + with _sync.Pool(2) as pool: + results = pool.run(_EXIT, range(2), bar=False) + assert all(proc.returncode == 0 for proc in results) diff --git a/tests/test_perf_budget.py b/tests/test_perf_budget.py index eb65cac..6ed750a 100644 --- a/tests/test_perf_budget.py +++ b/tests/test_perf_budget.py @@ -1,6 +1,8 @@ from __future__ import annotations +import concurrent.futures import io +import subprocess import sys import timeit @@ -82,3 +84,73 @@ def test_iterator_overhead_budget() -> None: f'({clock_ns:.1f} ns) - likely a regression to per-iteration ' f'clock reads' ) + + +def test_import_stays_lazy() -> None: + # Deterministic, not timing-based: `import progressbar` must load + # nothing beyond the package itself and the version module. Anything + # else appearing here means an eager import crept into __init__.py + # and the ~1.6 ms import time regressed for every consumer. + probe: str = ( + 'import sys, progressbar; ' + "print(','.join(sorted(" + "m for m in sys.modules if m.startswith('progressbar'))))" + ) + out: str = subprocess.run( + [sys.executable, '-c', probe], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + assert out == 'progressbar,progressbar.__about__', ( + f'import progressbar eagerly loaded: {out}' + ) + + +def _noop(value: int) -> int: + return value + + +def _parallel_map_us_per_item(n: int) -> float: + """Per-item wall cost of `progressbar.map` on a shared thread pool. + + The executor is created outside the measurement, so this isolates + the coordinator itself: chunking, submission windowing, the + done-queue and result assembly. + """ + import progressbar + + with concurrent.futures.ThreadPoolExecutor(4) as executor: + # Warm-up so pool spin-up and lazy imports land outside timing. + progressbar.map(_noop, range(64), pool=executor, bar=False) + elapsed: float = min( + timeit.timeit( + lambda: progressbar.map( + _noop, range(n), pool=executor, bar=False + ), + number=1, + ) + for _ in range(3) + ) + return elapsed / n * 1e6 + + +@pytest.mark.no_freezegun +def test_parallel_map_overhead_scales_linearly() -> None: + # Measure both before any early return so every line runs under + # coverage (same pattern as the iterator budget above). + small: float = _parallel_map_us_per_item(1_000) + large: float = _parallel_map_us_per_item(10_000) + if _coverage_active(): + return + # Machine-independent guard for the done-queue design: per-item cost + # must stay flat as the batch grows. The rejected coordinator design + # (re-registering a waiter on every pending future each poll) scales + # with batch size and blows past this immediately at 10x the items. + # A 3x ceiling tolerates noisy runners without letting an O(n) tick + # regime back in. + assert large < 3 * small, ( + f'parallel map per-item cost grew from {small:.2f} us at 1k items ' + f'to {large:.2f} us at 10k items - the coordinator is no longer ' + f'O(1) per completion' + )