File size: 1,539 Bytes
399944f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
from __future__ import annotations

"""
Progress and timing helpers.

This module provides small utilities to improve the user experience of long-running
pipeline stages by emitting:
- stage start/end messages
- elapsed time
- optional spinner UI while blocking library calls run

We intentionally keep this lightweight and dependency-minimal (Rich only).
"""

from contextlib import contextmanager
from time import perf_counter
from typing import Iterator, Optional

from rich.console import Console
from rich.status import Status


@contextmanager
def stage_status(
    title: str,
    *,
    console: Optional[Console] = None,
    spinner: str = "dots",
) -> Iterator[None]:
    """
    Context manager that shows a Rich spinner and prints elapsed time.

    Use this for stages where work happens inside a single blocking call
    (e.g., Docling loader.load()) and you cannot expose per-item progress.

    Parameters
    ----------
    title:
        Human readable stage label shown in the console.

    console:
        Rich Console instance. If None, a default Console is created.

    spinner:
        Rich spinner name. Examples: "dots", "earth", "bouncingBall".

    Example
    -------
    >>> with stage_status("🦆 Docling: PDF → paragraphs"):
    ...     docs = loader.load()
    """
    c = console or Console()
    start = perf_counter()

    with Status(title, console=c, spinner=spinner):
        yield

    elapsed = perf_counter() - start
    c.print(f"[green]✅ {title}[/green] [dim](took {elapsed:.2f}s)[/dim]")