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

123 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-21 09:55 +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/>. 

21 

22from __future__ import annotations 

23 

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

25 

26import logging 

27from abc import ABC, abstractmethod 

28from collections.abc import Generator, Iterable, Iterator, Sized 

29from contextlib import contextmanager 

30from typing import ClassVar, ContextManager, Protocol, TypeVar 

31 

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

33_K = TypeVar("_K") 

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

35 

36 

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

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

39 

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

41 `Progress.bar` method. 

42 

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

49 

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

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

52 

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 

61 

62 

63class Progress: 

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

65 related tools. 

66 

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. 

72 

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. 

83 

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

92 

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

94 self._name = name 

95 self._level = level 

96 

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 

102 

103 @classmethod 

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

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

106 

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. 

110 

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 

118 

119 def is_enabled(self) -> bool: 

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

121 

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 

133 

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

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

136 

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. 

143 

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) 

151 

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. 

160 

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. 

179 

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) 

203 

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. 

212 

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. 

229 

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 

237 

238 def iter_chunks( 

239 self, 

240 chunks: Iterable[_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. 

246 

247 Parameters 

248 ---------- 

249 chunks : `~collections.abc.Iterable` 

250 An iterable whose elements are themselves iterable. 

251 desc: `str`, optional 

252 A user-friendly description for this progress bar; usually appears 

253 next to it. If not provided, ``self.name`` is used (which is not 

254 usually a user-friendly string, but may be appropriate for 

255 debug-level progress). 

256 total : `int`, optional 

257 The total number of steps in this progress bar; defaults to the sum 

258 of the lengths of the chunks if this can be computed. If this is 

259 provided or `True`, each element in ``chunks`` must be sized but 

260 ``chunks`` itself need not be (and may be a single-pass iterable). 

261 skip_scalar: `bool`, optional 

262 If `True` and there are zero or one chunks, do not report progress. 

263 

264 Yields 

265 ------ 

266 chunk 

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

268 

269 Notes 

270 ----- 

271 This attempts to display as much progress as possible given the 

272 limitations of the iterables, assuming that sized iterables are also 

273 multi-pass (as is true of all built-in collections and lazy iterators). 

274 In detail, if ``total`` is `None`: 

275 

276 - if ``chunks`` and its elements are both sized, ``total`` is computed 

277 from them and full progress is reported, and ``chunks`` must be a 

278 multi-pass iterable. 

279 - if ``chunks`` is sized but its elements are not, a progress bar over 

280 the number of chunks is shown, and ``chunks`` must be a multi-pass 

281 iterable. 

282 - if ``chunks`` is not sized, the progress bar just shows when updates 

283 occur. 

284 

285 If ``total`` is `True` or an integer, ``chunks`` need not be sized, but 

286 its elements must be, ``chunks`` must be a multi-pass iterable, and 

287 full progress is shown. 

288 

289 If ``total`` is `False`, ``chunks`` and its elements need not be sized, 

290 and the progress bar just shows when updates occur. 

291 """ 

292 if isinstance(chunks, Sized): 

293 n_chunks = len(chunks) 

294 else: 

295 n_chunks = None 

296 if total is None: 

297 total = False 

298 if skip_scalar and n_chunks == 1: 

299 yield from chunks 

300 return 

301 use_n_chunks = False 

302 if total is True or total is None: 

303 total = 0 

304 for c in chunks: 

305 if total is True or isinstance(c, Sized): 

306 total += len(c) # type: ignore 

307 else: 

308 use_n_chunks = True 

309 total = None 

310 break 

311 if total is False: 

312 total = None 

313 use_n_chunks = True 

314 if use_n_chunks: 

315 with self.bar(desc=desc, total=n_chunks) as bar: # type: ignore 

316 for chunk in chunks: 

317 yield chunk 

318 bar.update(1) 

319 else: 

320 with self.bar(desc=desc, total=total) as bar: # type: ignore 

321 for chunk in chunks: 

322 yield chunk 

323 bar.update(len(chunk)) # type: ignore 

324 

325 def iter_item_chunks( 

326 self, 

327 items: Iterable[tuple[_K, _V]], 

328 desc: str | None = None, 

329 total: int | None = None, 

330 skip_scalar: bool = True, 

331 ) -> Generator[tuple[_K, _V], None, None]: 

332 """Wrap iteration over chunks of items in a progress bar. 

333 

334 Parameters 

335 ---------- 

336 items : `~collections.abc.Iterable` 

337 An iterable whose elements are (key, value) tuples, where the 

338 values are themselves iterable. 

339 desc: `str`, optional 

340 A user-friendly description for this progress bar; usually appears 

341 next to it. If not provided, ``self.name`` is used (which is not 

342 usually a user-friendly string, but may be appropriate for 

343 debug-level progress). 

344 total : `int`, optional 

345 The total number of steps in this progress bar; defaults to the sum 

346 of the lengths of the chunks if this can be computed. If this is 

347 provided or `True`, each element in ``chunks`` must be sized but 

348 ``chunks`` itself need not be (and may be a single-pass iterable). 

349 skip_scalar: `bool`, optional 

350 If `True` and there are zero or one items, do not report progress. 

351 

352 Yields 

353 ------ 

354 chunk 

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

356 

357 Notes 

358 ----- 

359 This attempts to display as much progress as possible given the 

360 limitations of the iterables, assuming that sized iterables are also 

361 multi-pass (as is true of all built-in collections and lazy iterators). 

362 In detail, if ``total`` is `None`: 

363 

364 - if ``chunks`` and its values elements are both sized, ``total`` is 

365 computed from them and full progress is reported, and ``chunks`` must 

366 be a multi-pass iterable. 

367 - if ``chunks`` is sized but its value elements are not, a progress bar 

368 over the number of chunks is shown, and ``chunks`` must be a 

369 multi-pass iterable. 

370 - if ``chunks`` is not sized, the progress bar just shows when updates 

371 occur. 

372 

373 If ``total`` is `True` or an integer, ``chunks`` need not be sized, but 

374 its value elements must be, ``chunks`` must be a multi-pass iterable, 

375 and full progress is shown. 

376 

377 If ``total`` is `False`, ``chunks`` and its values elements need not be 

378 sized, and the progress bar just shows when updates occur. 

379 """ 

380 if isinstance(items, Sized): 

381 n_items = len(items) 

382 if skip_scalar and n_items == 1: 

383 yield from items 

384 return 

385 else: 

386 n_items = None 

387 if total is None: 

388 total = False 

389 use_n_items = False 

390 if total is True or total is None: 

391 total = 0 

392 for _, v in items: 

393 if total is True or isinstance(v, Sized): 

394 total += len(v) # type: ignore 

395 else: 

396 use_n_items = True 

397 total = None 

398 break 

399 if total is False: 

400 total = None 

401 use_n_items = True 

402 if use_n_items: 

403 with self.bar(desc=desc, total=n_items) as bar: # type: ignore 

404 for chunk in items: 

405 yield chunk 

406 bar.update(1) 

407 else: 

408 with self.bar(desc=desc, total=total) as bar: # type: ignore 

409 for k, v in items: 

410 yield k, v 

411 bar.update(len(v)) # type: ignore 

412 

413 

414class ProgressHandler(ABC): 

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

416 

417 @abstractmethod 

418 def get_progress_bar( 

419 self, iterable: Iterable[_T] | None, desc: str, total: int | None, level: int 

420 ) -> ContextManager[ProgressBar[_T]]: 

421 """Create a new progress bar. 

422 

423 Parameters 

424 ---------- 

425 iterable : `~collections.abc.Iterable` or `None` 

426 An arbitrary Python iterable that will be iterated over when the 

427 returned `ProgressBar` is. If `None`, whether the progress bar is 

428 iterable is handler-defined, but it may be updated manually. 

429 desc: `str` 

430 A user-friendly description for this progress bar; usually appears 

431 next to it 

432 total : `int` or `None` 

433 The total number of steps in this progress bar. If `None``, 

434 ``len(iterable)`` should be used. If that does not work, whether 

435 the progress bar works at all is handler-defined. 

436 level : `int` 

437 A `logging` level value (defaults to `logging.INFO`) associated 

438 with the process reporting progress. Handlers are not responsible 

439 for disabling progress reporting on levels, but may utilize level 

440 information to annotate them differently. 

441 """ 

442 raise NotImplementedError() 

443 

444 

445class _NullProgressBar(Iterable[_T]): 

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

447 through its iterable's elements. 

448 

449 Parameters 

450 ---------- 

451 iterable : `~collections.abc.Iterable` or `None` 

452 An arbitrary Python iterable that will be iterated over when ``self`` 

453 is. 

454 """ 

455 

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

457 self._iterable = iterable 

458 

459 @classmethod 

460 @contextmanager 

461 def context(cls, iterable: Iterable[_T] | None) -> Generator[_NullProgressBar[_T], None, None]: 

462 """Return a trivial context manager that wraps an instance of this 

463 class. 

464 

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

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

467 

468 Parameters 

469 ---------- 

470 iterable : `~collections.abc.Iterable` or `None` 

471 An arbitrary Python iterable that will be iterated over when the 

472 returned object is. 

473 """ 

474 yield cls(iterable) 

475 

476 def __iter__(self) -> Iterator[_T]: 

477 assert self._iterable is not None, "Cannot iterate over progress bar initialized without iterable." 

478 return iter(self._iterable) 

479 

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

481 pass