Coverage for python/lsst/daf/butler/core/progress.py: 27%
81 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-07 02:05 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-07 02:05 -0800
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 contextlib import contextmanager
29from typing import (
30 ClassVar,
31 Collection,
32 ContextManager,
33 Generator,
34 Iterable,
35 Iterator,
36 Optional,
37 Protocol,
38 Tuple,
39 TypeVar,
40)
42_T = TypeVar("_T", covariant=True)
43_K = TypeVar("_K")
44_V = TypeVar("_V", bound=Collection)
47class ProgressBar(Iterable[_T], Protocol):
48 """A structural interface for progress bars that wrap iterables.
50 An object conforming to this interface can be obtained from the
51 `Progress.bar` method.
53 Notes
54 -----
55 This interface is intentionally defined as the intersection of the progress
56 bar objects returned by the ``click`` and ``tqdm`` packages, allowing those
57 to directly satisfy code that uses this interface.
58 """
60 def update(self, n: int = 1) -> None:
61 """Increment the progress bar by the given amount.
63 Parameters
64 ----------
65 n : `int`, optional
66 Increment the progress bar by this many steps (defaults to ``1``).
67 Note that this is a relative increment, not an absolute progress
68 value.
69 """
70 pass
73class Progress:
74 """Public interface for reporting incremental progress in the butler and
75 related tools.
77 This class automatically creates progress bars (or not) depending on
78 whether a handle (see `ProgressHandler`) has been installed and the given
79 name and level are enabled. When progress reporting is not enabled, it
80 returns dummy objects that can be used just like progress bars by calling
81 code.
83 Parameters
84 ----------
85 name : `str`
86 Name of the process whose progress is being reported. This is in
87 general the name of a group of progress bars, not necessarily a single
88 one, and it should have the same form as a logger name.
89 level : `int`, optional
90 A `logging` level value (defaults to `logging.INFO`). Progress
91 reporting is enabled if a logger with ``name`` is enabled for this
92 level, and a `ProgressHandler` has been installed.
94 Notes
95 -----
96 The progress system inspects the level for a name using the Python built-in
97 `logging` module, and may not respect level-setting done via the
98 ``lsst.log`` interface. But while `logging` may be necessary to control
99 progress bar visibility, the progress system can still be used together
100 with either system for actual logging.
101 """
103 def __init__(self, name: str, level: int = logging.INFO) -> None:
104 self._name = name
105 self._level = level
107 # The active handler is held in a ContextVar to isolate unit tests run
108 # by pytest-xdist. If butler codes is ever used in a real multithreaded
109 # or asyncio application _and_ we want progress bars, we'll have to set
110 # up per-thread handlers or similar.
111 _active_handler: ClassVar[Optional[ProgressHandler]] = None
113 @classmethod
114 def set_handler(cls, handler: Optional[ProgressHandler]) -> None:
115 """Set the (global) progress handler to the given instance.
117 This should only be called in very high-level code that can be
118 reasonably confident that it will dominate its current process, e.g.
119 at the initialization of a command-line script or Jupyter notebook.
121 Parameters
122 ----------
123 handler : `ProgressHandler` or `None`
124 Object that will handle all progress reporting. May be set to
125 `None` to disable progress reporting.
126 """
127 cls._active_handler = handler
129 def is_enabled(self) -> bool:
130 """Check whether this process should report progress.
132 Returns
133 -------
134 enabled : `bool`
135 `True` if there is a `ProgressHandler` set and a logger with the
136 same name and level as ``self`` is enabled.
137 """
138 if self._active_handler is not None:
139 logger = logging.getLogger(self._name)
140 if logger.isEnabledFor(self._level):
141 return True
142 return False
144 def at(self, level: int) -> Progress:
145 """Return a copy of this progress interface with a different level.
147 Parameters
148 ----------
149 level : `int`
150 A `logging` level value. Progress reporting is enabled if a logger
151 with ``name`` is enabled for this level, and a `ProgressHandler`
152 has been installed.
154 Returns
155 -------
156 progress : `Progress`
157 A new `Progress` object with the same name as ``self`` and the
158 given ``level``.
159 """
160 return Progress(self._name, level)
162 def bar(
163 self,
164 iterable: Optional[Iterable[_T]] = None,
165 desc: Optional[str] = None,
166 total: Optional[int] = None,
167 skip_scalar: bool = True,
168 ) -> ContextManager[ProgressBar[_T]]:
169 """Return a new progress bar context manager.
171 Parameters
172 ----------
173 iterable : `Iterable`, optional
174 An arbitrary Python iterable that will be iterated over when the
175 returned `ProgressBar` is. If not provided, whether the progress
176 bar is iterable is handler-defined, but it may be updated manually.
177 desc: `str`, optional
178 A user-friendly description for this progress bar; usually appears
179 next to it. If not provided, ``self.name`` is used (which is not
180 usually a user-friendly string, but may be appropriate for
181 debug-level progress).
182 total : `int`, optional
183 The total number of steps in this progress bar. If not provided,
184 ``len(iterable)`` is used. If that does not work, whether the
185 progress bar works at all is handler-defined, and hence this mode
186 should not be relied upon.
187 skip_scalar: `bool`, optional
188 If `True` and ``total`` is zero or one, do not report progress.
190 Returns
191 -------
192 bar : `ContextManager` [ `ProgressBar` ]
193 A context manager that returns an object satisfying the
194 `ProgressBar` interface when it is entered.
195 """
196 if self.is_enabled():
197 if desc is None:
198 desc = self._name
199 handler = self._active_handler
200 assert handler, "Guaranteed by `is_enabled` check above."
201 if skip_scalar:
202 if total is None:
203 try:
204 # static typing says len() won't but that's why
205 # we're doing it inside a try block.
206 total = len(iterable) # type: ignore
207 except TypeError:
208 pass
209 if total is not None and total <= 1:
210 return _NullProgressBar.context(iterable)
211 return handler.get_progress_bar(iterable, desc=desc, total=total, level=self._level)
212 return _NullProgressBar.context(iterable)
214 def wrap(
215 self,
216 iterable: Iterable[_T],
217 desc: Optional[str] = None,
218 total: Optional[int] = None,
219 skip_scalar: bool = True,
220 ) -> Generator[_T, None, None]:
221 """Iterate over an object while reporting progress.
223 Parameters
224 ----------
225 iterable : `Iterable`
226 An arbitrary Python iterable to iterate over.
227 desc: `str`, optional
228 A user-friendly description for this progress bar; usually appears
229 next to it. If not provided, ``self.name`` is used (which is not
230 usually a user-friendly string, but may be appropriate for
231 debug-level progress).
232 total : `int`, optional
233 The total number of steps in this progress bar. If not provided,
234 ``len(iterable)`` is used. If that does not work, whether the
235 progress bar works at all is handler-defined, and hence this mode
236 should not be relied upon.
237 skip_scalar: `bool`, optional
238 If `True` and ``total`` is zero or one, do not report progress.
240 Yields
241 ------
242 element
243 The same objects that iteration over ``iterable`` would yield.
244 """
245 with self.bar(iterable, desc=desc, total=total, skip_scalar=skip_scalar) as bar:
246 yield from bar
248 def iter_chunks(
249 self,
250 chunks: Collection[_V],
251 desc: Optional[str] = None,
252 total: Optional[int] = None,
253 skip_scalar: bool = True,
254 ) -> Generator[_V, None, None]:
255 """Wrap iteration over chunks of elements in a progress bar.
257 Parameters
258 ----------
259 chunks : `Collection`
260 A sized iterable whose elements are themselves both iterable and
261 sized (i.e. ``len(item)`` works). If ``total`` is not provided,
262 this may not be a single-pass iteration, because an initial pass to
263 estimate the total number of elements is required.
264 desc: `str`, optional
265 A user-friendly description for this progress bar; usually appears
266 next to it. If not provided, ``self.name`` is used (which is not
267 usually a user-friendly string, but may be appropriate for
268 debug-level progress).
269 total : `int`, optional
270 The total number of steps in this progress bar; defaults to the
271 sum of the lengths of the chunks.
272 skip_scalar: `bool`, optional
273 If `True` and there are zero or one chunks, do not report progress.
275 Yields
276 ------
277 chunk
278 The same objects that iteration over ``chunks`` would yield.
279 """
280 if skip_scalar and len(chunks) <= 1:
281 yield from chunks
282 else:
283 if total is None:
284 total = sum(len(c) for c in chunks)
285 with self.bar(desc=desc, total=total) as bar: # type: ignore
286 for chunk in chunks:
287 yield chunk
288 bar.update(len(chunk))
290 def iter_item_chunks(
291 self,
292 items: Collection[Tuple[_K, _V]],
293 desc: Optional[str] = None,
294 total: Optional[int] = None,
295 skip_scalar: bool = True,
296 ) -> Generator[Tuple[_K, _V], None, None]:
297 """Wrap iteration over chunks of items in a progress bar.
299 Parameters
300 ----------
301 items : `Iterable`
302 A sized iterable whose elements are (key, value) tuples, where the
303 values are themselves both iterable and sized (i.e. ``len(item)``
304 works). If ``total`` is not provided, this may not be a
305 single-pass iteration, because an initial pass to estimate the
306 total number of elements is required.
307 desc: `str`, optional
308 A user-friendly description for this progress bar; usually appears
309 next to it. If not provided, ``self.name`` is used (which is not
310 usually a user-friendly string, but may be appropriate for
311 debug-level progress).
312 total : `int`, optional
313 The total number of values in this progress bar; defaults to the
314 sum of the lengths of the chunks.
315 skip_scalar: `bool`, optional
316 If `True` and there are zero or one items, do not report progress.
318 Yields
319 ------
320 chunk
321 The same items that iteration over ``chunks`` would yield.
322 """
323 if skip_scalar and len(items) <= 1:
324 yield from items
325 else:
326 if total is None:
327 total = sum(len(v) for _, v in items)
328 with self.bar(desc=desc, total=total) as bar: # type: ignore
329 for key, chunk in items:
330 yield key, chunk
331 bar.update(len(chunk))
334class ProgressHandler(ABC):
335 """An interface for objects that can create progress bars."""
337 @abstractmethod
338 def get_progress_bar(
339 self, iterable: Optional[Iterable[_T]], desc: str, total: Optional[int], level: int
340 ) -> ContextManager[ProgressBar[_T]]:
341 """Create a new progress bar.
343 Parameters
344 ----------
345 iterable : `Iterable` or `None`
346 An arbitrary Python iterable that will be iterated over when the
347 returned `ProgressBar` is. If `None`, whether the progress bar is
348 iterable is handler-defined, but it may be updated manually.
349 desc: `str`
350 A user-friendly description for this progress bar; usually appears
351 next to it
352 total : `int` or `None`
353 The total number of steps in this progress bar. If `None``,
354 ``len(iterable)`` should be used. If that does not work, whether
355 the progress bar works at all is handler-defined.
356 level : `int`
357 A `logging` level value (defaults to `logging.INFO`) associated
358 with the process reporting progress. Handlers are not responsible
359 for disabling progress reporting on levels, but may utilize level
360 information to annotate them differently.
361 """
362 raise NotImplementedError()
365class _NullProgressBar(Iterable[_T]):
366 """A trivial implementation of `ProgressBar` that does nothing but pass
367 through its iterable's elements.
369 Parameters
370 ----------
371 iterable : `Iterable` or `None`
372 An arbitrary Python iterable that will be iterated over when ``self``
373 is.
374 """
376 def __init__(self, iterable: Optional[Iterable[_T]]):
377 self._iterable = iterable
379 @classmethod
380 @contextmanager
381 def context(cls, iterable: Optional[Iterable[_T]]) -> Generator[_NullProgressBar[_T], None, None]:
382 """Return a trivial context manager that wraps an instance of this
383 class.
385 This context manager doesn't actually do anything other than allow this
386 do-nothing implementation to be used in `Progress.bar`.
388 Parameters
389 ----------
390 iterable : `Iterable` or `None`
391 An arbitrary Python iterable that will be iterated over when the
392 returned object is.
393 """
394 yield cls(iterable)
396 def __iter__(self) -> Iterator[_T]:
397 assert self._iterable is not None, "Cannot iterate over progress bar initialized without iterable."
398 return iter(self._iterable)
400 def update(self, n: int = 1) -> None:
401 pass