Coverage for python/lsst/daf/butler/core/progress.py: 34%
81 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-14 19:21 +0000
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-14 19:21 +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 Collection, Generator, Iterable, Iterator
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=Collection)
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: Collection[_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.Collection`
250 A sized iterable whose elements are themselves both iterable and
251 sized (i.e. ``len(item)`` works). If ``total`` is not provided,
252 this may not be a single-pass iteration, because an initial pass to
253 estimate the total number of elements is required.
254 desc: `str`, optional
255 A user-friendly description for this progress bar; usually appears
256 next to it. If not provided, ``self.name`` is used (which is not
257 usually a user-friendly string, but may be appropriate for
258 debug-level progress).
259 total : `int`, optional
260 The total number of steps in this progress bar; defaults to the
261 sum of the lengths of the chunks.
262 skip_scalar: `bool`, optional
263 If `True` and there are zero or one chunks, do not report progress.
265 Yields
266 ------
267 chunk
268 The same objects that iteration over ``chunks`` would yield.
269 """
270 if skip_scalar and len(chunks) <= 1:
271 yield from chunks
272 else:
273 if total is None:
274 total = sum(len(c) for c in chunks)
275 with self.bar(desc=desc, total=total) as bar: # type: ignore
276 for chunk in chunks:
277 yield chunk
278 bar.update(len(chunk))
280 def iter_item_chunks(
281 self,
282 items: Collection[tuple[_K, _V]],
283 desc: str | None = None,
284 total: int | None = None,
285 skip_scalar: bool = True,
286 ) -> Generator[tuple[_K, _V], None, None]:
287 """Wrap iteration over chunks of items in a progress bar.
289 Parameters
290 ----------
291 items : `~collections.abc.Iterable`
292 A sized iterable whose elements are (key, value) tuples, where the
293 values are themselves both iterable and sized (i.e. ``len(item)``
294 works). If ``total`` is not provided, this may not be a
295 single-pass iteration, because an initial pass to estimate the
296 total number of elements is required.
297 desc: `str`, optional
298 A user-friendly description for this progress bar; usually appears
299 next to it. If not provided, ``self.name`` is used (which is not
300 usually a user-friendly string, but may be appropriate for
301 debug-level progress).
302 total : `int`, optional
303 The total number of values in this progress bar; defaults to the
304 sum of the lengths of the chunks.
305 skip_scalar: `bool`, optional
306 If `True` and there are zero or one items, do not report progress.
308 Yields
309 ------
310 chunk
311 The same items that iteration over ``chunks`` would yield.
312 """
313 if skip_scalar and len(items) <= 1:
314 yield from items
315 else:
316 if total is None:
317 total = sum(len(v) for _, v in items)
318 with self.bar(desc=desc, total=total) as bar: # type: ignore
319 for key, chunk in items:
320 yield key, chunk
321 bar.update(len(chunk))
324class ProgressHandler(ABC):
325 """An interface for objects that can create progress bars."""
327 @abstractmethod
328 def get_progress_bar(
329 self, iterable: Iterable[_T] | None, desc: str, total: int | None, level: int
330 ) -> ContextManager[ProgressBar[_T]]:
331 """Create a new progress bar.
333 Parameters
334 ----------
335 iterable : `~collections.abc.Iterable` or `None`
336 An arbitrary Python iterable that will be iterated over when the
337 returned `ProgressBar` is. If `None`, whether the progress bar is
338 iterable is handler-defined, but it may be updated manually.
339 desc: `str`
340 A user-friendly description for this progress bar; usually appears
341 next to it
342 total : `int` or `None`
343 The total number of steps in this progress bar. If `None``,
344 ``len(iterable)`` should be used. If that does not work, whether
345 the progress bar works at all is handler-defined.
346 level : `int`
347 A `logging` level value (defaults to `logging.INFO`) associated
348 with the process reporting progress. Handlers are not responsible
349 for disabling progress reporting on levels, but may utilize level
350 information to annotate them differently.
351 """
352 raise NotImplementedError()
355class _NullProgressBar(Iterable[_T]):
356 """A trivial implementation of `ProgressBar` that does nothing but pass
357 through its iterable's elements.
359 Parameters
360 ----------
361 iterable : `~collections.abc.Iterable` or `None`
362 An arbitrary Python iterable that will be iterated over when ``self``
363 is.
364 """
366 def __init__(self, iterable: Iterable[_T] | None):
367 self._iterable = iterable
369 @classmethod
370 @contextmanager
371 def context(cls, iterable: Iterable[_T] | None) -> Generator[_NullProgressBar[_T], None, None]:
372 """Return a trivial context manager that wraps an instance of this
373 class.
375 This context manager doesn't actually do anything other than allow this
376 do-nothing implementation to be used in `Progress.bar`.
378 Parameters
379 ----------
380 iterable : `~collections.abc.Iterable` or `None`
381 An arbitrary Python iterable that will be iterated over when the
382 returned object is.
383 """
384 yield cls(iterable)
386 def __iter__(self) -> Iterator[_T]:
387 assert self._iterable is not None, "Cannot iterate over progress bar initialized without iterable."
388 return iter(self._iterable)
390 def update(self, n: int = 1) -> None:
391 pass