import io
from typing import List, Optional
from rich.progress import BarColumn, DownloadColumn, Progress, TaskID, TimeElapsedColumn
class QuietProgress:
"""
A mock `Progress` class that does absolutely nothing.
We use this when users pass `quiet=True` since rich's `Progress` still
prints empty lines with `quiet=True`.
"""
def start(self, *args, **kwargs):
del args, kwargs
def stop(self, *args, **kwargs):
del args, kwargs
def update(self, *args, **kwargs):
del args, kwargs
def add_task(self, *args, **kwargs):
del args, kwargs
def advance(self, *args, **kwargs):
del args, kwargs
def stop_task(self, *args, **kwargs):
del args, kwargs
def __enter__(self):
return self
def __exit__(self, *args, **kwargs): # type: ignore
del args, kwargs
class BufferedWriterWithProgress(io.BufferedWriter):
def __init__(self, handle: io.BufferedWriter, progress: Progress, task_id: TaskID):
self.handle = handle
self.progress = progress
self.task_id = task_id
self.total_written = 0
def __enter__(self) -> "BufferedWriterWithProgress":
self.handle.__enter__()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
@property
def closed(self) -> bool:
return self.handle.closed
def close(self):
self.handle.close()
def fileno(self):
return self.handle.fileno()
def flush(self):
self.handle.flush()
def isatty(self) -> bool:
return self.handle.isatty()
def readable(self) -> bool:
return self.handle.readable()
def seekable(self) -> bool:
return self.handle.seekable()
def writable(self) -> bool:
return True
def read(self, size: Optional[int] = -1) -> bytes:
return self.handle.read(size)
def read1(self, size: Optional[int] = -1) -> bytes:
return self.handle.read1()
def readinto(self, b):
return self.handle.readinto(b)
def readinto1(self, b):
return self.handle.readinto1(b)
def readline(self, size: Optional[int] = -1) -> bytes:
return self.handle.readline(size)
def readlines(self, hint: int = -1) -> List[bytes]:
return self.handle.readlines(hint)
def write(self, b) -> int:
n = self.handle.write(b)
self.total_written += n
self.progress.advance(self.task_id, n)
return n
def writelines(self, lines):
return self.handle.writelines(lines)
def seek(self, offset: int, whence: int = 0) -> int:
pos = self.handle.seek(offset, whence)
# self.progress.update(self.task_id, completed=pos)
return pos
def tell(self) -> int:
return self.handle.tell()
@property
def raw(self):
return self.handle.raw
def detach(self):
return self.handle.detach()
[docs]def get_download_progress(quiet: bool = False) -> Progress:
if quiet:
return QuietProgress() # type: ignore
else:
return Progress(
"[progress.description]{task.description}",
BarColumn(),
"[progress.percentage]{task.percentage:>3.0f}%",
TimeElapsedColumn(),
DownloadColumn(),
# disable=quiet,
)