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

122 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-25 15:14 +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 contextlib 

27import logging 

28from abc import ABC, abstractmethod 

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

30from contextlib import AbstractContextManager, contextmanager 

31from typing import ClassVar, Protocol, TypeVar 

32 

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

34_K = TypeVar("_K") 

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

36 

37 

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

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

40 

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

42 `Progress.bar` method. 

43 

44 Notes 

45 ----- 

46 This interface is intentionally defined as the intersection of the progress 

47 bar objects returned by the ``click`` and ``tqdm`` packages, allowing those 

48 to directly satisfy code that uses this interface. 

49 """ 

50 

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

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

53 

54 Parameters 

55 ---------- 

56 n : `int`, optional 

57 Increment the progress bar by this many steps (defaults to ``1``). 

58 Note that this is a relative increment, not an absolute progress 

59 value. 

60 """ 

61 pass 

62 

63 

64class Progress: 

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

66 related tools. 

67 

68 This class automatically creates progress bars (or not) depending on 

69 whether a handle (see `ProgressHandler`) has been installed and the given 

70 name and level are enabled. When progress reporting is not enabled, it 

71 returns dummy objects that can be used just like progress bars by calling 

72 code. 

73 

74 Parameters 

75 ---------- 

76 name : `str` 

77 Name of the process whose progress is being reported. This is in 

78 general the name of a group of progress bars, not necessarily a single 

79 one, and it should have the same form as a logger name. 

80 level : `int`, optional 

81 A `logging` level value (defaults to `logging.INFO`). Progress 

82 reporting is enabled if a logger with ``name`` is enabled for this 

83 level, and a `ProgressHandler` has been installed. 

84 

85 Notes 

86 ----- 

87 The progress system inspects the level for a name using the Python built-in 

88 `logging` module, and may not respect level-setting done via the 

89 ``lsst.log`` interface. But while `logging` may be necessary to control 

90 progress bar visibility, the progress system can still be used together 

91 with either system for actual logging. 

92 """ 

93 

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

95 self._name = name 

96 self._level = level 

97 

98 # The active handler is held in a ContextVar to isolate unit tests run 

99 # by pytest-xdist. If butler codes is ever used in a real multithreaded 

100 # or asyncio application _and_ we want progress bars, we'll have to set 

101 # up per-thread handlers or similar. 

102 _active_handler: ClassVar[ProgressHandler | None] = None 

103 

104 @classmethod 

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

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

107 

108 This should only be called in very high-level code that can be 

109 reasonably confident that it will dominate its current process, e.g. 

110 at the initialization of a command-line script or Jupyter notebook. 

111 

112 Parameters 

113 ---------- 

114 handler : `ProgressHandler` or `None` 

115 Object that will handle all progress reporting. May be set to 

116 `None` to disable progress reporting. 

117 """ 

118 cls._active_handler = handler 

119 

120 def is_enabled(self) -> bool: 

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

122 

123 Returns 

124 ------- 

125 enabled : `bool` 

126 `True` if there is a `ProgressHandler` set and a logger with the 

127 same name and level as ``self`` is enabled. 

128 """ 

129 if self._active_handler is not None: 

130 logger = logging.getLogger(self._name) 

131 if logger.isEnabledFor(self._level): 

132 return True 

133 return False 

134 

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

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

137 

138 Parameters 

139 ---------- 

140 level : `int` 

141 A `logging` level value. Progress reporting is enabled if a logger 

142 with ``name`` is enabled for this level, and a `ProgressHandler` 

143 has been installed. 

144 

145 Returns 

146 ------- 

147 progress : `Progress` 

148 A new `Progress` object with the same name as ``self`` and the 

149 given ``level``. 

150 """ 

151 return Progress(self._name, level) 

152 

153 def bar( 

154 self, 

155 iterable: Iterable[_T] | None = None, 

156 desc: str | None = None, 

157 total: int | None = None, 

158 skip_scalar: bool = True, 

159 ) -> AbstractContextManager[ProgressBar[_T]]: 

160 """Return a new progress bar context manager. 

161 

162 Parameters 

163 ---------- 

164 iterable : `~collections.abc.Iterable`, optional 

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

166 returned `ProgressBar` is. If not provided, whether the progress 

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

168 desc: `str`, optional 

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

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

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

172 debug-level progress). 

173 total : `int`, optional 

174 The total number of steps in this progress bar. If not provided, 

175 ``len(iterable)`` is used. If that does not work, whether the 

176 progress bar works at all is handler-defined, and hence this mode 

177 should not be relied upon. 

178 skip_scalar: `bool`, optional 

179 If `True` and ``total`` is zero or one, do not report progress. 

180 

181 Returns 

182 ------- 

183 bar : `contextlib.AbstractContextManager` [ `ProgressBar` ] 

184 A context manager that returns an object satisfying the 

185 `ProgressBar` interface when it is entered. 

186 """ 

187 if self.is_enabled(): 

188 if desc is None: 

189 desc = self._name 

190 handler = self._active_handler 

191 assert handler, "Guaranteed by `is_enabled` check above." 

192 if skip_scalar: 

193 if total is None: 

194 with contextlib.suppress(TypeError): 

195 # static typing says len() won't but that's why 

196 # we're doing it inside a try block. 

197 total = len(iterable) # type: ignore 

198 if total is not None and total <= 1: 

199 return _NullProgressBar.context(iterable) 

200 return handler.get_progress_bar(iterable, desc=desc, total=total, level=self._level) 

201 return _NullProgressBar.context(iterable) 

202 

203 def wrap( 

204 self, 

205 iterable: Iterable[_T], 

206 desc: str | None = None, 

207 total: int | None = None, 

208 skip_scalar: bool = True, 

209 ) -> Generator[_T, None, None]: 

210 """Iterate over an object while reporting progress. 

211 

212 Parameters 

213 ---------- 

214 iterable : `~collections.abc.Iterable` 

215 An arbitrary Python iterable to iterate over. 

216 desc: `str`, optional 

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

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

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

220 debug-level progress). 

221 total : `int`, optional 

222 The total number of steps in this progress bar. If not provided, 

223 ``len(iterable)`` is used. If that does not work, whether the 

224 progress bar works at all is handler-defined, and hence this mode 

225 should not be relied upon. 

226 skip_scalar: `bool`, optional 

227 If `True` and ``total`` is zero or one, do not report progress. 

228 

229 Yields 

230 ------ 

231 element 

232 The same objects that iteration over ``iterable`` would yield. 

233 """ 

234 with self.bar(iterable, desc=desc, total=total, skip_scalar=skip_scalar) as bar: 

235 yield from bar 

236 

237 def iter_chunks( 

238 self, 

239 chunks: Iterable[_V], 

240 desc: str | None = None, 

241 total: int | None = None, 

242 skip_scalar: bool = True, 

243 ) -> Generator[_V, None, None]: 

244 """Wrap iteration over chunks of elements in a progress bar. 

245 

246 Parameters 

247 ---------- 

248 chunks : `~collections.abc.Iterable` 

249 An iterable whose elements are themselves iterable. 

250 desc: `str`, optional 

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

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

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

254 debug-level progress). 

255 total : `int`, optional 

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

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

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

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

260 skip_scalar: `bool`, optional 

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

262 

263 Yields 

264 ------ 

265 chunk 

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

267 

268 Notes 

269 ----- 

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

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

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

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

274 

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

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

277 multi-pass iterable. 

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

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

280 iterable. 

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

282 occur. 

283 

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

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

286 full progress is shown. 

287 

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

289 and the progress bar just shows when updates occur. 

290 """ 

291 if isinstance(chunks, Sized): 

292 n_chunks = len(chunks) 

293 else: 

294 n_chunks = None 

295 if total is None: 

296 total = False 

297 if skip_scalar and n_chunks == 1: 

298 yield from chunks 

299 return 

300 use_n_chunks = False 

301 if total is True or total is None: 

302 total = 0 

303 for c in chunks: 

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

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

306 else: 

307 use_n_chunks = True 

308 total = None 

309 break 

310 if total is False: 

311 total = None 

312 use_n_chunks = True 

313 if use_n_chunks: 

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

315 for chunk in chunks: 

316 yield chunk 

317 bar.update(1) 

318 else: 

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

320 for chunk in chunks: 

321 yield chunk 

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

323 

324 def iter_item_chunks( 

325 self, 

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

327 desc: str | None = None, 

328 total: int | None = None, 

329 skip_scalar: bool = True, 

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

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

332 

333 Parameters 

334 ---------- 

335 items : `~collections.abc.Iterable` 

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

337 values are themselves iterable. 

338 desc: `str`, optional 

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

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

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

342 debug-level progress). 

343 total : `int`, optional 

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

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

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

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

348 skip_scalar: `bool`, optional 

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

350 

351 Yields 

352 ------ 

353 chunk 

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

355 

356 Notes 

357 ----- 

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

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

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

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

362 

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

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

365 be a multi-pass iterable. 

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

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

368 multi-pass iterable. 

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

370 occur. 

371 

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

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

374 and full progress is shown. 

375 

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

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

378 """ 

379 if isinstance(items, Sized): 

380 n_items = len(items) 

381 if skip_scalar and n_items == 1: 

382 yield from items 

383 return 

384 else: 

385 n_items = None 

386 if total is None: 

387 total = False 

388 use_n_items = False 

389 if total is True or total is None: 

390 total = 0 

391 for _, v in items: 

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

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

394 else: 

395 use_n_items = True 

396 total = None 

397 break 

398 if total is False: 

399 total = None 

400 use_n_items = True 

401 if use_n_items: 

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

403 for chunk in items: 

404 yield chunk 

405 bar.update(1) 

406 else: 

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

408 for k, v in items: 

409 yield k, v 

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

411 

412 

413class ProgressHandler(ABC): 

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

415 

416 @abstractmethod 

417 def get_progress_bar( 

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

419 ) -> AbstractContextManager[ProgressBar[_T]]: 

420 """Create a new progress bar. 

421 

422 Parameters 

423 ---------- 

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

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

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

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

428 desc: `str` 

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

430 next to it 

431 total : `int` or `None` 

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

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

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

435 level : `int` 

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

437 with the process reporting progress. Handlers are not responsible 

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

439 information to annotate them differently. 

440 """ 

441 raise NotImplementedError() 

442 

443 

444class _NullProgressBar(Iterable[_T]): 

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

446 through its iterable's elements. 

447 

448 Parameters 

449 ---------- 

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

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

452 is. 

453 """ 

454 

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

456 self._iterable = iterable 

457 

458 @classmethod 

459 @contextmanager 

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

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

462 class. 

463 

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

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

466 

467 Parameters 

468 ---------- 

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

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

471 returned object is. 

472 """ 

473 yield cls(iterable) 

474 

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

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

477 return iter(self._iterable) 

478 

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

480 pass