[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
|
# 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
|
||||||
|
|
Loading…
Add table
Reference in a new issue