mirror of
https://github.com/Cog-Creators/Red-DiscordBot.git
synced 2025-12-07 01:42:30 -05:00
Use rich.progress instead of tqdm (#5064)
* Use rich progress instead of tqdm * Remove tqdm from deps
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user