[MERGE][CLI] Allow for rich progress bars (24578)
This commit is contained in:
parent
9c397f665b
commit
56e94b9f38
1 changed files with 8 additions and 1 deletions
|
@ -70,13 +70,20 @@ def parallelize():
|
|||
|
||||
# Prefer mpire's `WorkerPool` if it's available
|
||||
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.utils import make_single_arguments
|
||||
with WorkerPool() as pool:
|
||||
|
||||
def _worker(func, *args):
|
||||
# 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 _worker
|
||||
|
|
Loading…
Add table
Reference in a new issue