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

81 statements  

« prev     ^ index     » next       coverage.py v6.4.2, created at 2022-07-16 02:09 -0700

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

41 

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

43_K = TypeVar("_K") 

44_V = TypeVar("_V", bound=Collection) 

45 

46 

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

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

49 

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

51 `Progress.bar` method. 

52 

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

59 

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

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

62 

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 

71 

72 

73class Progress: 

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

75 related tools. 

76 

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. 

82 

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. 

93 

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

102 

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

104 self._name = name 

105 self._level = level 

106 

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 

112 

113 @classmethod 

114 def set_handler(cls, handler: Optional[ProgressHandler]) -> None: 

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

116 

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. 

120 

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 

128 

129 def is_enabled(self) -> bool: 

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

131 

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 

143 

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

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

146 

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. 

153 

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) 

161 

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. 

170 

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. 

189 

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) 

213 

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. 

222 

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. 

239 

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 

247 

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. 

256 

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. 

274 

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

289 

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. 

298 

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. 

317 

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

332 

333 

334class ProgressHandler(ABC): 

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

336 

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. 

342 

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

363 

364 

365class _NullProgressBar(Iterable[_T]): 

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

367 through its iterable's elements. 

368 

369 Parameters 

370 ---------- 

371 iterable : `Iterable` or `None` 

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

373 is. 

374 """ 

375 

376 def __init__(self, iterable: Optional[Iterable[_T]]): 

377 self._iterable = iterable 

378 

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. 

384 

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

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

387 

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) 

395 

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) 

399 

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

401 pass