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