Coverage for python/lsst/daf/butler/core/progress.py: 22%
123 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-21 09:55 +0000
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-21 09:55 +0000
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22from __future__ import annotations
24__all__ = ("Progress", "ProgressBar", "ProgressHandler")
26import logging
27from abc import ABC, abstractmethod
28from collections.abc import Generator, Iterable, Iterator, Sized
29from contextlib import contextmanager
30from typing import ClassVar, ContextManager, Protocol, TypeVar
32_T = TypeVar("_T", covariant=True)
33_K = TypeVar("_K")
34_V = TypeVar("_V", bound=Iterable)
37class ProgressBar(Iterable[_T], Protocol):
38 """A structural interface for progress bars that wrap iterables.
40 An object conforming to this interface can be obtained from the
41 `Progress.bar` method.
43 Notes
44 -----
45 This interface is intentionally defined as the intersection of the progress
46 bar objects returned by the ``click`` and ``tqdm`` packages, allowing those
47 to directly satisfy code that uses this interface.
48 """
50 def update(self, n: int = 1) -> None:
51 """Increment the progress bar by the given amount.
53 Parameters
54 ----------
55 n : `int`, optional
56 Increment the progress bar by this many steps (defaults to ``1``).
57 Note that this is a relative increment, not an absolute progress
58 value.
59 """
60 pass
63class Progress:
64 """Public interface for reporting incremental progress in the butler and
65 related tools.
67 This class automatically creates progress bars (or not) depending on
68 whether a handle (see `ProgressHandler`) has been installed and the given
69 name and level are enabled. When progress reporting is not enabled, it
70 returns dummy objects that can be used just like progress bars by calling
71 code.
73 Parameters
74 ----------
75 name : `str`
76 Name of the process whose progress is being reported. This is in
77 general the name of a group of progress bars, not necessarily a single
78 one, and it should have the same form as a logger name.
79 level : `int`, optional
80 A `logging` level value (defaults to `logging.INFO`). Progress
81 reporting is enabled if a logger with ``name`` is enabled for this
82 level, and a `ProgressHandler` has been installed.
84 Notes
85 -----
86 The progress system inspects the level for a name using the Python built-in
87 `logging` module, and may not respect level-setting done via the
88 ``lsst.log`` interface. But while `logging` may be necessary to control
89 progress bar visibility, the progress system can still be used together
90 with either system for actual logging.
91 """
93 def __init__(self, name: str, level: int = logging.INFO) -> None:
94 self._name = name
95 self._level = level
97 # The active handler is held in a ContextVar to isolate unit tests run
98 # by pytest-xdist. If butler codes is ever used in a real multithreaded
99 # or asyncio application _and_ we want progress bars, we'll have to set
100 # up per-thread handlers or similar.
101 _active_handler: ClassVar[ProgressHandler | None] = None
103 @classmethod
104 def set_handler(cls, handler: ProgressHandler | None) -> None:
105 """Set the (global) progress handler to the given instance.
107 This should only be called in very high-level code that can be
108 reasonably confident that it will dominate its current process, e.g.
109 at the initialization of a command-line script or Jupyter notebook.
111 Parameters
112 ----------
113 handler : `ProgressHandler` or `None`
114 Object that will handle all progress reporting. May be set to
115 `None` to disable progress reporting.
116 """
117 cls._active_handler = handler
119 def is_enabled(self) -> bool:
120 """Check whether this process should report progress.
122 Returns
123 -------
124 enabled : `bool`
125 `True` if there is a `ProgressHandler` set and a logger with the
126 same name and level as ``self`` is enabled.
127 """
128 if self._active_handler is not None:
129 logger = logging.getLogger(self._name)
130 if logger.isEnabledFor(self._level):
131 return True
132 return False
134 def at(self, level: int) -> Progress:
135 """Return a copy of this progress interface with a different level.
137 Parameters
138 ----------
139 level : `int`
140 A `logging` level value. Progress reporting is enabled if a logger
141 with ``name`` is enabled for this level, and a `ProgressHandler`
142 has been installed.
144 Returns
145 -------
146 progress : `Progress`
147 A new `Progress` object with the same name as ``self`` and the
148 given ``level``.
149 """
150 return Progress(self._name, level)
152 def bar(
153 self,
154 iterable: Iterable[_T] | None = None,
155 desc: str | None = None,
156 total: int | None = None,
157 skip_scalar: bool = True,
158 ) -> ContextManager[ProgressBar[_T]]:
159 """Return a new progress bar context manager.
161 Parameters
162 ----------
163 iterable : `~collections.abc.Iterable`, optional
164 An arbitrary Python iterable that will be iterated over when the
165 returned `ProgressBar` is. If not provided, whether the progress
166 bar is iterable is handler-defined, but it may be updated manually.
167 desc: `str`, optional
168 A user-friendly description for this progress bar; usually appears
169 next to it. If not provided, ``self.name`` is used (which is not
170 usually a user-friendly string, but may be appropriate for
171 debug-level progress).
172 total : `int`, optional
173 The total number of steps in this progress bar. If not provided,
174 ``len(iterable)`` is used. If that does not work, whether the
175 progress bar works at all is handler-defined, and hence this mode
176 should not be relied upon.
177 skip_scalar: `bool`, optional
178 If `True` and ``total`` is zero or one, do not report progress.
180 Returns
181 -------
182 bar : `ContextManager` [ `ProgressBar` ]
183 A context manager that returns an object satisfying the
184 `ProgressBar` interface when it is entered.
185 """
186 if self.is_enabled():
187 if desc is None:
188 desc = self._name
189 handler = self._active_handler
190 assert handler, "Guaranteed by `is_enabled` check above."
191 if skip_scalar:
192 if total is None:
193 try:
194 # static typing says len() won't but that's why
195 # we're doing it inside a try block.
196 total = len(iterable) # type: ignore
197 except TypeError:
198 pass
199 if total is not None and total <= 1:
200 return _NullProgressBar.context(iterable)
201 return handler.get_progress_bar(iterable, desc=desc, total=total, level=self._level)
202 return _NullProgressBar.context(iterable)
204 def wrap(
205 self,
206 iterable: Iterable[_T],
207 desc: str | None = None,
208 total: int | None = None,
209 skip_scalar: bool = True,
210 ) -> Generator[_T, None, None]:
211 """Iterate over an object while reporting progress.
213 Parameters
214 ----------
215 iterable : `~collections.abc.Iterable`
216 An arbitrary Python iterable to iterate over.
217 desc: `str`, optional
218 A user-friendly description for this progress bar; usually appears
219 next to it. If not provided, ``self.name`` is used (which is not
220 usually a user-friendly string, but may be appropriate for
221 debug-level progress).
222 total : `int`, optional
223 The total number of steps in this progress bar. If not provided,
224 ``len(iterable)`` is used. If that does not work, whether the
225 progress bar works at all is handler-defined, and hence this mode
226 should not be relied upon.
227 skip_scalar: `bool`, optional
228 If `True` and ``total`` is zero or one, do not report progress.
230 Yields
231 ------
232 element
233 The same objects that iteration over ``iterable`` would yield.
234 """
235 with self.bar(iterable, desc=desc, total=total, skip_scalar=skip_scalar) as bar:
236 yield from bar
238 def iter_chunks(
239 self,
240 chunks: Iterable[_V],
241 desc: str | None = None,
242 total: int | None = None,
243 skip_scalar: bool = True,
244 ) -> Generator[_V, None, None]:
245 """Wrap iteration over chunks of elements in a progress bar.
247 Parameters
248 ----------
249 chunks : `~collections.abc.Iterable`
250 An iterable whose elements are themselves iterable.
251 desc: `str`, optional
252 A user-friendly description for this progress bar; usually appears
253 next to it. If not provided, ``self.name`` is used (which is not
254 usually a user-friendly string, but may be appropriate for
255 debug-level progress).
256 total : `int`, optional
257 The total number of steps in this progress bar; defaults to the sum
258 of the lengths of the chunks if this can be computed. If this is
259 provided or `True`, each element in ``chunks`` must be sized but
260 ``chunks`` itself need not be (and may be a single-pass iterable).
261 skip_scalar: `bool`, optional
262 If `True` and there are zero or one chunks, do not report progress.
264 Yields
265 ------
266 chunk
267 The same objects that iteration over ``chunks`` would yield.
269 Notes
270 -----
271 This attempts to display as much progress as possible given the
272 limitations of the iterables, assuming that sized iterables are also
273 multi-pass (as is true of all built-in collections and lazy iterators).
274 In detail, if ``total`` is `None`:
276 - if ``chunks`` and its elements are both sized, ``total`` is computed
277 from them and full progress is reported, and ``chunks`` must be a
278 multi-pass iterable.
279 - if ``chunks`` is sized but its elements are not, a progress bar over
280 the number of chunks is shown, and ``chunks`` must be a multi-pass
281 iterable.
282 - if ``chunks`` is not sized, the progress bar just shows when updates
283 occur.
285 If ``total`` is `True` or an integer, ``chunks`` need not be sized, but
286 its elements must be, ``chunks`` must be a multi-pass iterable, and
287 full progress is shown.
289 If ``total`` is `False`, ``chunks`` and its elements need not be sized,
290 and the progress bar just shows when updates occur.
291 """
292 if isinstance(chunks, Sized):
293 n_chunks = len(chunks)
294 else:
295 n_chunks = None
296 if total is None:
297 total = False
298 if skip_scalar and n_chunks == 1:
299 yield from chunks
300 return
301 use_n_chunks = False
302 if total is True or total is None:
303 total = 0
304 for c in chunks:
305 if total is True or isinstance(c, Sized):
306 total += len(c) # type: ignore
307 else:
308 use_n_chunks = True
309 total = None
310 break
311 if total is False:
312 total = None
313 use_n_chunks = True
314 if use_n_chunks:
315 with self.bar(desc=desc, total=n_chunks) as bar: # type: ignore
316 for chunk in chunks:
317 yield chunk
318 bar.update(1)
319 else:
320 with self.bar(desc=desc, total=total) as bar: # type: ignore
321 for chunk in chunks:
322 yield chunk
323 bar.update(len(chunk)) # type: ignore
325 def iter_item_chunks(
326 self,
327 items: Iterable[tuple[_K, _V]],
328 desc: str | None = None,
329 total: int | None = None,
330 skip_scalar: bool = True,
331 ) -> Generator[tuple[_K, _V], None, None]:
332 """Wrap iteration over chunks of items in a progress bar.
334 Parameters
335 ----------
336 items : `~collections.abc.Iterable`
337 An iterable whose elements are (key, value) tuples, where the
338 values are themselves iterable.
339 desc: `str`, optional
340 A user-friendly description for this progress bar; usually appears
341 next to it. If not provided, ``self.name`` is used (which is not
342 usually a user-friendly string, but may be appropriate for
343 debug-level progress).
344 total : `int`, optional
345 The total number of steps in this progress bar; defaults to the sum
346 of the lengths of the chunks if this can be computed. If this is
347 provided or `True`, each element in ``chunks`` must be sized but
348 ``chunks`` itself need not be (and may be a single-pass iterable).
349 skip_scalar: `bool`, optional
350 If `True` and there are zero or one items, do not report progress.
352 Yields
353 ------
354 chunk
355 The same 2-tuples that iteration over ``items`` would yield.
357 Notes
358 -----
359 This attempts to display as much progress as possible given the
360 limitations of the iterables, assuming that sized iterables are also
361 multi-pass (as is true of all built-in collections and lazy iterators).
362 In detail, if ``total`` is `None`:
364 - if ``chunks`` and its values elements are both sized, ``total`` is
365 computed from them and full progress is reported, and ``chunks`` must
366 be a multi-pass iterable.
367 - if ``chunks`` is sized but its value elements are not, a progress bar
368 over the number of chunks is shown, and ``chunks`` must be a
369 multi-pass iterable.
370 - if ``chunks`` is not sized, the progress bar just shows when updates
371 occur.
373 If ``total`` is `True` or an integer, ``chunks`` need not be sized, but
374 its value elements must be, ``chunks`` must be a multi-pass iterable,
375 and full progress is shown.
377 If ``total`` is `False`, ``chunks`` and its values elements need not be
378 sized, and the progress bar just shows when updates occur.
379 """
380 if isinstance(items, Sized):
381 n_items = len(items)
382 if skip_scalar and n_items == 1:
383 yield from items
384 return
385 else:
386 n_items = None
387 if total is None:
388 total = False
389 use_n_items = False
390 if total is True or total is None:
391 total = 0
392 for _, v in items:
393 if total is True or isinstance(v, Sized):
394 total += len(v) # type: ignore
395 else:
396 use_n_items = True
397 total = None
398 break
399 if total is False:
400 total = None
401 use_n_items = True
402 if use_n_items:
403 with self.bar(desc=desc, total=n_items) as bar: # type: ignore
404 for chunk in items:
405 yield chunk
406 bar.update(1)
407 else:
408 with self.bar(desc=desc, total=total) as bar: # type: ignore
409 for k, v in items:
410 yield k, v
411 bar.update(len(v)) # type: ignore
414class ProgressHandler(ABC):
415 """An interface for objects that can create progress bars."""
417 @abstractmethod
418 def get_progress_bar(
419 self, iterable: Iterable[_T] | None, desc: str, total: int | None, level: int
420 ) -> ContextManager[ProgressBar[_T]]:
421 """Create a new progress bar.
423 Parameters
424 ----------
425 iterable : `~collections.abc.Iterable` or `None`
426 An arbitrary Python iterable that will be iterated over when the
427 returned `ProgressBar` is. If `None`, whether the progress bar is
428 iterable is handler-defined, but it may be updated manually.
429 desc: `str`
430 A user-friendly description for this progress bar; usually appears
431 next to it
432 total : `int` or `None`
433 The total number of steps in this progress bar. If `None``,
434 ``len(iterable)`` should be used. If that does not work, whether
435 the progress bar works at all is handler-defined.
436 level : `int`
437 A `logging` level value (defaults to `logging.INFO`) associated
438 with the process reporting progress. Handlers are not responsible
439 for disabling progress reporting on levels, but may utilize level
440 information to annotate them differently.
441 """
442 raise NotImplementedError()
445class _NullProgressBar(Iterable[_T]):
446 """A trivial implementation of `ProgressBar` that does nothing but pass
447 through its iterable's elements.
449 Parameters
450 ----------
451 iterable : `~collections.abc.Iterable` or `None`
452 An arbitrary Python iterable that will be iterated over when ``self``
453 is.
454 """
456 def __init__(self, iterable: Iterable[_T] | None):
457 self._iterable = iterable
459 @classmethod
460 @contextmanager
461 def context(cls, iterable: Iterable[_T] | None) -> Generator[_NullProgressBar[_T], None, None]:
462 """Return a trivial context manager that wraps an instance of this
463 class.
465 This context manager doesn't actually do anything other than allow this
466 do-nothing implementation to be used in `Progress.bar`.
468 Parameters
469 ----------
470 iterable : `~collections.abc.Iterable` or `None`
471 An arbitrary Python iterable that will be iterated over when the
472 returned object is.
473 """
474 yield cls(iterable)
476 def __iter__(self) -> Iterator[_T]:
477 assert self._iterable is not None, "Cannot iterate over progress bar initialized without iterable."
478 return iter(self._iterable)
480 def update(self, n: int = 1) -> None:
481 pass