Hykilpikonna
39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
"""
|
|
Importing this file requires installing tqdm.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from functools import partial
|
|
from typing import Callable, Iterable
|
|
|
|
import tqdm
|
|
from tqdm.contrib.concurrent import process_map, thread_map
|
|
|
|
|
|
def smap(fn: Callable, lst: Iterable, *args, **kwargs) -> list:
|
|
return [fn(i) for i in tqdm.tqdm(lst, position=0, leave=True)]
|
|
|
|
|
|
def pmap(fn: Callable, lst: Iterable, *args, **kwargs) -> list:
|
|
tqdm_args = dict(position=0, leave=True, chunksize=1, tqdm_class=tqdm.tqdm, max_workers=os.cpu_count())
|
|
return process_map(fn, lst, *args, **{**tqdm_args, **kwargs})
|
|
|
|
|
|
def tmap(fn: Callable, lst: Iterable, *args, **kwargs) -> list:
|
|
tqdm_args = dict(position=0, leave=True, chunksize=1, tqdm_class=tqdm.tqdm, max_workers=os.cpu_count())
|
|
return thread_map(fn, lst, *args, **{**tqdm_args, **kwargs})
|
|
|
|
|
|
def tq(it: Iterable, desc: str, *args, **kwargs) -> tqdm:
|
|
tqdm_args = dict(position=0, leave=True)
|
|
return tqdm.tqdm(it, desc, *args, **{**tqdm_args, **kwargs})
|
|
|
|
|
|
def patch_tqdm():
|
|
tqdm_args = dict(chunksize=1, position=0, leave=True, tqdm_class=tqdm.tqdm, max_workers=os.cpu_count())
|
|
tq: Callable[[Iterable], tqdm.tqdm] = partial(tqdm.tqdm, position=0, leave=True)
|
|
pmap = partial(process_map, **tqdm_args)
|
|
tmap = partial(thread_map, **tqdm_args)
|
|
return tq, pmap, tmap
|