progress.py 59 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715
  1. import io
  2. import sys
  3. import typing
  4. import warnings
  5. from abc import ABC, abstractmethod
  6. from collections import deque
  7. from dataclasses import dataclass, field
  8. from datetime import timedelta
  9. from io import RawIOBase, UnsupportedOperation
  10. from math import ceil
  11. from mmap import mmap
  12. from operator import length_hint
  13. from os import PathLike, stat
  14. from threading import Event, RLock, Thread
  15. from types import TracebackType
  16. from typing import (
  17. Any,
  18. BinaryIO,
  19. Callable,
  20. ContextManager,
  21. Deque,
  22. Dict,
  23. Generic,
  24. Iterable,
  25. List,
  26. NamedTuple,
  27. NewType,
  28. Optional,
  29. Sequence,
  30. TextIO,
  31. Tuple,
  32. Type,
  33. TypeVar,
  34. Union,
  35. )
  36. if sys.version_info >= (3, 8):
  37. from typing import Literal
  38. else:
  39. from typing_extensions import Literal # pragma: no cover
  40. if sys.version_info >= (3, 11):
  41. from typing import Self
  42. else:
  43. from typing_extensions import Self # pragma: no cover
  44. from . import filesize, get_console
  45. from .console import Console, Group, JustifyMethod, RenderableType
  46. from .highlighter import Highlighter
  47. from .jupyter import JupyterMixin
  48. from .live import Live
  49. from .progress_bar import ProgressBar
  50. from .spinner import Spinner
  51. from .style import StyleType
  52. from .table import Column, Table
  53. from .text import Text, TextType
  54. TaskID = NewType("TaskID", int)
  55. ProgressType = TypeVar("ProgressType")
  56. GetTimeCallable = Callable[[], float]
  57. _I = typing.TypeVar("_I", TextIO, BinaryIO)
  58. class _TrackThread(Thread):
  59. """A thread to periodically update progress."""
  60. def __init__(self, progress: "Progress", task_id: "TaskID", update_period: float):
  61. self.progress = progress
  62. self.task_id = task_id
  63. self.update_period = update_period
  64. self.done = Event()
  65. self.completed = 0
  66. super().__init__(daemon=True)
  67. def run(self) -> None:
  68. task_id = self.task_id
  69. advance = self.progress.advance
  70. update_period = self.update_period
  71. last_completed = 0
  72. wait = self.done.wait
  73. while not wait(update_period) and self.progress.live.is_started:
  74. completed = self.completed
  75. if last_completed != completed:
  76. advance(task_id, completed - last_completed)
  77. last_completed = completed
  78. self.progress.update(self.task_id, completed=self.completed, refresh=True)
  79. def __enter__(self) -> "_TrackThread":
  80. self.start()
  81. return self
  82. def __exit__(
  83. self,
  84. exc_type: Optional[Type[BaseException]],
  85. exc_val: Optional[BaseException],
  86. exc_tb: Optional[TracebackType],
  87. ) -> None:
  88. self.done.set()
  89. self.join()
  90. def track(
  91. sequence: Union[Sequence[ProgressType], Iterable[ProgressType]],
  92. description: str = "Working...",
  93. total: Optional[float] = None,
  94. completed: int = 0,
  95. auto_refresh: bool = True,
  96. console: Optional[Console] = None,
  97. transient: bool = False,
  98. get_time: Optional[Callable[[], float]] = None,
  99. refresh_per_second: float = 10,
  100. style: StyleType = "bar.back",
  101. complete_style: StyleType = "bar.complete",
  102. finished_style: StyleType = "bar.finished",
  103. pulse_style: StyleType = "bar.pulse",
  104. update_period: float = 0.1,
  105. disable: bool = False,
  106. show_speed: bool = True,
  107. ) -> Iterable[ProgressType]:
  108. """Track progress by iterating over a sequence.
  109. Args:
  110. sequence (Iterable[ProgressType]): A sequence (must support "len") you wish to iterate over.
  111. description (str, optional): Description of task show next to progress bar. Defaults to "Working".
  112. total: (float, optional): Total number of steps. Default is len(sequence).
  113. completed (int, optional): Number of steps completed so far. Defaults to 0.
  114. auto_refresh (bool, optional): Automatic refresh, disable to force a refresh after each iteration. Default is True.
  115. transient: (bool, optional): Clear the progress on exit. Defaults to False.
  116. console (Console, optional): Console to write to. Default creates internal Console instance.
  117. refresh_per_second (float): Number of times per second to refresh the progress information. Defaults to 10.
  118. style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
  119. complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
  120. finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
  121. pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
  122. update_period (float, optional): Minimum time (in seconds) between calls to update(). Defaults to 0.1.
  123. disable (bool, optional): Disable display of progress.
  124. show_speed (bool, optional): Show speed if total isn't known. Defaults to True.
  125. Returns:
  126. Iterable[ProgressType]: An iterable of the values in the sequence.
  127. """
  128. columns: List["ProgressColumn"] = (
  129. [TextColumn("[progress.description]{task.description}")] if description else []
  130. )
  131. columns.extend(
  132. (
  133. BarColumn(
  134. style=style,
  135. complete_style=complete_style,
  136. finished_style=finished_style,
  137. pulse_style=pulse_style,
  138. ),
  139. TaskProgressColumn(show_speed=show_speed),
  140. TimeRemainingColumn(elapsed_when_finished=True),
  141. )
  142. )
  143. progress = Progress(
  144. *columns,
  145. auto_refresh=auto_refresh,
  146. console=console,
  147. transient=transient,
  148. get_time=get_time,
  149. refresh_per_second=refresh_per_second or 10,
  150. disable=disable,
  151. )
  152. with progress:
  153. yield from progress.track(
  154. sequence,
  155. total=total,
  156. completed=completed,
  157. description=description,
  158. update_period=update_period,
  159. )
  160. class _Reader(RawIOBase, BinaryIO):
  161. """A reader that tracks progress while it's being read from."""
  162. def __init__(
  163. self,
  164. handle: BinaryIO,
  165. progress: "Progress",
  166. task: TaskID,
  167. close_handle: bool = True,
  168. ) -> None:
  169. self.handle = handle
  170. self.progress = progress
  171. self.task = task
  172. self.close_handle = close_handle
  173. self._closed = False
  174. def __enter__(self) -> "_Reader":
  175. self.handle.__enter__()
  176. return self
  177. def __exit__(
  178. self,
  179. exc_type: Optional[Type[BaseException]],
  180. exc_val: Optional[BaseException],
  181. exc_tb: Optional[TracebackType],
  182. ) -> None:
  183. self.close()
  184. def __iter__(self) -> BinaryIO:
  185. return self
  186. def __next__(self) -> bytes:
  187. line = next(self.handle)
  188. self.progress.advance(self.task, advance=len(line))
  189. return line
  190. @property
  191. def closed(self) -> bool:
  192. return self._closed
  193. def fileno(self) -> int:
  194. return self.handle.fileno()
  195. def isatty(self) -> bool:
  196. return self.handle.isatty()
  197. @property
  198. def mode(self) -> str:
  199. return self.handle.mode
  200. @property
  201. def name(self) -> str:
  202. return self.handle.name
  203. def readable(self) -> bool:
  204. return self.handle.readable()
  205. def seekable(self) -> bool:
  206. return self.handle.seekable()
  207. def writable(self) -> bool:
  208. return False
  209. def read(self, size: int = -1) -> bytes:
  210. block = self.handle.read(size)
  211. self.progress.advance(self.task, advance=len(block))
  212. return block
  213. def readinto(self, b: Union[bytearray, memoryview, mmap]): # type: ignore[no-untyped-def, override]
  214. n = self.handle.readinto(b) # type: ignore[attr-defined]
  215. self.progress.advance(self.task, advance=n)
  216. return n
  217. def readline(self, size: int = -1) -> bytes: # type: ignore[override]
  218. line = self.handle.readline(size)
  219. self.progress.advance(self.task, advance=len(line))
  220. return line
  221. def readlines(self, hint: int = -1) -> List[bytes]:
  222. lines = self.handle.readlines(hint)
  223. self.progress.advance(self.task, advance=sum(map(len, lines)))
  224. return lines
  225. def close(self) -> None:
  226. if self.close_handle:
  227. self.handle.close()
  228. self._closed = True
  229. def seek(self, offset: int, whence: int = 0) -> int:
  230. pos = self.handle.seek(offset, whence)
  231. self.progress.update(self.task, completed=pos)
  232. return pos
  233. def tell(self) -> int:
  234. return self.handle.tell()
  235. def write(self, s: Any) -> int:
  236. raise UnsupportedOperation("write")
  237. def writelines(self, lines: Iterable[Any]) -> None:
  238. raise UnsupportedOperation("writelines")
  239. class _ReadContext(ContextManager[_I], Generic[_I]):
  240. """A utility class to handle a context for both a reader and a progress."""
  241. def __init__(self, progress: "Progress", reader: _I) -> None:
  242. self.progress = progress
  243. self.reader: _I = reader
  244. def __enter__(self) -> _I:
  245. self.progress.start()
  246. return self.reader.__enter__()
  247. def __exit__(
  248. self,
  249. exc_type: Optional[Type[BaseException]],
  250. exc_val: Optional[BaseException],
  251. exc_tb: Optional[TracebackType],
  252. ) -> None:
  253. self.progress.stop()
  254. self.reader.__exit__(exc_type, exc_val, exc_tb)
  255. def wrap_file(
  256. file: BinaryIO,
  257. total: int,
  258. *,
  259. description: str = "Reading...",
  260. auto_refresh: bool = True,
  261. console: Optional[Console] = None,
  262. transient: bool = False,
  263. get_time: Optional[Callable[[], float]] = None,
  264. refresh_per_second: float = 10,
  265. style: StyleType = "bar.back",
  266. complete_style: StyleType = "bar.complete",
  267. finished_style: StyleType = "bar.finished",
  268. pulse_style: StyleType = "bar.pulse",
  269. disable: bool = False,
  270. ) -> ContextManager[BinaryIO]:
  271. """Read bytes from a file while tracking progress.
  272. Args:
  273. file (Union[str, PathLike[str], BinaryIO]): The path to the file to read, or a file-like object in binary mode.
  274. total (int): Total number of bytes to read.
  275. description (str, optional): Description of task show next to progress bar. Defaults to "Reading".
  276. auto_refresh (bool, optional): Automatic refresh, disable to force a refresh after each iteration. Default is True.
  277. transient: (bool, optional): Clear the progress on exit. Defaults to False.
  278. console (Console, optional): Console to write to. Default creates internal Console instance.
  279. refresh_per_second (float): Number of times per second to refresh the progress information. Defaults to 10.
  280. style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
  281. complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
  282. finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
  283. pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
  284. disable (bool, optional): Disable display of progress.
  285. Returns:
  286. ContextManager[BinaryIO]: A context manager yielding a progress reader.
  287. """
  288. columns: List["ProgressColumn"] = (
  289. [TextColumn("[progress.description]{task.description}")] if description else []
  290. )
  291. columns.extend(
  292. (
  293. BarColumn(
  294. style=style,
  295. complete_style=complete_style,
  296. finished_style=finished_style,
  297. pulse_style=pulse_style,
  298. ),
  299. DownloadColumn(),
  300. TimeRemainingColumn(),
  301. )
  302. )
  303. progress = Progress(
  304. *columns,
  305. auto_refresh=auto_refresh,
  306. console=console,
  307. transient=transient,
  308. get_time=get_time,
  309. refresh_per_second=refresh_per_second or 10,
  310. disable=disable,
  311. )
  312. reader = progress.wrap_file(file, total=total, description=description)
  313. return _ReadContext(progress, reader)
  314. @typing.overload
  315. def open(
  316. file: Union[str, "PathLike[str]", bytes],
  317. mode: Union[Literal["rt"], Literal["r"]],
  318. buffering: int = -1,
  319. encoding: Optional[str] = None,
  320. errors: Optional[str] = None,
  321. newline: Optional[str] = None,
  322. *,
  323. total: Optional[int] = None,
  324. description: str = "Reading...",
  325. auto_refresh: bool = True,
  326. console: Optional[Console] = None,
  327. transient: bool = False,
  328. get_time: Optional[Callable[[], float]] = None,
  329. refresh_per_second: float = 10,
  330. style: StyleType = "bar.back",
  331. complete_style: StyleType = "bar.complete",
  332. finished_style: StyleType = "bar.finished",
  333. pulse_style: StyleType = "bar.pulse",
  334. disable: bool = False,
  335. ) -> ContextManager[TextIO]:
  336. pass
  337. @typing.overload
  338. def open(
  339. file: Union[str, "PathLike[str]", bytes],
  340. mode: Literal["rb"],
  341. buffering: int = -1,
  342. encoding: Optional[str] = None,
  343. errors: Optional[str] = None,
  344. newline: Optional[str] = None,
  345. *,
  346. total: Optional[int] = None,
  347. description: str = "Reading...",
  348. auto_refresh: bool = True,
  349. console: Optional[Console] = None,
  350. transient: bool = False,
  351. get_time: Optional[Callable[[], float]] = None,
  352. refresh_per_second: float = 10,
  353. style: StyleType = "bar.back",
  354. complete_style: StyleType = "bar.complete",
  355. finished_style: StyleType = "bar.finished",
  356. pulse_style: StyleType = "bar.pulse",
  357. disable: bool = False,
  358. ) -> ContextManager[BinaryIO]:
  359. pass
  360. def open(
  361. file: Union[str, "PathLike[str]", bytes],
  362. mode: Union[Literal["rb"], Literal["rt"], Literal["r"]] = "r",
  363. buffering: int = -1,
  364. encoding: Optional[str] = None,
  365. errors: Optional[str] = None,
  366. newline: Optional[str] = None,
  367. *,
  368. total: Optional[int] = None,
  369. description: str = "Reading...",
  370. auto_refresh: bool = True,
  371. console: Optional[Console] = None,
  372. transient: bool = False,
  373. get_time: Optional[Callable[[], float]] = None,
  374. refresh_per_second: float = 10,
  375. style: StyleType = "bar.back",
  376. complete_style: StyleType = "bar.complete",
  377. finished_style: StyleType = "bar.finished",
  378. pulse_style: StyleType = "bar.pulse",
  379. disable: bool = False,
  380. ) -> Union[ContextManager[BinaryIO], ContextManager[TextIO]]:
  381. """Read bytes from a file while tracking progress.
  382. Args:
  383. path (Union[str, PathLike[str], BinaryIO]): The path to the file to read, or a file-like object in binary mode.
  384. mode (str): The mode to use to open the file. Only supports "r", "rb" or "rt".
  385. buffering (int): The buffering strategy to use, see :func:`io.open`.
  386. encoding (str, optional): The encoding to use when reading in text mode, see :func:`io.open`.
  387. errors (str, optional): The error handling strategy for decoding errors, see :func:`io.open`.
  388. newline (str, optional): The strategy for handling newlines in text mode, see :func:`io.open`
  389. total: (int, optional): Total number of bytes to read. Must be provided if reading from a file handle. Default for a path is os.stat(file).st_size.
  390. description (str, optional): Description of task show next to progress bar. Defaults to "Reading".
  391. auto_refresh (bool, optional): Automatic refresh, disable to force a refresh after each iteration. Default is True.
  392. transient: (bool, optional): Clear the progress on exit. Defaults to False.
  393. console (Console, optional): Console to write to. Default creates internal Console instance.
  394. refresh_per_second (float): Number of times per second to refresh the progress information. Defaults to 10.
  395. style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
  396. complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
  397. finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
  398. pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
  399. disable (bool, optional): Disable display of progress.
  400. encoding (str, optional): The encoding to use when reading in text mode.
  401. Returns:
  402. ContextManager[BinaryIO]: A context manager yielding a progress reader.
  403. """
  404. columns: List["ProgressColumn"] = (
  405. [TextColumn("[progress.description]{task.description}")] if description else []
  406. )
  407. columns.extend(
  408. (
  409. BarColumn(
  410. style=style,
  411. complete_style=complete_style,
  412. finished_style=finished_style,
  413. pulse_style=pulse_style,
  414. ),
  415. DownloadColumn(),
  416. TimeRemainingColumn(),
  417. )
  418. )
  419. progress = Progress(
  420. *columns,
  421. auto_refresh=auto_refresh,
  422. console=console,
  423. transient=transient,
  424. get_time=get_time,
  425. refresh_per_second=refresh_per_second or 10,
  426. disable=disable,
  427. )
  428. reader = progress.open(
  429. file,
  430. mode=mode,
  431. buffering=buffering,
  432. encoding=encoding,
  433. errors=errors,
  434. newline=newline,
  435. total=total,
  436. description=description,
  437. )
  438. return _ReadContext(progress, reader) # type: ignore[return-value, type-var]
  439. class ProgressColumn(ABC):
  440. """Base class for a widget to use in progress display."""
  441. max_refresh: Optional[float] = None
  442. def __init__(self, table_column: Optional[Column] = None) -> None:
  443. self._table_column = table_column
  444. self._renderable_cache: Dict[TaskID, Tuple[float, RenderableType]] = {}
  445. self._update_time: Optional[float] = None
  446. def get_table_column(self) -> Column:
  447. """Get a table column, used to build tasks table."""
  448. return self._table_column or Column()
  449. def __call__(self, task: "Task") -> RenderableType:
  450. """Called by the Progress object to return a renderable for the given task.
  451. Args:
  452. task (Task): An object containing information regarding the task.
  453. Returns:
  454. RenderableType: Anything renderable (including str).
  455. """
  456. current_time = task.get_time()
  457. if self.max_refresh is not None and not task.completed:
  458. try:
  459. timestamp, renderable = self._renderable_cache[task.id]
  460. except KeyError:
  461. pass
  462. else:
  463. if timestamp + self.max_refresh > current_time:
  464. return renderable
  465. renderable = self.render(task)
  466. self._renderable_cache[task.id] = (current_time, renderable)
  467. return renderable
  468. @abstractmethod
  469. def render(self, task: "Task") -> RenderableType:
  470. """Should return a renderable object."""
  471. class RenderableColumn(ProgressColumn):
  472. """A column to insert an arbitrary column.
  473. Args:
  474. renderable (RenderableType, optional): Any renderable. Defaults to empty string.
  475. """
  476. def __init__(
  477. self, renderable: RenderableType = "", *, table_column: Optional[Column] = None
  478. ):
  479. self.renderable = renderable
  480. super().__init__(table_column=table_column)
  481. def render(self, task: "Task") -> RenderableType:
  482. return self.renderable
  483. class SpinnerColumn(ProgressColumn):
  484. """A column with a 'spinner' animation.
  485. Args:
  486. spinner_name (str, optional): Name of spinner animation. Defaults to "dots".
  487. style (StyleType, optional): Style of spinner. Defaults to "progress.spinner".
  488. speed (float, optional): Speed factor of spinner. Defaults to 1.0.
  489. finished_text (TextType, optional): Text used when task is finished. Defaults to " ".
  490. """
  491. def __init__(
  492. self,
  493. spinner_name: str = "dots",
  494. style: Optional[StyleType] = "progress.spinner",
  495. speed: float = 1.0,
  496. finished_text: TextType = " ",
  497. table_column: Optional[Column] = None,
  498. ):
  499. self.spinner = Spinner(spinner_name, style=style, speed=speed)
  500. self.finished_text = (
  501. Text.from_markup(finished_text)
  502. if isinstance(finished_text, str)
  503. else finished_text
  504. )
  505. super().__init__(table_column=table_column)
  506. def set_spinner(
  507. self,
  508. spinner_name: str,
  509. spinner_style: Optional[StyleType] = "progress.spinner",
  510. speed: float = 1.0,
  511. ) -> None:
  512. """Set a new spinner.
  513. Args:
  514. spinner_name (str): Spinner name, see python -m rich.spinner.
  515. spinner_style (Optional[StyleType], optional): Spinner style. Defaults to "progress.spinner".
  516. speed (float, optional): Speed factor of spinner. Defaults to 1.0.
  517. """
  518. self.spinner = Spinner(spinner_name, style=spinner_style, speed=speed)
  519. def render(self, task: "Task") -> RenderableType:
  520. text = (
  521. self.finished_text
  522. if task.finished
  523. else self.spinner.render(task.get_time())
  524. )
  525. return text
  526. class TextColumn(ProgressColumn):
  527. """A column containing text."""
  528. def __init__(
  529. self,
  530. text_format: str,
  531. style: StyleType = "none",
  532. justify: JustifyMethod = "left",
  533. markup: bool = True,
  534. highlighter: Optional[Highlighter] = None,
  535. table_column: Optional[Column] = None,
  536. ) -> None:
  537. self.text_format = text_format
  538. self.justify: JustifyMethod = justify
  539. self.style = style
  540. self.markup = markup
  541. self.highlighter = highlighter
  542. super().__init__(table_column=table_column or Column(no_wrap=True))
  543. def render(self, task: "Task") -> Text:
  544. _text = self.text_format.format(task=task)
  545. if self.markup:
  546. text = Text.from_markup(_text, style=self.style, justify=self.justify)
  547. else:
  548. text = Text(_text, style=self.style, justify=self.justify)
  549. if self.highlighter:
  550. self.highlighter.highlight(text)
  551. return text
  552. class BarColumn(ProgressColumn):
  553. """Renders a visual progress bar.
  554. Args:
  555. bar_width (Optional[int], optional): Width of bar or None for full width. Defaults to 40.
  556. style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
  557. complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
  558. finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
  559. pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
  560. """
  561. def __init__(
  562. self,
  563. bar_width: Optional[int] = 40,
  564. style: StyleType = "bar.back",
  565. complete_style: StyleType = "bar.complete",
  566. finished_style: StyleType = "bar.finished",
  567. pulse_style: StyleType = "bar.pulse",
  568. table_column: Optional[Column] = None,
  569. ) -> None:
  570. self.bar_width = bar_width
  571. self.style = style
  572. self.complete_style = complete_style
  573. self.finished_style = finished_style
  574. self.pulse_style = pulse_style
  575. super().__init__(table_column=table_column)
  576. def render(self, task: "Task") -> ProgressBar:
  577. """Gets a progress bar widget for a task."""
  578. return ProgressBar(
  579. total=max(0, task.total) if task.total is not None else None,
  580. completed=max(0, task.completed),
  581. width=None if self.bar_width is None else max(1, self.bar_width),
  582. pulse=not task.started,
  583. animation_time=task.get_time(),
  584. style=self.style,
  585. complete_style=self.complete_style,
  586. finished_style=self.finished_style,
  587. pulse_style=self.pulse_style,
  588. )
  589. class TimeElapsedColumn(ProgressColumn):
  590. """Renders time elapsed."""
  591. def render(self, task: "Task") -> Text:
  592. """Show time elapsed."""
  593. elapsed = task.finished_time if task.finished else task.elapsed
  594. if elapsed is None:
  595. return Text("-:--:--", style="progress.elapsed")
  596. delta = timedelta(seconds=max(0, int(elapsed)))
  597. return Text(str(delta), style="progress.elapsed")
  598. class TaskProgressColumn(TextColumn):
  599. """Show task progress as a percentage.
  600. Args:
  601. text_format (str, optional): Format for percentage display. Defaults to "[progress.percentage]{task.percentage:>3.0f}%".
  602. text_format_no_percentage (str, optional): Format if percentage is unknown. Defaults to "".
  603. style (StyleType, optional): Style of output. Defaults to "none".
  604. justify (JustifyMethod, optional): Text justification. Defaults to "left".
  605. markup (bool, optional): Enable markup. Defaults to True.
  606. highlighter (Optional[Highlighter], optional): Highlighter to apply to output. Defaults to None.
  607. table_column (Optional[Column], optional): Table Column to use. Defaults to None.
  608. show_speed (bool, optional): Show speed if total is unknown. Defaults to False.
  609. """
  610. def __init__(
  611. self,
  612. text_format: str = "[progress.percentage]{task.percentage:>3.0f}%",
  613. text_format_no_percentage: str = "",
  614. style: StyleType = "none",
  615. justify: JustifyMethod = "left",
  616. markup: bool = True,
  617. highlighter: Optional[Highlighter] = None,
  618. table_column: Optional[Column] = None,
  619. show_speed: bool = False,
  620. ) -> None:
  621. self.text_format_no_percentage = text_format_no_percentage
  622. self.show_speed = show_speed
  623. super().__init__(
  624. text_format=text_format,
  625. style=style,
  626. justify=justify,
  627. markup=markup,
  628. highlighter=highlighter,
  629. table_column=table_column,
  630. )
  631. @classmethod
  632. def render_speed(cls, speed: Optional[float]) -> Text:
  633. """Render the speed in iterations per second.
  634. Args:
  635. task (Task): A Task object.
  636. Returns:
  637. Text: Text object containing the task speed.
  638. """
  639. if speed is None:
  640. return Text("", style="progress.percentage")
  641. unit, suffix = filesize.pick_unit_and_suffix(
  642. int(speed),
  643. ["", "×10³", "×10⁶", "×10⁹", "×10¹²"],
  644. 1000,
  645. )
  646. data_speed = speed / unit
  647. return Text(f"{data_speed:.1f}{suffix} it/s", style="progress.percentage")
  648. def render(self, task: "Task") -> Text:
  649. if task.total is None and self.show_speed:
  650. return self.render_speed(task.finished_speed or task.speed)
  651. text_format = (
  652. self.text_format_no_percentage if task.total is None else self.text_format
  653. )
  654. _text = text_format.format(task=task)
  655. if self.markup:
  656. text = Text.from_markup(_text, style=self.style, justify=self.justify)
  657. else:
  658. text = Text(_text, style=self.style, justify=self.justify)
  659. if self.highlighter:
  660. self.highlighter.highlight(text)
  661. return text
  662. class TimeRemainingColumn(ProgressColumn):
  663. """Renders estimated time remaining.
  664. Args:
  665. compact (bool, optional): Render MM:SS when time remaining is less than an hour. Defaults to False.
  666. elapsed_when_finished (bool, optional): Render time elapsed when the task is finished. Defaults to False.
  667. """
  668. # Only refresh twice a second to prevent jitter
  669. max_refresh = 0.5
  670. def __init__(
  671. self,
  672. compact: bool = False,
  673. elapsed_when_finished: bool = False,
  674. table_column: Optional[Column] = None,
  675. ):
  676. self.compact = compact
  677. self.elapsed_when_finished = elapsed_when_finished
  678. super().__init__(table_column=table_column)
  679. def render(self, task: "Task") -> Text:
  680. """Show time remaining."""
  681. if self.elapsed_when_finished and task.finished:
  682. task_time = task.finished_time
  683. style = "progress.elapsed"
  684. else:
  685. task_time = task.time_remaining
  686. style = "progress.remaining"
  687. if task.total is None:
  688. return Text("", style=style)
  689. if task_time is None:
  690. return Text("--:--" if self.compact else "-:--:--", style=style)
  691. # Based on https://github.com/tqdm/tqdm/blob/master/tqdm/std.py
  692. minutes, seconds = divmod(int(task_time), 60)
  693. hours, minutes = divmod(minutes, 60)
  694. if self.compact and not hours:
  695. formatted = f"{minutes:02d}:{seconds:02d}"
  696. else:
  697. formatted = f"{hours:d}:{minutes:02d}:{seconds:02d}"
  698. return Text(formatted, style=style)
  699. class FileSizeColumn(ProgressColumn):
  700. """Renders completed filesize."""
  701. def render(self, task: "Task") -> Text:
  702. """Show data completed."""
  703. data_size = filesize.decimal(int(task.completed))
  704. return Text(data_size, style="progress.filesize")
  705. class TotalFileSizeColumn(ProgressColumn):
  706. """Renders total filesize."""
  707. def render(self, task: "Task") -> Text:
  708. """Show data completed."""
  709. data_size = filesize.decimal(int(task.total)) if task.total is not None else ""
  710. return Text(data_size, style="progress.filesize.total")
  711. class MofNCompleteColumn(ProgressColumn):
  712. """Renders completed count/total, e.g. ' 10/1000'.
  713. Best for bounded tasks with int quantities.
  714. Space pads the completed count so that progress length does not change as task progresses
  715. past powers of 10.
  716. Args:
  717. separator (str, optional): Text to separate completed and total values. Defaults to "/".
  718. """
  719. def __init__(self, separator: str = "/", table_column: Optional[Column] = None):
  720. self.separator = separator
  721. super().__init__(table_column=table_column)
  722. def render(self, task: "Task") -> Text:
  723. """Show completed/total."""
  724. completed = int(task.completed)
  725. total = int(task.total) if task.total is not None else "?"
  726. total_width = len(str(total))
  727. return Text(
  728. f"{completed:{total_width}d}{self.separator}{total}",
  729. style="progress.download",
  730. )
  731. class DownloadColumn(ProgressColumn):
  732. """Renders file size downloaded and total, e.g. '0.5/2.3 GB'.
  733. Args:
  734. binary_units (bool, optional): Use binary units, KiB, MiB etc. Defaults to False.
  735. """
  736. def __init__(
  737. self, binary_units: bool = False, table_column: Optional[Column] = None
  738. ) -> None:
  739. self.binary_units = binary_units
  740. super().__init__(table_column=table_column)
  741. def render(self, task: "Task") -> Text:
  742. """Calculate common unit for completed and total."""
  743. completed = int(task.completed)
  744. unit_and_suffix_calculation_base = (
  745. int(task.total) if task.total is not None else completed
  746. )
  747. if self.binary_units:
  748. unit, suffix = filesize.pick_unit_and_suffix(
  749. unit_and_suffix_calculation_base,
  750. ["bytes", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"],
  751. 1024,
  752. )
  753. else:
  754. unit, suffix = filesize.pick_unit_and_suffix(
  755. unit_and_suffix_calculation_base,
  756. ["bytes", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"],
  757. 1000,
  758. )
  759. precision = 0 if unit == 1 else 1
  760. completed_ratio = completed / unit
  761. completed_str = f"{completed_ratio:,.{precision}f}"
  762. if task.total is not None:
  763. total = int(task.total)
  764. total_ratio = total / unit
  765. total_str = f"{total_ratio:,.{precision}f}"
  766. else:
  767. total_str = "?"
  768. download_status = f"{completed_str}/{total_str} {suffix}"
  769. download_text = Text(download_status, style="progress.download")
  770. return download_text
  771. class TransferSpeedColumn(ProgressColumn):
  772. """Renders human readable transfer speed."""
  773. def render(self, task: "Task") -> Text:
  774. """Show data transfer speed."""
  775. speed = task.finished_speed or task.speed
  776. if speed is None:
  777. return Text("?", style="progress.data.speed")
  778. data_speed = filesize.decimal(int(speed))
  779. return Text(f"{data_speed}/s", style="progress.data.speed")
  780. class ProgressSample(NamedTuple):
  781. """Sample of progress for a given time."""
  782. timestamp: float
  783. """Timestamp of sample."""
  784. completed: float
  785. """Number of steps completed."""
  786. @dataclass
  787. class Task:
  788. """Information regarding a progress task.
  789. This object should be considered read-only outside of the :class:`~Progress` class.
  790. """
  791. id: TaskID
  792. """Task ID associated with this task (used in Progress methods)."""
  793. description: str
  794. """str: Description of the task."""
  795. total: Optional[float]
  796. """Optional[float]: Total number of steps in this task."""
  797. completed: float
  798. """float: Number of steps completed"""
  799. _get_time: GetTimeCallable
  800. """Callable to get the current time."""
  801. finished_time: Optional[float] = None
  802. """float: Time task was finished."""
  803. visible: bool = True
  804. """bool: Indicates if this task is visible in the progress display."""
  805. fields: Dict[str, Any] = field(default_factory=dict)
  806. """dict: Arbitrary fields passed in via Progress.update."""
  807. start_time: Optional[float] = field(default=None, init=False, repr=False)
  808. """Optional[float]: Time this task was started, or None if not started."""
  809. stop_time: Optional[float] = field(default=None, init=False, repr=False)
  810. """Optional[float]: Time this task was stopped, or None if not stopped."""
  811. finished_speed: Optional[float] = None
  812. """Optional[float]: The last speed for a finished task."""
  813. _progress: Deque[ProgressSample] = field(
  814. default_factory=lambda: deque(maxlen=1000), init=False, repr=False
  815. )
  816. _lock: RLock = field(repr=False, default_factory=RLock)
  817. """Thread lock."""
  818. def get_time(self) -> float:
  819. """float: Get the current time, in seconds."""
  820. return self._get_time()
  821. @property
  822. def started(self) -> bool:
  823. """bool: Check if the task as started."""
  824. return self.start_time is not None
  825. @property
  826. def remaining(self) -> Optional[float]:
  827. """Optional[float]: Get the number of steps remaining, if a non-None total was set."""
  828. if self.total is None:
  829. return None
  830. return self.total - self.completed
  831. @property
  832. def elapsed(self) -> Optional[float]:
  833. """Optional[float]: Time elapsed since task was started, or ``None`` if the task hasn't started."""
  834. if self.start_time is None:
  835. return None
  836. if self.stop_time is not None:
  837. return self.stop_time - self.start_time
  838. return self.get_time() - self.start_time
  839. @property
  840. def finished(self) -> bool:
  841. """Check if the task has finished."""
  842. return self.finished_time is not None
  843. @property
  844. def percentage(self) -> float:
  845. """float: Get progress of task as a percentage. If a None total was set, returns 0"""
  846. if not self.total:
  847. return 0.0
  848. completed = (self.completed / self.total) * 100.0
  849. completed = min(100.0, max(0.0, completed))
  850. return completed
  851. @property
  852. def speed(self) -> Optional[float]:
  853. """Optional[float]: Get the estimated speed in steps per second."""
  854. if self.start_time is None:
  855. return None
  856. with self._lock:
  857. progress = self._progress
  858. if not progress:
  859. return None
  860. total_time = progress[-1].timestamp - progress[0].timestamp
  861. if total_time == 0:
  862. return None
  863. iter_progress = iter(progress)
  864. next(iter_progress)
  865. total_completed = sum(sample.completed for sample in iter_progress)
  866. speed = total_completed / total_time
  867. return speed
  868. @property
  869. def time_remaining(self) -> Optional[float]:
  870. """Optional[float]: Get estimated time to completion, or ``None`` if no data."""
  871. if self.finished:
  872. return 0.0
  873. speed = self.speed
  874. if not speed:
  875. return None
  876. remaining = self.remaining
  877. if remaining is None:
  878. return None
  879. estimate = ceil(remaining / speed)
  880. return estimate
  881. def _reset(self) -> None:
  882. """Reset progress."""
  883. self._progress.clear()
  884. self.finished_time = None
  885. self.finished_speed = None
  886. class Progress(JupyterMixin):
  887. """Renders an auto-updating progress bar(s).
  888. Args:
  889. console (Console, optional): Optional Console instance. Defaults to an internal Console instance writing to stdout.
  890. auto_refresh (bool, optional): Enable auto refresh. If disabled, you will need to call `refresh()`.
  891. refresh_per_second (Optional[float], optional): Number of times per second to refresh the progress information or None to use default (10). Defaults to None.
  892. speed_estimate_period: (float, optional): Period (in seconds) used to calculate the speed estimate. Defaults to 30.
  893. transient: (bool, optional): Clear the progress on exit. Defaults to False.
  894. redirect_stdout: (bool, optional): Enable redirection of stdout, so ``print`` may be used. Defaults to True.
  895. redirect_stderr: (bool, optional): Enable redirection of stderr. Defaults to True.
  896. get_time: (Callable, optional): A callable that gets the current time, or None to use Console.get_time. Defaults to None.
  897. disable (bool, optional): Disable progress display. Defaults to False
  898. expand (bool, optional): Expand tasks table to fit width. Defaults to False.
  899. """
  900. def __init__(
  901. self,
  902. *columns: Union[str, ProgressColumn],
  903. console: Optional[Console] = None,
  904. auto_refresh: bool = True,
  905. refresh_per_second: float = 10,
  906. speed_estimate_period: float = 30.0,
  907. transient: bool = False,
  908. redirect_stdout: bool = True,
  909. redirect_stderr: bool = True,
  910. get_time: Optional[GetTimeCallable] = None,
  911. disable: bool = False,
  912. expand: bool = False,
  913. ) -> None:
  914. assert refresh_per_second > 0, "refresh_per_second must be > 0"
  915. self._lock = RLock()
  916. self.columns = columns or self.get_default_columns()
  917. self.speed_estimate_period = speed_estimate_period
  918. self.disable = disable
  919. self.expand = expand
  920. self._tasks: Dict[TaskID, Task] = {}
  921. self._task_index: TaskID = TaskID(0)
  922. self.live = Live(
  923. console=console or get_console(),
  924. auto_refresh=auto_refresh,
  925. refresh_per_second=refresh_per_second,
  926. transient=transient,
  927. redirect_stdout=redirect_stdout,
  928. redirect_stderr=redirect_stderr,
  929. get_renderable=self.get_renderable,
  930. )
  931. self.get_time = get_time or self.console.get_time
  932. self.print = self.console.print
  933. self.log = self.console.log
  934. @classmethod
  935. def get_default_columns(cls) -> Tuple[ProgressColumn, ...]:
  936. """Get the default columns used for a new Progress instance:
  937. - a text column for the description (TextColumn)
  938. - the bar itself (BarColumn)
  939. - a text column showing completion percentage (TextColumn)
  940. - an estimated-time-remaining column (TimeRemainingColumn)
  941. If the Progress instance is created without passing a columns argument,
  942. the default columns defined here will be used.
  943. You can also create a Progress instance using custom columns before
  944. and/or after the defaults, as in this example:
  945. progress = Progress(
  946. SpinnerColumn(),
  947. *Progress.get_default_columns(),
  948. "Elapsed:",
  949. TimeElapsedColumn(),
  950. )
  951. This code shows the creation of a Progress display, containing
  952. a spinner to the left, the default columns, and a labeled elapsed
  953. time column.
  954. """
  955. return (
  956. TextColumn("[progress.description]{task.description}"),
  957. BarColumn(),
  958. TaskProgressColumn(),
  959. TimeRemainingColumn(),
  960. )
  961. @property
  962. def console(self) -> Console:
  963. return self.live.console
  964. @property
  965. def tasks(self) -> List[Task]:
  966. """Get a list of Task instances."""
  967. with self._lock:
  968. return list(self._tasks.values())
  969. @property
  970. def task_ids(self) -> List[TaskID]:
  971. """A list of task IDs."""
  972. with self._lock:
  973. return list(self._tasks.keys())
  974. @property
  975. def finished(self) -> bool:
  976. """Check if all tasks have been completed."""
  977. with self._lock:
  978. if not self._tasks:
  979. return True
  980. return all(task.finished for task in self._tasks.values())
  981. def start(self) -> None:
  982. """Start the progress display."""
  983. if not self.disable:
  984. self.live.start(refresh=True)
  985. def stop(self) -> None:
  986. """Stop the progress display."""
  987. self.live.stop()
  988. if not self.console.is_interactive and not self.console.is_jupyter:
  989. self.console.print()
  990. def __enter__(self) -> Self:
  991. self.start()
  992. return self
  993. def __exit__(
  994. self,
  995. exc_type: Optional[Type[BaseException]],
  996. exc_val: Optional[BaseException],
  997. exc_tb: Optional[TracebackType],
  998. ) -> None:
  999. self.stop()
  1000. def track(
  1001. self,
  1002. sequence: Union[Iterable[ProgressType], Sequence[ProgressType]],
  1003. total: Optional[float] = None,
  1004. completed: int = 0,
  1005. task_id: Optional[TaskID] = None,
  1006. description: str = "Working...",
  1007. update_period: float = 0.1,
  1008. ) -> Iterable[ProgressType]:
  1009. """Track progress by iterating over a sequence.
  1010. Args:
  1011. sequence (Sequence[ProgressType]): A sequence of values you want to iterate over and track progress.
  1012. total: (float, optional): Total number of steps. Default is len(sequence).
  1013. completed (int, optional): Number of steps completed so far. Defaults to 0.
  1014. task_id: (TaskID): Task to track. Default is new task.
  1015. description: (str, optional): Description of task, if new task is created.
  1016. update_period (float, optional): Minimum time (in seconds) between calls to update(). Defaults to 0.1.
  1017. Returns:
  1018. Iterable[ProgressType]: An iterable of values taken from the provided sequence.
  1019. """
  1020. if total is None:
  1021. total = float(length_hint(sequence)) or None
  1022. if task_id is None:
  1023. task_id = self.add_task(description, total=total, completed=completed)
  1024. else:
  1025. self.update(task_id, total=total, completed=completed)
  1026. if self.live.auto_refresh:
  1027. with _TrackThread(self, task_id, update_period) as track_thread:
  1028. for value in sequence:
  1029. yield value
  1030. track_thread.completed += 1
  1031. else:
  1032. advance = self.advance
  1033. refresh = self.refresh
  1034. for value in sequence:
  1035. yield value
  1036. advance(task_id, 1)
  1037. refresh()
  1038. def wrap_file(
  1039. self,
  1040. file: BinaryIO,
  1041. total: Optional[int] = None,
  1042. *,
  1043. task_id: Optional[TaskID] = None,
  1044. description: str = "Reading...",
  1045. ) -> BinaryIO:
  1046. """Track progress file reading from a binary file.
  1047. Args:
  1048. file (BinaryIO): A file-like object opened in binary mode.
  1049. total (int, optional): Total number of bytes to read. This must be provided unless a task with a total is also given.
  1050. task_id (TaskID): Task to track. Default is new task.
  1051. description (str, optional): Description of task, if new task is created.
  1052. Returns:
  1053. BinaryIO: A readable file-like object in binary mode.
  1054. Raises:
  1055. ValueError: When no total value can be extracted from the arguments or the task.
  1056. """
  1057. # attempt to recover the total from the task
  1058. total_bytes: Optional[float] = None
  1059. if total is not None:
  1060. total_bytes = total
  1061. elif task_id is not None:
  1062. with self._lock:
  1063. total_bytes = self._tasks[task_id].total
  1064. if total_bytes is None:
  1065. raise ValueError(
  1066. f"unable to get the total number of bytes, please specify 'total'"
  1067. )
  1068. # update total of task or create new task
  1069. if task_id is None:
  1070. task_id = self.add_task(description, total=total_bytes)
  1071. else:
  1072. self.update(task_id, total=total_bytes)
  1073. return _Reader(file, self, task_id, close_handle=False)
  1074. @typing.overload
  1075. def open(
  1076. self,
  1077. file: Union[str, "PathLike[str]", bytes],
  1078. mode: Literal["rb"],
  1079. buffering: int = -1,
  1080. encoding: Optional[str] = None,
  1081. errors: Optional[str] = None,
  1082. newline: Optional[str] = None,
  1083. *,
  1084. total: Optional[int] = None,
  1085. task_id: Optional[TaskID] = None,
  1086. description: str = "Reading...",
  1087. ) -> BinaryIO:
  1088. pass
  1089. @typing.overload
  1090. def open(
  1091. self,
  1092. file: Union[str, "PathLike[str]", bytes],
  1093. mode: Union[Literal["r"], Literal["rt"]],
  1094. buffering: int = -1,
  1095. encoding: Optional[str] = None,
  1096. errors: Optional[str] = None,
  1097. newline: Optional[str] = None,
  1098. *,
  1099. total: Optional[int] = None,
  1100. task_id: Optional[TaskID] = None,
  1101. description: str = "Reading...",
  1102. ) -> TextIO:
  1103. pass
  1104. def open(
  1105. self,
  1106. file: Union[str, "PathLike[str]", bytes],
  1107. mode: Union[Literal["rb"], Literal["rt"], Literal["r"]] = "r",
  1108. buffering: int = -1,
  1109. encoding: Optional[str] = None,
  1110. errors: Optional[str] = None,
  1111. newline: Optional[str] = None,
  1112. *,
  1113. total: Optional[int] = None,
  1114. task_id: Optional[TaskID] = None,
  1115. description: str = "Reading...",
  1116. ) -> Union[BinaryIO, TextIO]:
  1117. """Track progress while reading from a binary file.
  1118. Args:
  1119. path (Union[str, PathLike[str]]): The path to the file to read.
  1120. mode (str): The mode to use to open the file. Only supports "r", "rb" or "rt".
  1121. buffering (int): The buffering strategy to use, see :func:`io.open`.
  1122. encoding (str, optional): The encoding to use when reading in text mode, see :func:`io.open`.
  1123. errors (str, optional): The error handling strategy for decoding errors, see :func:`io.open`.
  1124. newline (str, optional): The strategy for handling newlines in text mode, see :func:`io.open`.
  1125. total (int, optional): Total number of bytes to read. If none given, os.stat(path).st_size is used.
  1126. task_id (TaskID): Task to track. Default is new task.
  1127. description (str, optional): Description of task, if new task is created.
  1128. Returns:
  1129. BinaryIO: A readable file-like object in binary mode.
  1130. Raises:
  1131. ValueError: When an invalid mode is given.
  1132. """
  1133. # normalize the mode (always rb, rt)
  1134. _mode = "".join(sorted(mode, reverse=False))
  1135. if _mode not in ("br", "rt", "r"):
  1136. raise ValueError(f"invalid mode {mode!r}")
  1137. # patch buffering to provide the same behaviour as the builtin `open`
  1138. line_buffering = buffering == 1
  1139. if _mode == "br" and buffering == 1:
  1140. warnings.warn(
  1141. "line buffering (buffering=1) isn't supported in binary mode, the default buffer size will be used",
  1142. RuntimeWarning,
  1143. )
  1144. buffering = -1
  1145. elif _mode in ("rt", "r"):
  1146. if buffering == 0:
  1147. raise ValueError("can't have unbuffered text I/O")
  1148. elif buffering == 1:
  1149. buffering = -1
  1150. # attempt to get the total with `os.stat`
  1151. if total is None:
  1152. total = stat(file).st_size
  1153. # update total of task or create new task
  1154. if task_id is None:
  1155. task_id = self.add_task(description, total=total)
  1156. else:
  1157. self.update(task_id, total=total)
  1158. # open the file in binary mode,
  1159. handle = io.open(file, "rb", buffering=buffering)
  1160. reader = _Reader(handle, self, task_id, close_handle=True)
  1161. # wrap the reader in a `TextIOWrapper` if text mode
  1162. if mode in ("r", "rt"):
  1163. return io.TextIOWrapper(
  1164. reader,
  1165. encoding=encoding,
  1166. errors=errors,
  1167. newline=newline,
  1168. line_buffering=line_buffering,
  1169. )
  1170. return reader
  1171. def start_task(self, task_id: TaskID) -> None:
  1172. """Start a task.
  1173. Starts a task (used when calculating elapsed time). You may need to call this manually,
  1174. if you called ``add_task`` with ``start=False``.
  1175. Args:
  1176. task_id (TaskID): ID of task.
  1177. """
  1178. with self._lock:
  1179. task = self._tasks[task_id]
  1180. if task.start_time is None:
  1181. task.start_time = self.get_time()
  1182. def stop_task(self, task_id: TaskID) -> None:
  1183. """Stop a task.
  1184. This will freeze the elapsed time on the task.
  1185. Args:
  1186. task_id (TaskID): ID of task.
  1187. """
  1188. with self._lock:
  1189. task = self._tasks[task_id]
  1190. current_time = self.get_time()
  1191. if task.start_time is None:
  1192. task.start_time = current_time
  1193. task.stop_time = current_time
  1194. def update(
  1195. self,
  1196. task_id: TaskID,
  1197. *,
  1198. total: Optional[float] = None,
  1199. completed: Optional[float] = None,
  1200. advance: Optional[float] = None,
  1201. description: Optional[str] = None,
  1202. visible: Optional[bool] = None,
  1203. refresh: bool = False,
  1204. **fields: Any,
  1205. ) -> None:
  1206. """Update information associated with a task.
  1207. Args:
  1208. task_id (TaskID): Task id (returned by add_task).
  1209. total (float, optional): Updates task.total if not None.
  1210. completed (float, optional): Updates task.completed if not None.
  1211. advance (float, optional): Add a value to task.completed if not None.
  1212. description (str, optional): Change task description if not None.
  1213. visible (bool, optional): Set visible flag if not None.
  1214. refresh (bool): Force a refresh of progress information. Default is False.
  1215. **fields (Any): Additional data fields required for rendering.
  1216. """
  1217. with self._lock:
  1218. task = self._tasks[task_id]
  1219. completed_start = task.completed
  1220. if total is not None and total != task.total:
  1221. task.total = total
  1222. task._reset()
  1223. if advance is not None:
  1224. task.completed += advance
  1225. if completed is not None:
  1226. task.completed = completed
  1227. if description is not None:
  1228. task.description = description
  1229. if visible is not None:
  1230. task.visible = visible
  1231. task.fields.update(fields)
  1232. update_completed = task.completed - completed_start
  1233. current_time = self.get_time()
  1234. old_sample_time = current_time - self.speed_estimate_period
  1235. _progress = task._progress
  1236. popleft = _progress.popleft
  1237. while _progress and _progress[0].timestamp < old_sample_time:
  1238. popleft()
  1239. if update_completed > 0:
  1240. _progress.append(ProgressSample(current_time, update_completed))
  1241. if (
  1242. task.total is not None
  1243. and task.completed >= task.total
  1244. and task.finished_time is None
  1245. ):
  1246. task.finished_time = task.elapsed
  1247. if refresh:
  1248. self.refresh()
  1249. def reset(
  1250. self,
  1251. task_id: TaskID,
  1252. *,
  1253. start: bool = True,
  1254. total: Optional[float] = None,
  1255. completed: int = 0,
  1256. visible: Optional[bool] = None,
  1257. description: Optional[str] = None,
  1258. **fields: Any,
  1259. ) -> None:
  1260. """Reset a task so completed is 0 and the clock is reset.
  1261. Args:
  1262. task_id (TaskID): ID of task.
  1263. start (bool, optional): Start the task after reset. Defaults to True.
  1264. total (float, optional): New total steps in task, or None to use current total. Defaults to None.
  1265. completed (int, optional): Number of steps completed. Defaults to 0.
  1266. visible (bool, optional): Enable display of the task. Defaults to True.
  1267. description (str, optional): Change task description if not None. Defaults to None.
  1268. **fields (str): Additional data fields required for rendering.
  1269. """
  1270. current_time = self.get_time()
  1271. with self._lock:
  1272. task = self._tasks[task_id]
  1273. task._reset()
  1274. task.start_time = current_time if start else None
  1275. if total is not None:
  1276. task.total = total
  1277. task.completed = completed
  1278. if visible is not None:
  1279. task.visible = visible
  1280. if fields:
  1281. task.fields = fields
  1282. if description is not None:
  1283. task.description = description
  1284. task.finished_time = None
  1285. self.refresh()
  1286. def advance(self, task_id: TaskID, advance: float = 1) -> None:
  1287. """Advance task by a number of steps.
  1288. Args:
  1289. task_id (TaskID): ID of task.
  1290. advance (float): Number of steps to advance. Default is 1.
  1291. """
  1292. current_time = self.get_time()
  1293. with self._lock:
  1294. task = self._tasks[task_id]
  1295. completed_start = task.completed
  1296. task.completed += advance
  1297. update_completed = task.completed - completed_start
  1298. old_sample_time = current_time - self.speed_estimate_period
  1299. _progress = task._progress
  1300. popleft = _progress.popleft
  1301. while _progress and _progress[0].timestamp < old_sample_time:
  1302. popleft()
  1303. while len(_progress) > 1000:
  1304. popleft()
  1305. _progress.append(ProgressSample(current_time, update_completed))
  1306. if (
  1307. task.total is not None
  1308. and task.completed >= task.total
  1309. and task.finished_time is None
  1310. ):
  1311. task.finished_time = task.elapsed
  1312. task.finished_speed = task.speed
  1313. def refresh(self) -> None:
  1314. """Refresh (render) the progress information."""
  1315. if not self.disable and self.live.is_started:
  1316. self.live.refresh()
  1317. def get_renderable(self) -> RenderableType:
  1318. """Get a renderable for the progress display."""
  1319. renderable = Group(*self.get_renderables())
  1320. return renderable
  1321. def get_renderables(self) -> Iterable[RenderableType]:
  1322. """Get a number of renderables for the progress display."""
  1323. table = self.make_tasks_table(self.tasks)
  1324. yield table
  1325. def make_tasks_table(self, tasks: Iterable[Task]) -> Table:
  1326. """Get a table to render the Progress display.
  1327. Args:
  1328. tasks (Iterable[Task]): An iterable of Task instances, one per row of the table.
  1329. Returns:
  1330. Table: A table instance.
  1331. """
  1332. table_columns = (
  1333. (
  1334. Column(no_wrap=True)
  1335. if isinstance(_column, str)
  1336. else _column.get_table_column().copy()
  1337. )
  1338. for _column in self.columns
  1339. )
  1340. table = Table.grid(*table_columns, padding=(0, 1), expand=self.expand)
  1341. for task in tasks:
  1342. if task.visible:
  1343. table.add_row(
  1344. *(
  1345. (
  1346. column.format(task=task)
  1347. if isinstance(column, str)
  1348. else column(task)
  1349. )
  1350. for column in self.columns
  1351. )
  1352. )
  1353. return table
  1354. def __rich__(self) -> RenderableType:
  1355. """Makes the Progress class itself renderable."""
  1356. with self._lock:
  1357. return self.get_renderable()
  1358. def add_task(
  1359. self,
  1360. description: str,
  1361. start: bool = True,
  1362. total: Optional[float] = 100.0,
  1363. completed: int = 0,
  1364. visible: bool = True,
  1365. **fields: Any,
  1366. ) -> TaskID:
  1367. """Add a new 'task' to the Progress display.
  1368. Args:
  1369. description (str): A description of the task.
  1370. start (bool, optional): Start the task immediately (to calculate elapsed time). If set to False,
  1371. you will need to call `start` manually. Defaults to True.
  1372. total (float, optional): Number of total steps in the progress if known.
  1373. Set to None to render a pulsing animation. Defaults to 100.
  1374. completed (int, optional): Number of steps completed so far. Defaults to 0.
  1375. visible (bool, optional): Enable display of the task. Defaults to True.
  1376. **fields (str): Additional data fields required for rendering.
  1377. Returns:
  1378. TaskID: An ID you can use when calling `update`.
  1379. """
  1380. with self._lock:
  1381. task = Task(
  1382. self._task_index,
  1383. description,
  1384. total,
  1385. completed,
  1386. visible=visible,
  1387. fields=fields,
  1388. _get_time=self.get_time,
  1389. _lock=self._lock,
  1390. )
  1391. self._tasks[self._task_index] = task
  1392. if start:
  1393. self.start_task(self._task_index)
  1394. new_task_index = self._task_index
  1395. self._task_index = TaskID(int(self._task_index) + 1)
  1396. self.refresh()
  1397. return new_task_index
  1398. def remove_task(self, task_id: TaskID) -> None:
  1399. """Delete a task if it exists.
  1400. Args:
  1401. task_id (TaskID): A task ID.
  1402. """
  1403. with self._lock:
  1404. del self._tasks[task_id]
  1405. if __name__ == "__main__": # pragma: no coverage
  1406. import random
  1407. import time
  1408. from .panel import Panel
  1409. from .rule import Rule
  1410. from .syntax import Syntax
  1411. from .table import Table
  1412. syntax = Syntax(
  1413. '''def loop_last(values: Iterable[T]) -> Iterable[Tuple[bool, T]]:
  1414. """Iterate and generate a tuple with a flag for last value."""
  1415. iter_values = iter(values)
  1416. try:
  1417. previous_value = next(iter_values)
  1418. except StopIteration:
  1419. return
  1420. for value in iter_values:
  1421. yield False, previous_value
  1422. previous_value = value
  1423. yield True, previous_value''',
  1424. "python",
  1425. line_numbers=True,
  1426. )
  1427. table = Table("foo", "bar", "baz")
  1428. table.add_row("1", "2", "3")
  1429. progress_renderables = [
  1430. "Text may be printed while the progress bars are rendering.",
  1431. Panel("In fact, [i]any[/i] renderable will work"),
  1432. "Such as [magenta]tables[/]...",
  1433. table,
  1434. "Pretty printed structures...",
  1435. {"type": "example", "text": "Pretty printed"},
  1436. "Syntax...",
  1437. syntax,
  1438. Rule("Give it a try!"),
  1439. ]
  1440. from itertools import cycle
  1441. examples = cycle(progress_renderables)
  1442. console = Console(record=True)
  1443. with Progress(
  1444. SpinnerColumn(),
  1445. *Progress.get_default_columns(),
  1446. TimeElapsedColumn(),
  1447. console=console,
  1448. transient=False,
  1449. ) as progress:
  1450. task1 = progress.add_task("[red]Downloading", total=1000)
  1451. task2 = progress.add_task("[green]Processing", total=1000)
  1452. task3 = progress.add_task("[yellow]Thinking", total=None)
  1453. while not progress.finished:
  1454. progress.update(task1, advance=0.5)
  1455. progress.update(task2, advance=0.3)
  1456. time.sleep(0.01)
  1457. if random.randint(0, 100) < 1:
  1458. progress.log(next(examples))