Coverage for python/lsst/daf/butler/core/progress.py: 22%
122 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-08-05 01:26 +0000
« prev ^ index » next coverage.py v7.2.7, created at 2023-08-05 01:26 +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 contextlib
27import logging
28from abc import ABC, abstractmethod
29from collections.abc import Generator, Iterable, Iterator, Sized
30from contextlib import AbstractContextManager, contextmanager
31from typing import ClassVar, Protocol, TypeVar
33_T = TypeVar("_T", covariant=True)
34_K = TypeVar("_K")
35_V = TypeVar("_V", bound=Iterable)
38class ProgressBar(Iterable[_T], Protocol):
39 """A structural interface for progress bars that wrap iterables.
41 An object conforming to this interface can be obtained from the
42 `Progress.bar` method.
44 Notes
45 -----
46 This interface is intentionally defined as the intersection of the progress
47 bar objects returned by the ``click`` and ``tqdm`` packages, allowing those
48 to directly satisfy code that uses this interface.
49 """
51 def update(self, n: int = 1) -> None:
52 """Increment the progress bar by the given amount.
54 Parameters
55 ----------
56 n : `int`, optional
57 Increment the progress bar by this many steps (defaults to ``1``).
58 Note that this is a relative increment, not an absolute progress
59 value.
60 """
61 pass
64class Progress:
65 """Public interface for reporting incremental progress in the butler and
66 related tools.
68 This class automatically creates progress bars (or not) depending on
69 whether a handle (see `ProgressHandler`) has been installed and the given
70 name and level are enabled. When progress reporting is not enabled, it
71 returns dummy objects that can be used just like progress bars by calling
72 code.
74 Parameters
75 ----------
76 name : `str`
77 Name of the process whose progress is being reported. This is in
78 general the name of a group of progress bars, not necessarily a single
79 one, and it should have the same form as a logger name.
80 level : `int`, optional
81 A `logging` level value (defaults to `logging.INFO`). Progress
82 reporting is enabled if a logger with ``name`` is enabled for this
83 level, and a `ProgressHandler` has been installed.
85 Notes
86 -----
87 The progress system inspects the level for a name using the Python built-in
88 `logging` module, and may not respect level-setting done via the
89 ``lsst.log`` interface. But while `logging` may be necessary to control
90 progress bar visibility, the progress system can still be used together
91 with either system for actual logging.
92 """
94 def __init__(self, name: str, level: int = logging.INFO) -> None:
95 self._name = name
96 self._level = level
98 # The active handler is held in a ContextVar to isolate unit tests run
99 # by pytest-xdist. If butler codes is ever used in a real multithreaded
100 # or asyncio application _and_ we want progress bars, we'll have to set
101 # up per-thread handlers or similar.
102 _active_handler: ClassVar[ProgressHandler | None] = None
104 @classmethod
105 def set_handler(cls, handler: ProgressHandler | None) -> None:
106 """Set the (global) progress handler to the given instance.
108 This should only be called in very high-level code that can be
109 reasonably confident that it will dominate its current process, e.g.
110 at the initialization of a command-line script or Jupyter notebook.
112 Parameters
113 ----------
114 handler : `ProgressHandler` or `None`
115 Object that will handle all progress reporting. May be set to
116 `None` to disable progress reporting.
117 """
118 cls._active_handler = handler
120 def is_enabled(self) -> bool:
121 """Check whether this process should report progress.
123 Returns
124 -------
125 enabled : `bool`
126 `True` if there is a `ProgressHandler` set and a logger with the
127 same name and level as ``self`` is enabled.
128 """
129 if self._active_handler is not None:
130 logger = logging.getLogger(self._name)
131 if logger.isEnabledFor(self._level):
132 return True
133 return False
135 def at(self, level: int) -> Progress:
136 """Return a copy of this progress interface with a different level.
138 Parameters
139 ----------
140 level : `int`
141 A `logging` level value. Progress reporting is enabled if a logger
142 with ``name`` is enabled for this level, and a `ProgressHandler`
143 has been installed.
145 Returns
146 -------
147 progress : `Progress`
148 A new `Progress` object with the same name as ``self`` and the
149 given ``level``.
150 """
151 return Progress(self._name, level)
153 def bar(
154 self,
155 iterable: Iterable[_T] | None = None,
156 desc: str | None = None,
157 total: int | None = None,
158 skip_scalar: bool = True,
159 ) -> AbstractContextManager[ProgressBar[_T]]:
160 """Return a new progress bar context manager.
162 Parameters
163 ----------
164 iterable : `~collections.abc.Iterable`, optional
165 An arbitrary Python iterable that will be iterated over when the
166 returned `ProgressBar` is. If not provided, whether the progress
167 bar is iterable is handler-defined, but it may be updated manually.
168 desc: `str`, optional
169 A user-friendly description for this progress bar; usually appears
170 next to it. If not provided, ``self.name`` is used (which is not
171 usually a user-friendly string, but may be appropriate for
172 debug-level progress).
173 total : `int`, optional
174 The total number of steps in this progress bar. If not provided,
175 ``len(iterable)`` is used. If that does not work, whether the
176 progress bar works at all is handler-defined, and hence this mode
177 should not be relied upon.
178 skip_scalar: `bool`, optional
179 If `True` and ``total`` is zero or one, do not report progress.
181 Returns
182 -------
183 bar : `contextlib.AbstractContextManager` [ `ProgressBar` ]
184 A context manager that returns an object satisfying the
185 `ProgressBar` interface when it is entered.
186 """
187 if self.is_enabled():
188 if desc is None:
189 desc = self._name
190 handler = self._active_handler
191 assert handler, "Guaranteed by `is_enabled` check above."
192 if skip_scalar:
193 if total is None:
194 with contextlib.suppress(TypeError):
195 # static typing says len() won't but that's why
196 # we're doing it inside a try block.
197 total = len(iterable) # type: ignore
198 if total is not None and total <= 1:
199 return _NullProgressBar.context(iterable)
200 return handler.get_progress_bar(iterable, desc=desc, total=total, level=self._level)
201 return _NullProgressBar.context(iterable)
203 def wrap(
204 self,
205 iterable: Iterable[_T],
206 desc: str | None = None,
207 total: int | None = None,
208 skip_scalar: bool = True,
209 ) -> Generator[_T, None, None]:
210 """Iterate over an object while reporting progress.
212 Parameters
213 ----------
214 iterable : `~collections.abc.Iterable`
215 An arbitrary Python iterable to iterate over.
216 desc: `str`, optional
217 A user-friendly description for this progress bar; usually appears
218 next to it. If not provided, ``self.name`` is used (which is not
219 usually a user-friendly string, but may be appropriate for
220 debug-level progress).
221 total : `int`, optional
222 The total number of steps in this progress bar. If not provided,
223 ``len(iterable)`` is used. If that does not work, whether the
224 progress bar works at all is handler-defined, and hence this mode
225 should not be relied upon.
226 skip_scalar: `bool`, optional
227 If `True` and ``total`` is zero or one, do not report progress.
229 Yields
230 ------
231 element
232 The same objects that iteration over ``iterable`` would yield.
233 """
234 with self.bar(iterable, desc=desc, total=total, skip_scalar=skip_scalar) as bar:
235 yield from bar
237 def iter_chunks(
238 self,
239 chunks: Iterable[_V],
240 desc: str | None = None,
241 total: int | None = None,
242 skip_scalar: bool = True,
243 ) -> Generator[_V, None, None]:
244 """Wrap iteration over chunks of elements in a progress bar.
246 Parameters
247 ----------
248 chunks : `~collections.abc.Iterable`
249 An iterable whose elements are themselves iterable.
250 desc: `str`, optional
251 A user-friendly description for this progress bar; usually appears
252 next to it. If not provided, ``self.name`` is used (which is not
253 usually a user-friendly string, but may be appropriate for
254 debug-level progress).
255 total : `int`, optional
256 The total number of steps in this progress bar; defaults to the sum
257 of the lengths of the chunks if this can be computed. If this is
258 provided or `True`, each element in ``chunks`` must be sized but
259 ``chunks`` itself need not be (and may be a single-pass iterable).
260 skip_scalar: `bool`, optional
261 If `True` and there are zero or one chunks, do not report progress.
263 Yields
264 ------
265 chunk
266 The same objects that iteration over ``chunks`` would yield.
268 Notes
269 -----
270 This attempts to display as much progress as possible given the
271 limitations of the iterables, assuming that sized iterables are also
272 multi-pass (as is true of all built-in collections and lazy iterators).
273 In detail, if ``total`` is `None`:
275 - if ``chunks`` and its elements are both sized, ``total`` is computed
276 from them and full progress is reported, and ``chunks`` must be a
277 multi-pass iterable.
278 - if ``chunks`` is sized but its elements are not, a progress bar over
279 the number of chunks is shown, and ``chunks`` must be a multi-pass
280 iterable.
281 - if ``chunks`` is not sized, the progress bar just shows when updates
282 occur.
284 If ``total`` is `True` or an integer, ``chunks`` need not be sized, but
285 its elements must be, ``chunks`` must be a multi-pass iterable, and
286 full progress is shown.
288 If ``total`` is `False`, ``chunks`` and its elements need not be sized,
289 and the progress bar just shows when updates occur.
290 """
291 if isinstance(chunks, Sized):
292 n_chunks = len(chunks)
293 else:
294 n_chunks = None
295 if total is None:
296 total = False
297 if skip_scalar and n_chunks == 1:
298 yield from chunks
299 return
300 use_n_chunks = False
301 if total is True or total is None:
302 total = 0
303 for c in chunks:
304 if total is True or isinstance(c, Sized):
305 total += len(c) # type: ignore
306 else:
307 use_n_chunks = True
308 total = None
309 break
310 if total is False:
311 total = None
312 use_n_chunks = True
313 if use_n_chunks:
314 with self.bar(desc=desc, total=n_chunks) as bar: # type: ignore
315 for chunk in chunks:
316 yield chunk
317 bar.update(1)
318 else:
319 with self.bar(desc=desc, total=total) as bar: # type: ignore
320 for chunk in chunks:
321 yield chunk
322 bar.update(len(chunk)) # type: ignore
324 def iter_item_chunks(
325 self,
326 items: Iterable[tuple[_K, _V]],
327 desc: str | None = None,
328 total: int | None = None,
329 skip_scalar: bool = True,
330 ) -> Generator[tuple[_K, _V], None, None]:
331 """Wrap iteration over chunks of items in a progress bar.
333 Parameters
334 ----------
335 items : `~collections.abc.Iterable`
336 An iterable whose elements are (key, value) tuples, where the
337 values are themselves iterable.
338 desc: `str`, optional
339 A user-friendly description for this progress bar; usually appears
340 next to it. If not provided, ``self.name`` is used (which is not
341 usually a user-friendly string, but may be appropriate for
342 debug-level progress).
343 total : `int`, optional
344 The total number of steps in this progress bar; defaults to the sum
345 of the lengths of the chunks if this can be computed. If this is
346 provided or `True`, each element in ``chunks`` must be sized but
347 ``chunks`` itself need not be (and may be a single-pass iterable).
348 skip_scalar: `bool`, optional
349 If `True` and there are zero or one items, do not report progress.
351 Yields
352 ------
353 chunk
354 The same 2-tuples that iteration over ``items`` would yield.
356 Notes
357 -----
358 This attempts to display as much progress as possible given the
359 limitations of the iterables, assuming that sized iterables are also
360 multi-pass (as is true of all built-in collections and lazy iterators).
361 In detail, if ``total`` is `None`:
363 - if ``chunks`` and its values elements are both sized, ``total`` is
364 computed from them and full progress is reported, and ``chunks`` must
365 be a multi-pass iterable.
366 - if ``chunks`` is sized but its value elements are not, a progress bar
367 over the number of chunks is shown, and ``chunks`` must be a
368 multi-pass iterable.
369 - if ``chunks`` is not sized, the progress bar just shows when updates
370 occur.
372 If ``total`` is `True` or an integer, ``chunks`` need not be sized, but
373 its value elements must be, ``chunks`` must be a multi-pass iterable,
374 and full progress is shown.
376 If ``total`` is `False`, ``chunks`` and its values elements need not be
377 sized, and the progress bar just shows when updates occur.
378 """
379 if isinstance(items, Sized):
380 n_items = len(items)
381 if skip_scalar and n_items == 1:
382 yield from items
383 return
384 else:
385 n_items = None
386 if total is None:
387 total = False
388 use_n_items = False
389 if total is True or total is None:
390 total = 0
391 for _, v in items:
392 if total is True or isinstance(v, Sized):
393 total += len(v) # type: ignore
394 else:
395 use_n_items = True
396 total = None
397 break
398 if total is False:
399 total = None
400 use_n_items = True
401 if use_n_items:
402 with self.bar(desc=desc, total=n_items) as bar: # type: ignore
403 for chunk in items:
404 yield chunk
405 bar.update(1)
406 else:
407 with self.bar(desc=desc, total=total) as bar: # type: ignore
408 for k, v in items:
409 yield k, v
410 bar.update(len(v)) # type: ignore
413class ProgressHandler(ABC):
414 """An interface for objects that can create progress bars."""
416 @abstractmethod
417 def get_progress_bar(
418 self, iterable: Iterable[_T] | None, desc: str, total: int | None, level: int
419 ) -> AbstractContextManager[ProgressBar[_T]]:
420 """Create a new progress bar.
422 Parameters
423 ----------
424 iterable : `~collections.abc.Iterable` or `None`
425 An arbitrary Python iterable that will be iterated over when the
426 returned `ProgressBar` is. If `None`, whether the progress bar is
427 iterable is handler-defined, but it may be updated manually.
428 desc: `str`
429 A user-friendly description for this progress bar; usually appears
430 next to it
431 total : `int` or `None`
432 The total number of steps in this progress bar. If `None``,
433 ``len(iterable)`` should be used. If that does not work, whether
434 the progress bar works at all is handler-defined.
435 level : `int`
436 A `logging` level value (defaults to `logging.INFO`) associated
437 with the process reporting progress. Handlers are not responsible
438 for disabling progress reporting on levels, but may utilize level
439 information to annotate them differently.
440 """
441 raise NotImplementedError()
444class _NullProgressBar(Iterable[_T]):
445 """A trivial implementation of `ProgressBar` that does nothing but pass
446 through its iterable's elements.
448 Parameters
449 ----------
450 iterable : `~collections.abc.Iterable` or `None`
451 An arbitrary Python iterable that will be iterated over when ``self``
452 is.
453 """
455 def __init__(self, iterable: Iterable[_T] | None):
456 self._iterable = iterable
458 @classmethod
459 @contextmanager
460 def context(cls, iterable: Iterable[_T] | None) -> Generator[_NullProgressBar[_T], None, None]:
461 """Return a trivial context manager that wraps an instance of this
462 class.
464 This context manager doesn't actually do anything other than allow this
465 do-nothing implementation to be used in `Progress.bar`.
467 Parameters
468 ----------
469 iterable : `~collections.abc.Iterable` or `None`
470 An arbitrary Python iterable that will be iterated over when the
471 returned object is.
472 """
473 yield cls(iterable)
475 def __iter__(self) -> Iterator[_T]:
476 assert self._iterable is not None, "Cannot iterate over progress bar initialized without iterable."
477 return iter(self._iterable)
479 def update(self, n: int = 1) -> None:
480 pass