mqdm 2.0.0__py3-none-any.whl

mqdm/__init__.py

@@ -0,0 +1,96 @@
+from .runtime import Runtime, _current_runtime, configure
+
+
+def sustain():
+    """Keep the live progress display alive across this block.
+
+    A sequence of separate bars normally renders one at a time — each freezes
+    into the scrollback as the next begins. Inside ``sustain()`` they stack and
+    stay visible together as one growing panel, with ``print``/logging streaming
+    above them. Nestable.
+    """
+    return _current_runtime().sustain()
+
+
+def print(*args, **kw):
+    """Print above active progress bars."""
+    return _current_runtime().print(*args, **kw)
+
+
+def get(i=-1):
+    """Get an mqdm instance."""
+    return _current_runtime().get_instance(i)
+
+
+def set_description(desc, i=-1):
+    """Set the description of the last progress bar."""
+    return get(i).set_description(desc)
+
+
+def set(i=-1, **kw):
+    """Set the last progress bar."""
+    return get(i).set(**kw)
+
+
+def update(n=1, i=-1, **kw):
+    """Update the last progress bar."""
+    return get(i).update(n, **kw)
+
+
+def pause(paused=True):
+    """Pause the progress bars. Useful for opening an interactive shell or printing stack traces."""
+    return _current_runtime().pause(paused)
+
+
+# ---------------------------------- Logging --------------------------------- #
+
+
+def install_logging(logger=None, *, level=None, capture_warnings='process', markup=True, formatter=None):
+    """Install an MQDMHandler on a logger for a runtime."""
+    return _current_runtime().install_logging(
+        logger=logger,
+        level=level,
+        capture_warnings=capture_warnings,
+        markup=markup,
+        formatter=formatter,
+    )
+
+
+def uninstall_logging(*, logger=None):
+    """Remove an MQDMHandler from a logger for a runtime."""
+    return _current_runtime().uninstall_logging(logger=logger)
+
+# ----------------------------------- Utils ---------------------------------- #
+
+from . import utils
+from .parallel.executor import T_POOL_MODE, get_executor, Initializer
+from .utils import args, fn, fopen, ratelimit
+from .utils._logging import MQDMHandler
+
+# ---------------------------------- Events ---------------------------------- #
+
+from . import events
+
+# ----------------------------------- Core ----------------------------------- #
+
+from .bar import mqdm
+from .parallel.pool import ipool, pool, PoolError, Result
+from .parallel.apool import aipool, apool
+
+# ----------------------------- Development utils ---------------------------- #
+
+from .utils._dev import bp, embed, iex, profile, timeit
+
+__all__ = [
+    'mqdm',
+    'pool',
+    'ipool',
+    'apool',
+    'aipool',
+    'print',
+    'sustain',
+    'pause',
+    'args',
+    'PoolError',
+    'configure',
+]

mqdm/backend/__init__.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .protocols import (
+    ProgressBackend,
+    ProgressBackendFactory,
+    ProxyConvertibleBackend,
+    RichTaskState,
+    TaskState,
+)
+from .protocols import ProgressBackendFactory as _PF  # noqa: F401 — re-export
+
+
+
+def create_backend(
+    *,
+    runtime: Any,
+    **kw: Any,
+) -> ProgressBackend:
+    from . import rich
+
+    return rich.Progress(**kw)
+

mqdm/backend/protocols.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Protocol, TypedDict, runtime_checkable
+
+
+class TaskState(TypedDict, total=False):
+    """Minimum detached task snapshot contract shared across backends.
+
+    This is the state mqdm itself needs to reattach a logical task after a bar
+    is temporarily detached from a live backend.
+    """
+
+    id: int
+    description: str
+    total: float | None
+    completed: float
+    visible: bool
+    fields: dict[str, Any]
+    start_time: float | None
+
+
+class RichTaskState(TaskState, total=False):
+    """Rich-specific snapshot fields used for restore and progress snapshots."""
+
+    stop_time: float | None
+    finished_time: float | None
+    finished_speed: float | None
+    _progress: list[tuple[float, float]] | None
+
+
+#: Signature of a callable that creates a :class:`ProgressBackend`.
+#:
+#: ``runtime`` receives the calling :class:`mqdm.Runtime`.  ``columns``
+#: is an optional column-tuple that overrides the backend's own defaults.
+#: Extra keyword arguments are forwarded to the backend constructor.
+ProgressBackendFactory = Callable[..., "ProgressBackend"]
+
+
+class ProgressBackend(Protocol):
+    """Core runtime/backend contract used by mqdm bars.
+
+    This deliberately excludes optional capabilities such as process-mode
+    promotion or backend-specific snapshot extensions.
+    """
+
+    multiprocess: bool
+
+    # ---------------------------------- Control --------------------------------- #
+
+    def start(self) -> None: ...
+    def stop(self) -> None: ...
+    def refresh(self) -> None: ...
+
+    # ---------------------------------- Output ---------------------------------- #
+
+    def write(self, *args: Any, **kw: Any) -> Any: ...
+
+    # ------------------------------ Task Management ----------------------------- #
+
+    def add_task(self, **task_kw: Any) -> int: ...
+    def try_update(self, task_id: int, **task_update: Any) -> None: ...
+    def dump_task(self, task_id: int) -> TaskState | None: ...
+    def load_task(self, task: TaskState, start: bool = True) -> None: ...
+    def pop_task(self, task_id: int, remove: bool | None = None) -> TaskState | None: ...
+
+
+@runtime_checkable
+class ProxyConvertibleBackend(Protocol):
+    """Optional capability for backends that can promote themselves for IPC."""
+
+    def convert_proxy(self, command_dispatch=None) -> ProgressBackend: ...

mqdm/backend/rich.py

@@ -0,0 +1,302 @@
+from collections import deque
+import dataclasses
+from dataclasses import dataclass
+from typing import Any
+from functools import wraps
+from rich.console import Console
+from rich import progress
+import mqdm as M
+
+from .protocols import RichTaskState, TaskState
+from ..utils.proxy import TransportProxy, proxymethod
+
+
+class _DiscardFile:
+    """A write-only file-like object that discards all data.
+
+    Used as the Rich console file when ``silent=True`` so the in-memory
+    buffer doesn't grow without bounds across refreshes.
+    """
+
+    encoding = "utf-8"
+
+    def write(self, s: str) -> int:
+        return len(s)
+
+    def flush(self) -> None:
+        pass
+
+    def getvalue(self) -> str:
+        return ""
+
+    def isatty(self) -> bool:
+        return False
+
+
+class Task(progress.Task):
+    @property
+    def speed(self):
+        """Optional[float]: Get the estimated speed in steps per second."""
+        if self.start_time is None:
+            return None
+        with self._lock:
+            progress = self._progress
+            if progress:
+                total_time = progress[-1].timestamp - self.start_time
+            elif self.finished:
+                total_time = self.finished_time
+            else:
+                total_time = self.elapsed
+            return self.completed / total_time if total_time else total_time
+
+
+@dataclass
+class TaskSnapshot:
+    id: int
+    description: str
+    total: float | None
+    completed: int
+    visible: bool
+    fields: dict
+    start_time: float | None = None
+    stop_time: float | None = None
+    finished_time: float | None = None
+    finished_speed: float | None = None
+    _progress: list[tuple[float, float]] | None = None
+
+    @classmethod
+    def from_task(cls, task: progress.Task) -> 'TaskSnapshot':
+        data = {k.name: getattr(task, k.name) for k in dataclasses.fields(task) if not k.name.startswith('_')}
+        progress_samples = []
+        if task._progress:
+            progress_samples.append((task._progress[-1].timestamp, sum(s.completed for s in task._progress)))
+        return cls(**data, _progress=progress_samples)
+
+    def to_dict(self) -> RichTaskState:
+        return dataclasses.asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: TaskState) -> 'TaskSnapshot':
+        return cls(**data)
+
+
+class Progress(progress.Progress):
+    multiprocess = False
+
+    # ---------------------------------------------------------------------------- #
+    #                           Construction / Bootstrap                           #
+    #  these hooks preserve enough init state to restore detached tasks and keep   #
+    #  process-safe proxy promotion independent of Rich's in-process-only state.   #
+    # ---------------------------------------------------------------------------- #
+
+    def __init__(self, *, columns=None, _tasks=None, _task_index=None, _pause_event=None, silent=False, **kw):
+        if columns is None:
+            columns = self.default_progress_columns()
+        self._init_options = dict(kw)
+        if silent:
+            self._init_options['silent'] = True
+            kw.setdefault('console', Console(file=_DiscardFile(), force_terminal=True))
+        super().__init__(*(columns or ()), **kw)
+
+        if _tasks is not None:
+            self._tasks = {task_id: self._load_task(TaskSnapshot.from_dict(task)) for task_id, task in _tasks.items()}
+        self._task_index = progress.TaskID(_task_index or 0)
+
+    @classmethod
+    def default_progress_columns(cls) -> tuple[object, ...]:
+        """Return the default Rich column layout for mqdm progress displays.
+
+        Per-task fields such as ``bytes`` are resolved by the individual column
+        implementations, so the shared layout does not vary by runtime options.
+        """
+        from . import rich_progress_columns
+
+        return (
+            "[progress.description]{task.description}",
+            rich_progress_columns.TwoToneColumn(bar_width=None),
+            "[progress.percentage]{task.percentage:>3.0f}%",
+            rich_progress_columns.MofNColumn(),
+            rich_progress_columns.SpeedColumn(),
+            rich_progress_columns.TimeElapsedColumn(compact=True),
+            progress.TimeRemainingColumn(compact=True),
+            progress.SpinnerColumn(),
+        )
+
+    # -------------------------------- Rich API --------------------------------- #
+
+    def write(self, *args, **kw):
+        """Print above the live progress display.
+
+        Uses this Progress's own console so the output interleaves with the
+        live region. When called through a process-safe proxy, the write is
+        routed into the process that owns the real ``Progress`` (and its Live).
+        """
+        self.console.print(*args, **kw)
+
+    @wraps(progress.Progress.update)
+    def update(self, task_id, **kw):
+        if 'description' in kw and kw['description'] is None:  # ignore None descriptions
+            kw.pop('description')
+        return super().update(task_id, **kw)
+
+    @wraps(progress.Progress.update)
+    def try_update(self, task_id, **kw):
+        try:
+            return self.update(task_id, **kw)
+        except KeyError as e:
+            pass
+
+    update_ = try_update
+
+    # ---------------------------------------------------------------------------- #
+    #                           Task Lifecycle / Mutation                          #
+    #  these overrides funnel task creation and lifecycle through mqdm-owned       #
+    #  helpers so local and proxied bars can share the same task state model.      #
+    # ---------------------------------------------------------------------------- #
+
+    def _start_task(self, task):
+        if task.start_time is None:
+            task.start_time = self.get_time()
+
+    def _new_task(self,
+                  description: str='',
+                  start: bool = True,
+                  total: float|None = None,
+                  completed: int = 0,
+                  visible: bool = True,
+                  **fields):
+        task = Task(
+            self._task_index,
+            description,
+            total,
+            completed,
+            visible=visible,
+            fields=fields,
+            _get_time=self.get_time,
+            _lock=self._lock,
+        )
+        start and self._start_task(task)
+        self._task_index = progress.TaskID(int(self._task_index) + 1)
+        return task
+
+    def add_task(
+        self,
+        description: str='',
+        start: bool = True,
+        total: float|None = None,
+        completed: int = 0,
+        visible: bool = True,
+        **fields,
+    ) -> progress.TaskID:
+        with self._lock:
+            task = self._new_task(description or '', start, total, completed, visible, **fields)
+            self._tasks[task.id] = task
+        self.refresh()
+        return task.id
+
+    def start_task(self, task_id: progress.TaskID) -> None:
+        with self._lock:
+            self._start_task(self._tasks[task_id])
+
+    def new_task(self,
+                  description: str='',
+                  start: bool = True,
+                  total: float|None = None,
+                  completed: int = 0,
+                  visible: bool = True,
+                  **fields):
+        with self._lock:
+            return self._dump_task(self._new_task(description or '', start, total, completed, visible, **fields)).to_dict()
+
+    # ---------------------------------------------------------------------------- #
+    #                        Task Snapshot / Serialization                         #
+    #  these helpers reduce Rich tasks to transportable snapshots so state can be  #
+    #  exported, restored, and forwarded across process boundaries.                #
+    # ---------------------------------------------------------------------------- #
+
+    def dump_tasks(self) -> dict[int, RichTaskState]:
+        with self._lock:
+            return {task_id: self._dump_task(self._tasks[task_id]).to_dict() for task_id in self._tasks}
+
+    def dump_task(self, task_id) -> RichTaskState:
+        with self._lock:
+            return self._dump_task(self._tasks[task_id]).to_dict()
+
+    def _dump_task(self, task):
+        return TaskSnapshot.from_task(task)
+
+    def _load_task(self, snapshot: TaskSnapshot, start=True):
+        data = snapshot.to_dict()
+        _progress = data.pop('_progress') or []
+        start_time = data.pop('start_time', None)
+        stop_time = data.pop('stop_time', None)
+        task = Task(_get_time=self.get_time, _lock=self._lock, **data)
+        task.start_time = start_time
+        task.stop_time = stop_time
+        start and self._start_task(task)
+        if _progress:
+            task._progress = deque([progress.ProgressSample(*s) for s in _progress], maxlen=1000)
+        return task
+
+    def load_task(self, task: TaskState, start=True):
+        with self._lock:
+            task = self._load_task(TaskSnapshot.from_dict(task), start=start)
+            self._tasks[task.id] = task
+            if task.id >= self._task_index:
+                self._task_index = progress.TaskID(task.id+1)
+        self.refresh()
+
+    def pop_task(self, task_id, remove=None):
+        """Close a task and return its serialized data."""
+        try:
+            self.stop_task(task_id)
+            data = self.dump_task(task_id)
+            if remove is None:
+                remove = self._tasks[task_id].fields.get('transient', False)
+            if remove:
+                self.remove_task(task_id)
+            return data
+        except KeyError as e:
+            pass
+
+    # ---------------------------------------------------------------------------- #
+    #                        Multiprocessing Backend Upgrade                       #
+    #  this section upgrades a local Rich-backed Progress into a queue-backed      #
+    #  owner-rendered proxy for process workers.                                   #
+    # ---------------------------------------------------------------------------- #
+
+    def convert_proxy(self, command_dispatch=None) -> 'ProgressProxy':
+        """Convert to a multiprocessing-safe proxy object."""
+        return ProgressProxy.from_target(self, command_dispatch=command_dispatch)
+
+
+class ProgressProxy(TransportProxy[Progress]):
+    multiprocess = True
+
+    start = proxymethod(Progress.start, expect_reply=False, owner_only=True)
+    stop = proxymethod(Progress.stop, expect_reply=False, owner_only=True)
+    refresh = proxymethod(Progress.refresh, expect_reply=False, owner_only=True)
+
+    write = proxymethod(Progress.write, expect_reply=False)
+
+    add_task = proxymethod(Progress.add_task)
+    try_update = proxymethod(Progress.try_update, expect_reply=False)
+    dump_task = proxymethod(Progress.dump_task)
+    load_task = proxymethod(Progress.load_task, expect_reply=False)
+    pop_task = proxymethod(Progress.pop_task)
+    
+    # Unused methods (to deprecate) - not defined on ProgressBackend
+    
+    new_task = proxymethod(Progress.new_task)
+    update = proxymethod(Progress.update, expect_reply=False)
+    update_ = try_update
+    start_task = proxymethod(Progress.start_task, expect_reply=False)
+    stop_task = proxymethod(Progress.stop_task, expect_reply=False)
+    remove_task = proxymethod(Progress.remove_task, expect_reply=False)
+    dump_tasks = proxymethod(Progress.dump_tasks)
+
+    def __rich_console__(self, console, options):
+        target = self.target
+        if target is None:
+            raise RuntimeError("ProgressProxy can only render in the owner process.")
+        yield target.get_renderable()

mqdm/backend/rich_progress_columns.py

@@ -0,0 +1,158 @@
+import time
+from rich import progress
+from rich.segment import Segment
+from rich.style import Style
+from rich.progress_bar import ProgressBar
+
+
+class TwoToneBar(ProgressBar):
+    """A progress bar with a lighter fill for work that's started but not done.
+
+    Renders three zones — completed, started-but-pending, and idle — so you can
+    see how far the workers have run ahead of what has actually finished. Reads
+    the pending count from ``started`` (defaults to ``completed``, i.e. a plain
+    bar) and clamps it to ``completed <= started <= total``.
+
+    The started zone defaults to a dimmed shade of the complete colour (same hue,
+    lower intensity). That reads as "same work, in progress" and — because dim is
+    an SGR attribute rather than a colour — stays distinct from the grey
+    background even on 16-colour terminals, where a lighter grey would collapse
+    into it. Pass ``started_style`` to override.
+    """
+    def __init__(self, *args, started: float = 0, started_style: str | None = None, **kw):
+        super().__init__(*args, **kw)
+        self.started = started
+        self.started_style = started_style
+
+    def _resolve_started_style(self, console) -> Style:
+        if self.started_style is not None:
+            return console.get_style(self.started_style)
+        return console.get_style(self.complete_style) + Style(dim=True)
+
+    def __rich_console__(self, console, options):
+        width = min(self.width or options.max_width, options.max_width)
+        # Defer to the stock bar for indeterminate / pulsing tasks.
+        if self.total is None or self.pulse:
+            yield from super().__rich_console__(console, options)
+            return
+
+        bar = "-" if (options.legacy_windows or options.ascii_only) else "━"
+        done = max(0.0, min(self.total, self.completed))
+        started = max(done, min(self.total, self.started))
+        done_cells = int(round(width * done / self.total)) if self.total else 0
+        started_cells = int(round(width * started / self.total)) if self.total else 0
+
+        finished = self.total and self.completed >= self.total
+        zones = (
+            (done_cells, console.get_style(self.finished_style if finished else self.complete_style)),
+            (started_cells - done_cells, self._resolve_started_style(console)),
+            (width - started_cells, console.get_style(self.style)),
+        )
+        for count, style in zones:
+            if count > 0:
+                yield Segment(bar * count, style)
+
+
+class TwoToneColumn(progress.BarColumn):
+    """``BarColumn`` that shades a task's ``started`` field as a second tone."""
+    def __init__(self, *args, started_style: str | None = None, **kw):
+        self.started_style = started_style
+        super().__init__(*args, **kw)
+
+    def render(self, task) -> ProgressBar:
+        return TwoToneBar(
+            total=max(0, task.total) if task.total is not None else None,
+            completed=max(0, task.completed),
+            started=max(0, task.fields.get("started", task.completed)),
+            width=None if self.bar_width is None else max(1, self.bar_width),
+            pulse=not task.started,
+            animation_time=task.get_time(),
+            style=self.style,
+            complete_style=self.complete_style,
+            finished_style=self.finished_style,
+            pulse_style=self.pulse_style,
+            started_style=self.started_style,
+        )
+
+
+class MofNColumn(progress.DownloadColumn):
+    '''A progress column that shows the current vs. total count of items.'''
+    def __init__(self, bytes=False, separator="/", **kw):
+        self.bytes = bytes
+        self.separator = separator
+        super().__init__(**kw)
+
+    def render(self, task):
+        bytes = task.fields.get("bytes", self.bytes)
+        if bytes:
+            return super().render(task)
+        total = f'{int(task.total):,}' if task.total is not None else "?"
+        return progress.Text(
+            f"{int(task.completed):,d}{self.separator}{total}",
+            style="progress.download",
+            justify='right'
+        )
+
+
+class SpeedColumn(progress.TransferSpeedColumn):
+    """Renders human readable transfer speed."""
+    def __init__(self, bytes=False, unit_scale=1, **kw):
+        self.bytes = bytes
+        self.unit_scale = unit_scale
+        super().__init__(**kw)
+
+    def render(self, task):
+        """Show data transfer speed."""
+        bytes = task.fields.get("bytes", self.bytes)
+        if bytes:
+            return super().render(task)
+        speed = task.finished_speed or task.speed
+        if speed is None:
+            return progress.Text("", style="progress.data.speed")
+        end = '/s'
+        if 0 < speed < 1:
+            speed = 1 / speed
+            speed, suffix = time_units(speed)
+            end = '/x'
+            return progress.Text(f"{speed:.1f}{suffix}{end}", justify='right', style="progress.data.speed")
+        unit, suffix = progress.filesize.pick_unit_and_suffix(
+            int(speed), 
+            ["x", "K", "M", "B", "T"], 
+            # ["", "×10³", "×10⁶", "×10⁹", "×10¹²"], 
+            1000)
+        return progress.Text(f"{speed/unit:.1f}{suffix}{end}", justify='right', style="progress.data.speed")
+
+def time_units(seconds):
+    for d, u in [(60, 's'), (60, 'm'), (24, 'h'), (365, 'd')]:
+        if seconds < d:
+            return seconds, u
+        seconds /= d
+    return seconds, 'y'
+
+
+class TimeElapsedColumn(progress.TimeRemainingColumn):
+    """Renders time elapsed."""
+    def __init__(self, compact: bool = False, **kw):
+        self._compact = compact
+        super().__init__(**kw)
+
+    def render(self, task):
+        """Show time elapsed."""
+        elapsed = task.finished_time if task.finished else task.elapsed
+        if elapsed is None:
+            return progress.Text("--:--" if self._compact else "-:--:--", style="progress.elapsed")
+        minutes, seconds = divmod(int(elapsed), 60)
+        hours, minutes = divmod(minutes, 60)
+        return progress.Text(
+            (f"{hours:d}:" if hours or not self._compact else "") + 
+            f"{minutes:02d}:{seconds:02d}", style="progress.elapsed")
+
+
+
+# class LogBarColumn(progress.BarColumn):
+#     def render(self, task):
+#         return progress.Group(
+#             super().render(task),
+#             progress.Text(task.description, style="progress.description"),
+#         )
+