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

81 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-15 09:13 +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 Collection, Generator, Iterable, Iterator 

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

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: Collection[_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.Collection` 

250 A sized iterable whose elements are themselves both iterable and 

251 sized (i.e. ``len(item)`` works). If ``total`` is not provided, 

252 this may not be a single-pass iteration, because an initial pass to 

253 estimate the total number of elements is required. 

254 desc: `str`, optional 

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

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

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

258 debug-level progress). 

259 total : `int`, optional 

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

261 sum of the lengths of the chunks. 

262 skip_scalar: `bool`, optional 

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

264 

265 Yields 

266 ------ 

267 chunk 

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

269 """ 

270 if skip_scalar and len(chunks) <= 1: 

271 yield from chunks 

272 else: 

273 if total is None: 

274 total = sum(len(c) for c in chunks) 

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

276 for chunk in chunks: 

277 yield chunk 

278 bar.update(len(chunk)) 

279 

280 def iter_item_chunks( 

281 self, 

282 items: Collection[tuple[_K, _V]], 

283 desc: str | None = None, 

284 total: int | None = None, 

285 skip_scalar: bool = True, 

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

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

288 

289 Parameters 

290 ---------- 

291 items : `~collections.abc.Iterable` 

292 A sized iterable whose elements are (key, value) tuples, where the 

293 values are themselves both iterable and sized (i.e. ``len(item)`` 

294 works). If ``total`` is not provided, this may not be a 

295 single-pass iteration, because an initial pass to estimate the 

296 total number of elements is required. 

297 desc: `str`, optional 

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

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

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

301 debug-level progress). 

302 total : `int`, optional 

303 The total number of values in this progress bar; defaults to the 

304 sum of the lengths of the chunks. 

305 skip_scalar: `bool`, optional 

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

307 

308 Yields 

309 ------ 

310 chunk 

311 The same items that iteration over ``chunks`` would yield. 

312 """ 

313 if skip_scalar and len(items) <= 1: 

314 yield from items 

315 else: 

316 if total is None: 

317 total = sum(len(v) for _, v in items) 

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

319 for key, chunk in items: 

320 yield key, chunk 

321 bar.update(len(chunk)) 

322 

323 

324class ProgressHandler(ABC): 

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

326 

327 @abstractmethod 

328 def get_progress_bar( 

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

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

331 """Create a new progress bar. 

332 

333 Parameters 

334 ---------- 

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

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

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

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

339 desc: `str` 

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

341 next to it 

342 total : `int` or `None` 

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

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

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

346 level : `int` 

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

348 with the process reporting progress. Handlers are not responsible 

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

350 information to annotate them differently. 

351 """ 

352 raise NotImplementedError() 

353 

354 

355class _NullProgressBar(Iterable[_T]): 

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

357 through its iterable's elements. 

358 

359 Parameters 

360 ---------- 

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

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

363 is. 

364 """ 

365 

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

367 self._iterable = iterable 

368 

369 @classmethod 

370 @contextmanager 

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

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

373 class. 

374 

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

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

377 

378 Parameters 

379 ---------- 

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

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

382 returned object is. 

383 """ 

384 yield cls(iterable) 

385 

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

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

388 return iter(self._iterable) 

389 

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

391 pass