notebook.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321
  1. """
  2. IPython/Jupyter Notebook progressbar decorator for iterators.
  3. Includes a default `range` iterator printing to `stderr`.
  4. Usage:
  5. >>> from tqdm.notebook import trange, tqdm
  6. >>> for i in trange(10):
  7. ... ...
  8. """
  9. # import compatibility functions and utilities
  10. import re
  11. import sys
  12. from weakref import proxy
  13. # to inherit from the tqdm class
  14. from .std import tqdm as std_tqdm
  15. if True: # pragma: no cover
  16. # import IPython/Jupyter base widget and display utilities
  17. IPY = 0
  18. try: # IPython 4.x
  19. import ipywidgets
  20. IPY = 4
  21. except ImportError: # IPython 3.x / 2.x
  22. IPY = 32
  23. import warnings
  24. with warnings.catch_warnings():
  25. warnings.filterwarnings(
  26. 'ignore', message=".*The `IPython.html` package has been deprecated.*")
  27. try:
  28. import IPython.html.widgets as ipywidgets # NOQA: F401
  29. except ImportError:
  30. pass
  31. try: # IPython 4.x / 3.x
  32. if IPY == 32:
  33. from IPython.html.widgets import HTML
  34. from IPython.html.widgets import FloatProgress as IProgress
  35. from IPython.html.widgets import HBox
  36. IPY = 3
  37. else:
  38. from ipywidgets import HTML
  39. from ipywidgets import FloatProgress as IProgress
  40. from ipywidgets import HBox
  41. except ImportError:
  42. try: # IPython 2.x
  43. from IPython.html.widgets import HTML
  44. from IPython.html.widgets import ContainerWidget as HBox
  45. from IPython.html.widgets import FloatProgressWidget as IProgress
  46. IPY = 2
  47. except ImportError:
  48. IPY = 0
  49. IProgress = None
  50. HBox = object
  51. try:
  52. from IPython.display import display # , clear_output
  53. except ImportError:
  54. pass
  55. # HTML encoding
  56. try: # Py3
  57. from html import escape
  58. except ImportError: # Py2
  59. from cgi import escape
  60. __author__ = {"github.com/": ["lrq3000", "casperdcl", "alexanderkuk"]}
  61. __all__ = ['tqdm_notebook', 'tnrange', 'tqdm', 'trange']
  62. WARN_NOIPYW = ("IProgress not found. Please update jupyter and ipywidgets."
  63. " See https://ipywidgets.readthedocs.io/en/stable"
  64. "/user_install.html")
  65. class TqdmHBox(HBox):
  66. """`ipywidgets.HBox` with a pretty representation"""
  67. def _json_(self, pretty=None):
  68. pbar = getattr(self, 'pbar', None)
  69. if pbar is None:
  70. return {}
  71. d = pbar.format_dict
  72. if pretty is not None:
  73. d["ascii"] = not pretty
  74. return d
  75. def __repr__(self, pretty=False):
  76. pbar = getattr(self, 'pbar', None)
  77. if pbar is None:
  78. return super(TqdmHBox, self).__repr__()
  79. return pbar.format_meter(**self._json_(pretty))
  80. def _repr_pretty_(self, pp, *_, **__):
  81. pp.text(self.__repr__(True))
  82. class tqdm_notebook(std_tqdm):
  83. """
  84. Experimental IPython/Jupyter Notebook widget using tqdm!
  85. """
  86. @staticmethod
  87. def status_printer(_, total=None, desc=None, ncols=None):
  88. """
  89. Manage the printing of an IPython/Jupyter Notebook progress bar widget.
  90. """
  91. # Fallback to text bar if there's no total
  92. # DEPRECATED: replaced with an 'info' style bar
  93. # if not total:
  94. # return super(tqdm_notebook, tqdm_notebook).status_printer(file)
  95. # fp = file
  96. # Prepare IPython progress bar
  97. if IProgress is None: # #187 #451 #558 #872
  98. raise ImportError(WARN_NOIPYW)
  99. if total:
  100. pbar = IProgress(min=0, max=total)
  101. else: # No total? Show info style bar with no progress tqdm status
  102. pbar = IProgress(min=0, max=1)
  103. pbar.value = 1
  104. pbar.bar_style = 'info'
  105. if ncols is None:
  106. pbar.layout.width = "20px"
  107. ltext = HTML()
  108. rtext = HTML()
  109. if desc:
  110. ltext.value = desc
  111. container = TqdmHBox(children=[ltext, pbar, rtext])
  112. # Prepare layout
  113. if ncols is not None: # use default style of ipywidgets
  114. # ncols could be 100, "100px", "100%"
  115. ncols = str(ncols) # ipywidgets only accepts string
  116. try:
  117. if int(ncols) > 0: # isnumeric and positive
  118. ncols += 'px'
  119. except ValueError:
  120. pass
  121. pbar.layout.flex = '2'
  122. container.layout.width = ncols
  123. container.layout.display = 'inline-flex'
  124. container.layout.flex_flow = 'row wrap'
  125. return container
  126. def display(self, msg=None, pos=None,
  127. # additional signals
  128. close=False, bar_style=None, check_delay=True):
  129. # Note: contrary to native tqdm, msg='' does NOT clear bar
  130. # goal is to keep all infos if error happens so user knows
  131. # at which iteration the loop failed.
  132. # Clear previous output (really necessary?)
  133. # clear_output(wait=1)
  134. if not msg and not close:
  135. d = self.format_dict
  136. # remove {bar}
  137. d['bar_format'] = (d['bar_format'] or "{l_bar}<bar/>{r_bar}").replace(
  138. "{bar}", "<bar/>")
  139. msg = self.format_meter(**d)
  140. ltext, pbar, rtext = self.container.children
  141. pbar.value = self.n
  142. if msg:
  143. # html escape special characters (like '&')
  144. if '<bar/>' in msg:
  145. left, right = map(escape, re.split(r'\|?<bar/>\|?', msg, 1))
  146. else:
  147. left, right = '', escape(msg)
  148. # Update description
  149. ltext.value = left
  150. # never clear the bar (signal: msg='')
  151. if right:
  152. rtext.value = right
  153. # Change bar style
  154. if bar_style:
  155. # Hack-ish way to avoid the danger bar_style being overridden by
  156. # success because the bar gets closed after the error...
  157. if pbar.bar_style != 'danger' or bar_style != 'success':
  158. pbar.bar_style = bar_style
  159. # Special signal to close the bar
  160. if close and pbar.bar_style != 'danger': # hide only if no error
  161. try:
  162. self.container.close()
  163. except AttributeError:
  164. self.container.visible = False
  165. self.container.layout.visibility = 'hidden' # IPYW>=8
  166. if check_delay and self.delay > 0 and not self.displayed:
  167. display(self.container)
  168. self.displayed = True
  169. @property
  170. def colour(self):
  171. if hasattr(self, 'container'):
  172. return self.container.children[-2].style.bar_color
  173. @colour.setter
  174. def colour(self, bar_color):
  175. if hasattr(self, 'container'):
  176. self.container.children[-2].style.bar_color = bar_color
  177. def __init__(self, *args, **kwargs):
  178. """
  179. Supports the usual `tqdm.tqdm` parameters as well as those listed below.
  180. Parameters
  181. ----------
  182. display : Whether to call `display(self.container)` immediately
  183. [default: True].
  184. """
  185. kwargs = kwargs.copy()
  186. # Setup default output
  187. file_kwarg = kwargs.get('file', sys.stderr)
  188. if file_kwarg is sys.stderr or file_kwarg is None:
  189. kwargs['file'] = sys.stdout # avoid the red block in IPython
  190. # Initialize parent class + avoid printing by using gui=True
  191. kwargs['gui'] = True
  192. # convert disable = None to False
  193. kwargs['disable'] = bool(kwargs.get('disable', False))
  194. colour = kwargs.pop('colour', None)
  195. display_here = kwargs.pop('display', True)
  196. super(tqdm_notebook, self).__init__(*args, **kwargs)
  197. if self.disable or not kwargs['gui']:
  198. self.disp = lambda *_, **__: None
  199. return
  200. # Get bar width
  201. self.ncols = '100%' if self.dynamic_ncols else kwargs.get("ncols", None)
  202. # Replace with IPython progress bar display (with correct total)
  203. unit_scale = 1 if self.unit_scale is True else self.unit_scale or 1
  204. total = self.total * unit_scale if self.total else self.total
  205. self.container = self.status_printer(self.fp, total, self.desc, self.ncols)
  206. self.container.pbar = proxy(self)
  207. self.displayed = False
  208. if display_here and self.delay <= 0:
  209. display(self.container)
  210. self.displayed = True
  211. self.disp = self.display
  212. self.colour = colour
  213. # Print initial bar state
  214. if not self.disable:
  215. self.display(check_delay=False)
  216. def __iter__(self):
  217. try:
  218. it = super(tqdm_notebook, self).__iter__()
  219. for obj in it:
  220. # return super(tqdm...) will not catch exception
  221. yield obj
  222. # NB: except ... [ as ...] breaks IPython async KeyboardInterrupt
  223. except: # NOQA
  224. self.disp(bar_style='danger')
  225. raise
  226. # NB: don't `finally: close()`
  227. # since this could be a shared bar which the user will `reset()`
  228. def update(self, n=1):
  229. try:
  230. return super(tqdm_notebook, self).update(n=n)
  231. # NB: except ... [ as ...] breaks IPython async KeyboardInterrupt
  232. except: # NOQA
  233. # cannot catch KeyboardInterrupt when using manual tqdm
  234. # as the interrupt will most likely happen on another statement
  235. self.disp(bar_style='danger')
  236. raise
  237. # NB: don't `finally: close()`
  238. # since this could be a shared bar which the user will `reset()`
  239. def close(self):
  240. if self.disable:
  241. return
  242. super(tqdm_notebook, self).close()
  243. # Try to detect if there was an error or KeyboardInterrupt
  244. # in manual mode: if n < total, things probably got wrong
  245. if self.total and self.n < self.total:
  246. self.disp(bar_style='danger', check_delay=False)
  247. else:
  248. if self.leave:
  249. self.disp(bar_style='success', check_delay=False)
  250. else:
  251. self.disp(close=True, check_delay=False)
  252. def clear(self, *_, **__):
  253. pass
  254. def reset(self, total=None):
  255. """
  256. Resets to 0 iterations for repeated use.
  257. Consider combining with `leave=True`.
  258. Parameters
  259. ----------
  260. total : int or float, optional. Total to use for the new bar.
  261. """
  262. if self.disable:
  263. return super(tqdm_notebook, self).reset(total=total)
  264. _, pbar, _ = self.container.children
  265. pbar.bar_style = ''
  266. if total is not None:
  267. pbar.max = total
  268. if not self.total and self.ncols is None: # no longer unknown total
  269. pbar.layout.width = None # reset width
  270. return super(tqdm_notebook, self).reset(total=total)
  271. def tnrange(*args, **kwargs):
  272. """Shortcut for `tqdm.notebook.tqdm(range(*args), **kwargs)`."""
  273. return tqdm_notebook(range(*args), **kwargs)
  274. # Aliases
  275. tqdm = tqdm_notebook
  276. trange = tnrange