Coverage for src/lilbee/catalog/download_progress.py: 100%

1"""Download progress callback plumbing shared by all surfaces. 

2 

3The TUI runs under Textual which owns the terminal; tqdm output to 

4stderr/stdout corrupts the screen. This module provides a tqdm subclass 

5(``_CallbackProgressBar``) that suppresses terminal output and tracks 

6cumulative bytes; the ``_ProgressTracker`` wrapper produces a further subclass 

7that forwards progress to a plain ``Callable[[int, int], None]`` callback 

8(aggregating across split shards) and detects whether progress events actually 

9fired so the TUI can detect a cache-hit (no progress events) and render 

10``"already downloaded"`` instead of leaving the bar at 0%. 

11 

12``make_download_callback`` is the public entry point used by every 

13surface to convert raw bytes-progress into ``DownloadProgress`` events. 

14""" 

15 

16from __future__ import annotations 

17 

18import io 

19import threading 

20import time 

21from collections.abc import Callable 

22from typing import Any 

23 

24from tqdm.auto import tqdm as _base_tqdm 

25 

26from lilbee.catalog.models import DownloadProgress 

27 

28ProgressCallback = Callable[[int, int], None] 

29_BYTES_PER_MB = 1024 * 1024 

30 

31 

32def make_download_callback( 

33 on_update: Callable[[DownloadProgress], None], 

34 *, 

35 throttle_interval: float = 0.1, 

36) -> ProgressCallback: 

37 """Build a download progress callback that converts bytes to human-readable state. 

38 *on_update(progress: DownloadProgress)* is called at most once per 

39 ``throttle_interval`` seconds with a float percentage (0.0 to 100.0), a 

40 ``"<done>/<total> MB"`` detail string, and a cache-hit flag. Both the 

41 catalog and setup screens use this so byte-to-MB conversion and 

42 cache-hit detection aren't duplicated. 

43 """ 

44 last_update_time = 0.0 

45 seen_partial = False 

46 

47 def _on_progress(downloaded: int, total: int) -> None: 

48 nonlocal last_update_time, seen_partial 

49 

50 if total > 0 and downloaded >= total and not seen_partial: 

51 on_update( 

52 DownloadProgress(percent=100.0, detail="already downloaded", is_cache_hit=True) 

53 ) 

54 return 

55 seen_partial = True 

56 

57 now = time.monotonic() 

58 if now - last_update_time < throttle_interval: 

59 return 

60 last_update_time = now 

61 

62 mb_done = downloaded / _BYTES_PER_MB 

63 if total > 0: 

64 pct = min(downloaded * 100.0 / total, 100.0) 

65 mb_total = total / _BYTES_PER_MB 

66 on_update( 

67 DownloadProgress( 

68 percent=pct, 

69 detail=f"{mb_done:.0f}/{mb_total:.0f} MB", 

70 is_cache_hit=False, 

71 ) 

72 ) 

73 else: 

74 on_update(DownloadProgress(percent=0.0, detail=f"{mb_done:.0f} MB", is_cache_hit=False)) 

75 

76 return _on_progress 

77 

78 

79class _CallbackProgressBar(_base_tqdm): 

80 """tqdm subclass that suppresses terminal output and tracks cumulative bytes. 

81 

82 ``_ProgressTracker`` produces a further subclass of this that forwards the 

83 progress to a callback. 

84 Fully suppresses terminal output by disabling tqdm rendering and redirecting 

85 its file handle to a devnull sink: prevents ANSI escape sequences from leaking 

86 into Textual's managed terminal. 

87 

88 Overrides ``get_lock`` to return a threading lock instead of tqdm's default 

89 multiprocessing lock. Vanilla tqdm acquires ``self._lock`` even on the 

90 ``disable=True`` path (std.py:988), and the multiprocessing lock's lazy init 

91 raises ``ValueError`` when ``sys.stderr.fileno() == -1`` (Textual, Jupyter, 

92 pytest capture). A thread lock sidesteps that fd handling entirely. 

93 

94 ``update`` carries bytes written to disk, ``update_transfer`` bytes off the 

95 network. Both fire for the same bytes on HTTP, so the two are tracked apart 

96 and the larger reported rather than summed. 

97 """ 

98 

99 _lock = threading.RLock() 

100 

101 @classmethod 

102 def get_lock(cls) -> threading.RLock: 

103 return cls._lock 

104 

105 def __init__(self, *args: Any, **kwargs: Any): 

106 kwargs["disable"] = True 

107 kwargs["file"] = io.StringIO() # absorb any accidental tqdm output 

108 super().__init__(*args, **kwargs) 

109 # `initial` is the resume offset; seeding both keeps a resumed percentage absolute. 

110 self._written = int(self.n) 

111 self._transferred = int(self.n) 

112 

113 @property 

114 def _cumulative(self) -> int: 

115 return max(self._written, self._transferred) 

116 

117 def update(self, n: float = 1) -> bool | None: 

118 # The base only tracks cumulative bytes and suppresses output; forwarding 

119 # to the callback (with split-shard aggregation) lives in the 

120 # _ProgressTracker subclass below. 

121 self._written += int(n) 

122 return None 

123 

124 def update_transfer(self, n: float = 1) -> bool | None: 

125 """Absorb the network-bytes stream. 

126 

127 huggingface_hub routes xet transfer progress here only when this method 

128 exists; otherwise it opens a second tqdm on stderr. 

129 """ 

130 self._transferred += int(n) 

131 return None 

132 

133 def set_transfer_postfix_str(self, s: str = "", refresh: bool = True) -> None: 

134 """No-op: huggingface_hub sets a transfer rate on bars that take transfer.""" 

135 

136 

137class _ProgressTracker: 

138 """Wraps a tqdm_class to detect updates and aggregate across split shards. 

139 

140 For a multi-shard GGUF, each shard gets its own tqdm. Reporting each shard's 

141 own ``(done, total)`` would show N separate 0->100% cycles against the wrong 

142 total; instead the tracker carries ``grand_total`` (all shards) and a 

143 ``completed_base`` (bytes from finished shards), so the callback sees one 

144 monotonic 0->100% over the whole download. ``grand_total`` of 0 means 

145 single-file: fall back to the shard's own tqdm total (unchanged behavior). 

146 """ 

147 

148 def __init__(self, callback: Any, grand_total: int = 0) -> None: 

149 self.was_used = False 

150 self._callback = callback 

151 self.grand_total = grand_total 

152 self._completed_base = 0 

153 

154 def shard_done(self, shard_size: int) -> None: 

155 """Roll a finished shard's bytes into the base for the next shard.""" 

156 self._completed_base += shard_size 

157 

158 def make_tqdm_class(self) -> type[_base_tqdm]: 

159 tracker = self 

160 

161 class _Cls(_CallbackProgressBar): 

162 def _report(self) -> None: 

163 tracker.was_used = True 

164 done = tracker._completed_base + self._cumulative 

165 shard_total = self.total if self.total is not None else 0 

166 total = tracker.grand_total or (tracker._completed_base + shard_total) 

167 tracker._callback(int(done), int(total)) 

168 

169 def update(self, n: float = 1) -> bool | None: 

170 super().update(n) 

171 self._report() 

172 return None 

173 

174 def update_transfer(self, n: float = 1) -> bool | None: 

175 super().update_transfer(n) 

176 self._report() 

177 return None 

178 

179 return _Cls