Use rich.progress instead of tqdm (#5064)

* Use rich progress instead of tqdm

* Remove tqdm from deps
This commit is contained in:
jack1142
2021-06-03 21:37:53 +02:00
committed by GitHub
parent 0ce2634bb3
commit 8f390147c1
5 changed files with 49 additions and 106 deletions

View File

@@ -33,7 +33,8 @@ import aiohttp
import discord
import pkg_resources
from fuzzywuzzy import fuzz, process
from tqdm import tqdm
from rich.progress import ProgressColumn
from rich.progress_bar import ProgressBar
from redbot import VersionInfo
from redbot.core import data_manager
@@ -55,7 +56,7 @@ __all__ = (
"expected_version",
"fetch_latest_red_version_info",
"deprecated_removed",
"async_tqdm",
"RichIndefiniteBarColumn",
)
_T = TypeVar("_T")
@@ -347,76 +348,12 @@ def deprecated_removed(
)
class _AsyncTqdm(AsyncIterator[_T], tqdm):
def __init__(self, iterable: AsyncIterable[_T], *args, **kwargs) -> None:
self.async_iterator = iterable.__aiter__()
super().__init__(self.infinite_generator(), *args, **kwargs)
self.iterator = cast(Generator[None, bool, None], iter(self))
@staticmethod
def infinite_generator() -> Generator[None, bool, None]:
while True:
# Generator can be forced to raise StopIteration by calling `g.send(True)`
current = yield
if current:
break
async def __anext__(self) -> _T:
try:
result = await self.async_iterator.__anext__()
except StopAsyncIteration:
# If the async iterator is exhausted, force-stop the tqdm iterator
with contextlib.suppress(StopIteration):
self.iterator.send(True)
raise
else:
next(self.iterator)
return result
def __aiter__(self) -> _AsyncTqdm[_T]:
return self
def async_tqdm(
iterable: Optional[Union[Iterable, AsyncIterable]] = None,
*args,
refresh_interval: float = 0.5,
**kwargs,
) -> Union[tqdm, _AsyncTqdm]:
"""Same as `tqdm() <https://tqdm.github.io>`_, except it can be used
in ``async for`` loops, and a task can be spawned to asynchronously
refresh the progress bar every ``refresh_interval`` seconds.
This should only be used for ``async for`` loops, or ``for`` loops
which ``await`` something slow between iterations.
Parameters
----------
iterable: Optional[Union[Iterable, AsyncIterable]]
The iterable to pass to ``tqdm()``. If this is an async
iterable, this function will return a wrapper
*args
Other positional arguments to ``tqdm()``.
refresh_interval : float
The sleep interval between the progress bar being refreshed, in
seconds. Defaults to 0.5. Set to 0 to disable the auto-
refresher.
**kwargs
Keyword arguments to ``tqdm()``.
"""
if isinstance(iterable, AsyncIterable):
progress_bar = _AsyncTqdm(iterable, *args, **kwargs)
else:
progress_bar = tqdm(iterable, *args, **kwargs)
if refresh_interval:
# The background task that refreshes the progress bar
async def _progress_bar_refresher() -> None:
while not progress_bar.disable:
await asyncio.sleep(refresh_interval)
progress_bar.refresh()
asyncio.create_task(_progress_bar_refresher())
return progress_bar
class RichIndefiniteBarColumn(ProgressColumn):
def render(self, task):
return ProgressBar(
pulse=task.completed < task.total,
animation_time=task.get_time(),
width=40,
total=task.total,
completed=task.completed,
)