File size: 3,358 Bytes
6d63e5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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()
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,
)
|