Coverage for /home/crpier/Projects/snektest/snektest/output.py: 64%

1from __future__ import annotations 

2 

3import inspect 

4import pdb # noqa: T100 

5import sys 

6import warnings 

7from collections.abc import Callable, Generator 

8from contextlib import contextmanager 

9from dataclasses import dataclass 

10from io import StringIO 

11from typing import Any, TextIO 

12 

13 

14class StdinProxy: 

15 """Proxy for sys.stdin that disables output capture when read from. 

16 

17 This allows breakpoint()/pdb to work correctly even when output is being 

18 captured. When stdin is read from (e.g., pdb waiting for user input), 

19 the proxy calls a callback to restore stdout/stderr so the debugger 

20 can display its prompt and receive user input. 

21 """ 

22 

23 def __init__( 

24 self, 

25 original_stdin: TextIO, 

26 disable_capture_callback: Callable[[], None], 

27 ) -> None: 

28 self._original_stdin = original_stdin 

29 self._disable_capture = disable_capture_callback 

30 

31 def read(self, size: int = -1) -> str: 

32 self._disable_capture() 

33 return self._original_stdin.read(size) 

34 

35 def readline(self, size: int = -1) -> str: 

36 self._disable_capture() 

37 return self._original_stdin.readline(size) 

38 

39 def readlines(self, hint: int = -1) -> list[str]: 

40 self._disable_capture() 

41 return self._original_stdin.readlines(hint) 

42 

43 def __iter__(self) -> StdinProxy: 

44 self._disable_capture() 

45 return self 

46 

47 def __next__(self) -> str: 

48 self._disable_capture() 

49 return next(self._original_stdin) 

50 

51 def fileno(self) -> int: 

52 return self._original_stdin.fileno() 

53 

54 def isatty(self) -> bool: 

55 return self._original_stdin.isatty() 

56 

57 def close(self) -> None: 

58 pass # Don't close the original stdin 

59 

60 @property 

61 def closed(self) -> bool: 

62 return self._original_stdin.closed 

63 

64 @property 

65 def encoding(self) -> str: 

66 return self._original_stdin.encoding 

67 

68 @property 

69 def mode(self) -> str: 

70 return self._original_stdin.mode 

71 

72 @property 

73 def name(self) -> str: 

74 return self._original_stdin.name 

75 

76 def __getattr__(self, name: str) -> object: 

77 return getattr(self._original_stdin, name) 

78 

79 

80@dataclass(frozen=True) 

81class _OriginalSysState: 

82 stdout: Any 

83 stderr: Any 

84 stdin: Any 

85 settrace: Any 

86 breakpointhook: Any 

87 

88 

89def _make_disable_capture( 

90 *, 

91 system_stdout: Any, 

92 system_stderr: Any, 

93) -> Callable[[], None]: 

94 capture_disabled = False 

95 

96 def disable_capture() -> None: 

97 nonlocal capture_disabled 

98 if capture_disabled: 

99 return 

100 sys.stdout = system_stdout 

101 sys.stderr = system_stderr 

102 capture_disabled = True 

103 

104 return disable_capture 

105 

106 

107def _make_settrace_wrapper( 

108 *, 

109 system_settrace: Any, 

110 disable_capture: Callable[[], None], 

111) -> Callable[[Any], None]: 

112 def settrace_wrapper(func: Any) -> None: 

113 disable_capture() 

114 system_settrace(func) 

115 

116 return settrace_wrapper 

117 

118 

119def _maybe_run_inline_pdb_breakpoint( 

120 *, 

121 system_breakpointhook: Any, 

122 caller_frame: Any, 

123 args: tuple[Any, ...], 

124 kwargs: dict[str, Any], 

125 pdb_factory: Callable[..., Any] = pdb.Pdb, 

126) -> Any: 

127 if system_breakpointhook is not sys.__breakpointhook__: 

128 return system_breakpointhook(*args, **kwargs) 

129 

130 if caller_frame is None: 

131 return system_breakpointhook(*args, **kwargs) 

132 if args: 

133 return system_breakpointhook(*args, **kwargs) 

134 

135 header = kwargs.get("header") 

136 commands = kwargs.get("commands") 

137 if set(kwargs) - {"header", "commands"}: 

138 return system_breakpointhook(*args, **kwargs) 

139 

140 debugger: pdb.Pdb = pdb_factory( 

141 mode="inline", 

142 backend="monitoring", 

143 colorize=True, 

144 ) 

145 if header is not None: 

146 _ = debugger.message(header) 

147 return debugger.set_trace(caller_frame, commands=commands) 

148 

149 

150def _make_breakpointhook_wrapper( 

151 *, 

152 system_breakpointhook: Any, 

153 disable_capture: Callable[[], None], 

154 frame_provider: Callable[[], Any] = inspect.currentframe, 

155 pdb_factory: Callable[..., Any] = pdb.Pdb, 

156) -> Callable[..., Any]: 

157 def breakpointhook_wrapper(*args: Any, **kwargs: Any) -> Any: 

158 disable_capture() 

159 sys.breakpointhook = system_breakpointhook 

160 

161 frame = frame_provider() 

162 caller_frame = frame.f_back if frame else None 

163 

164 return _maybe_run_inline_pdb_breakpoint( 

165 system_breakpointhook=system_breakpointhook, 

166 caller_frame=caller_frame, 

167 args=args, 

168 kwargs=kwargs, 

169 pdb_factory=pdb_factory, 

170 ) 

171 

172 return breakpointhook_wrapper 

173 

174 

175def _format_warnings(warnings_list: list[warnings.WarningMessage]) -> list[str]: 

176 return [ 

177 f"{warning.filename}:{warning.lineno}: {warning.category.__name__}: {warning.message}" 

178 for warning in warnings_list 

179 ] 

180 

181 

182def _install_capture( 

183 *, 

184 output_buffer: StringIO, 

185 system_stdin: Any, 

186 disable_capture: Callable[[], None], 

187 settrace_wrapper: Any, 

188 breakpointhook_wrapper: Any, 

189) -> None: 

190 sys.stdout = output_buffer 

191 sys.stderr = output_buffer 

192 sys.stdin = StdinProxy(system_stdin, disable_capture) 

193 sys.settrace = settrace_wrapper 

194 sys.breakpointhook = breakpointhook_wrapper 

195 

196 

197def _restore_system_state(system: _OriginalSysState) -> None: 

198 sys.stdout = system.stdout 

199 sys.stderr = system.stderr 

200 sys.stdin = system.stdin 

201 sys.settrace = system.settrace 

202 sys.breakpointhook = system.breakpointhook 

203 

204 

205@contextmanager 

206def capture_output( 

207 *, 

208 frame_provider: Callable[[], Any] = inspect.currentframe, 

209 pdb_factory: Callable[..., Any] = pdb.Pdb, 

210) -> Generator[tuple[StringIO, list[str]]]: 

211 """Context manager to capture stdout, stderr, and warnings. 

212 

213 If stdin is read from (e.g., by pdb/breakpoint), output capture is 

214 automatically disabled so the debugger can function properly. 

215 """ 

216 output_buffer = StringIO() 

217 captured_warnings: list[str] = [] 

218 

219 original_sys = _OriginalSysState( 

220 stdout=sys.stdout, 

221 stderr=sys.stderr, 

222 stdin=sys.stdin, 

223 settrace=sys.settrace, 

224 breakpointhook=sys.breakpointhook, 

225 ) 

226 

227 disable_capture = _make_disable_capture( 

228 system_stdout=original_sys.stdout, 

229 system_stderr=original_sys.stderr, 

230 ) 

231 settrace_wrapper = _make_settrace_wrapper( 

232 system_settrace=original_sys.settrace, 

233 disable_capture=disable_capture, 

234 ) 

235 breakpointhook_wrapper = _make_breakpointhook_wrapper( 

236 system_breakpointhook=original_sys.breakpointhook, 

237 disable_capture=disable_capture, 

238 frame_provider=frame_provider, 

239 pdb_factory=pdb_factory, 

240 ) 

241 

242 _install_capture( 

243 output_buffer=output_buffer, 

244 system_stdin=original_sys.stdin, 

245 disable_capture=disable_capture, 

246 settrace_wrapper=settrace_wrapper, 

247 breakpointhook_wrapper=breakpointhook_wrapper, 

248 ) 

249 

250 with warnings.catch_warnings(record=True) as warning_list: 

251 warnings.simplefilter("always") 

252 try: 

253 yield output_buffer, captured_warnings 

254 finally: 

255 captured_warnings.extend(_format_warnings(warning_list)) 

256 _restore_system_state(original_sys) 

257 

258 

259@contextmanager 

260def maybe_capture_output( 

261 capture: bool, 

262) -> Generator[tuple[StringIO, list[str]]]: 

263 """Conditionally capture output based on a flag.""" 

264 if capture: 

265 with capture_output() as (buffer, warnings_list): 

266 yield buffer, warnings_list 

267 else: 

268 buffer = StringIO() 

269 warnings_list: list[str] = [] 

270 yield buffer, warnings_list