58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
import sys
|
|
from datetime import datetime, timedelta
|
|
from typing import Iterable, Iterator, Sequence, TextIO, TypeVar
|
|
|
|
from .base import dump_api_message
|
|
|
|
_T = TypeVar("_T")
|
|
|
|
|
|
class ProgressAPI:
|
|
def __init__(
|
|
self,
|
|
min_step: int = 1,
|
|
min_time: timedelta = timedelta(milliseconds=100),
|
|
output: TextIO = sys.stdout,
|
|
):
|
|
super().__init__()
|
|
|
|
self.counter = 0
|
|
self.output = output
|
|
self.min_step = min_step
|
|
self.min_time = min_time
|
|
|
|
def wrap(
|
|
self, values: Sequence[_T] | Iterable[_T], total: int | None = None
|
|
) -> Iterator[_T]:
|
|
total = total or len(values) # type: ignore
|
|
|
|
current = self.counter
|
|
self.counter += 1
|
|
|
|
dump_api_message("progress-start", {"counter": current, "total": total})
|
|
|
|
try:
|
|
percent = 0
|
|
time = datetime.now()
|
|
|
|
for i_value, value in enumerate(values):
|
|
yield value
|
|
|
|
if datetime.now() - time < self.min_time:
|
|
continue
|
|
|
|
time = datetime.now()
|
|
|
|
c_percent = round(i_value / total * 100)
|
|
|
|
if c_percent >= percent + self.min_step:
|
|
dump_api_message(
|
|
"progress-step", {"counter": current, "percent": c_percent}
|
|
)
|
|
percent = c_percent
|
|
finally:
|
|
dump_api_message(
|
|
"progress-end",
|
|
{"counter": current},
|
|
)
|