1
0
Fork 0

[MERGE][CLI] Allow for rich progress bars (24578)

This commit is contained in:
Drashna Jael're 2024-11-08 11:13:56 -08:00
parent 9c397f665b
commit 56e94b9f38
Signed by: drashna
GPG key ID: DBA1FD3A860D1B11

View file

@ -70,13 +70,20 @@ def parallelize():
# Prefer mpire's `WorkerPool` if it's available # Prefer mpire's `WorkerPool` if it's available
with contextlib.suppress(ImportError): with contextlib.suppress(ImportError):
style = 'std'
try:
import rich # noqa: F401 -- we're just testing if it's available
style = 'rich'
except ImportError:
pass
from mpire import WorkerPool from mpire import WorkerPool
from mpire.utils import make_single_arguments from mpire.utils import make_single_arguments
with WorkerPool() as pool: with WorkerPool() as pool:
def _worker(func, *args): def _worker(func, *args):
# Ensure we don't unpack tuples -- mpire's `WorkerPool` tries to do so normally so we tell it not to. # Ensure we don't unpack tuples -- mpire's `WorkerPool` tries to do so normally so we tell it not to.
for r in pool.imap_unordered(func, make_single_arguments(*args, generator=False), progress_bar=True): for r in pool.imap_unordered(func, make_single_arguments(*args, generator=False), progress_bar=True, progress_bar_style=style):
yield r yield r
yield _worker yield _worker