Coverage for python/lsst/daf/butler/progress.py: 22%

122 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-01 11:20 +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/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ("Progress", "ProgressBar", "ProgressHandler") 

31 

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 

38 

39_T = TypeVar("_T", covariant=True) 

40_K = TypeVar("_K") 

41_V = TypeVar("_V", bound=Iterable) 

42 

43 

44class ProgressBar(Iterable[_T], Protocol): 

45 """A structural interface for progress bars that wrap iterables. 

46 

47 An object conforming to this interface can be obtained from the 

48 `Progress.bar` method. 

49 

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 """ 

56 

57 def update(self, n: int = 1) -> None: 

58 """Increment the progress bar by the given amount. 

59 

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 

68 

69 

70class Progress: 

71 """Public interface for reporting incremental progress in the butler and 

72 related tools. 

73 

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. 

79 

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. 

90 

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 """ 

99 

100 def __init__(self, name: str, level: int = logging.INFO) -> None: 

101 self._name = name 

102 self._level = level 

103 

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 

109 

110 @classmethod 

111 def set_handler(cls, handler: ProgressHandler | None) -> None: 

112 """Set the (global) progress handler to the given instance. 

113 

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. 

117 

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 

125 

126 def is_enabled(self) -> bool: 

127 """Check whether this process should report progress. 

128 

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 

140 

141 def at(self, level: int) -> Progress: 

142 """Return a copy of this progress interface with a different level. 

143 

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. 

150 

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) 

158 

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. 

167 

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. 

186 

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) 

208 

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. 

217 

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. 

234 

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 

242 

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. 

251 

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. 

268 

269 Yields 

270 ------ 

271 chunk 

272 The same objects that iteration over ``chunks`` would yield. 

273 

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`: 

280 

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. 

289 

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. 

293 

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 

329 

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. 

338 

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. 

356 

357 Yields 

358 ------ 

359 chunk 

360 The same 2-tuples that iteration over ``items`` would yield. 

361 

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`: 

368 

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. 

377 

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. 

381 

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 

417 

418 

419class ProgressHandler(ABC): 

420 """An interface for objects that can create progress bars.""" 

421 

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. 

427 

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() 

448 

449 

450class _NullProgressBar(Iterable[_T]): 

451 """A trivial implementation of `ProgressBar` that does nothing but pass 

452 through its iterable's elements. 

453 

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 """ 

460 

461 def __init__(self, iterable: Iterable[_T] | None): 

462 self._iterable = iterable 

463 

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. 

469 

470 This context manager doesn't actually do anything other than allow this 

471 do-nothing implementation to be used in `Progress.bar`. 

472 

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. 

478 

479 Yields 

480 ------ 

481 _NullProgressBar 

482 Progress bar that does nothing. 

483 """ 

484 yield cls(iterable) 

485 

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) 

489 

490 def update(self, n: int = 1) -> None: 

491 pass