Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/black/handle_ipynb_magics.py: 27%

1"""Functions to process IPython magics with.""" 

2 

3import ast 

4import collections 

5import dataclasses 

6import secrets 

7import sys 

8from functools import lru_cache 

9from typing import Dict, List, Optional, Tuple 

10 

11if sys.version_info >= (3, 10): 

12 from typing import TypeGuard 

13else: 

14 from typing_extensions import TypeGuard 

15 

16from black.output import out 

17from black.report import NothingChanged 

18 

19TRANSFORMED_MAGICS = frozenset( 

20 ( 

21 "get_ipython().run_cell_magic", 

22 "get_ipython().system", 

23 "get_ipython().getoutput", 

24 "get_ipython().run_line_magic", 

25 ) 

26) 

27TOKENS_TO_IGNORE = frozenset( 

28 ( 

29 "ENDMARKER", 

30 "NL", 

31 "NEWLINE", 

32 "COMMENT", 

33 "DEDENT", 

34 "UNIMPORTANT_WS", 

35 "ESCAPED_NL", 

36 ) 

37) 

38PYTHON_CELL_MAGICS = frozenset( 

39 ( 

40 "capture", 

41 "prun", 

42 "pypy", 

43 "python", 

44 "python3", 

45 "time", 

46 "timeit", 

47 ) 

48) 

49TOKEN_HEX = secrets.token_hex 

50 

51 

52@dataclasses.dataclass(frozen=True) 

53class Replacement: 

54 mask: str 

55 src: str 

56 

57 

58@lru_cache() 

59def jupyter_dependencies_are_installed(*, verbose: bool, quiet: bool) -> bool: 

60 try: 

61 import IPython # noqa:F401 

62 import tokenize_rt # noqa:F401 

63 except ModuleNotFoundError: 

64 if verbose or not quiet: 

65 msg = ( 

66 "Skipping .ipynb files as Jupyter dependencies are not installed.\n" 

67 'You can fix this by running ``pip install "black[jupyter]"``' 

68 ) 

69 out(msg) 

70 return False 

71 else: 

72 return True 

73 

74 

75def remove_trailing_semicolon(src: str) -> Tuple[str, bool]: 

76 """Remove trailing semicolon from Jupyter notebook cell. 

77 

78 For example, 

79 

80 fig, ax = plt.subplots() 

81 ax.plot(x_data, y_data); # plot data 

82 

83 would become 

84 

85 fig, ax = plt.subplots() 

86 ax.plot(x_data, y_data) # plot data 

87 

88 Mirrors the logic in `quiet` from `IPython.core.displayhook`, but uses 

89 ``tokenize_rt`` so that round-tripping works fine. 

90 """ 

91 from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src 

92 

93 tokens = src_to_tokens(src) 

94 trailing_semicolon = False 

95 for idx, token in reversed_enumerate(tokens): 

96 if token.name in TOKENS_TO_IGNORE: 

97 continue 

98 if token.name == "OP" and token.src == ";": 

99 del tokens[idx] 

100 trailing_semicolon = True 

101 break 

102 if not trailing_semicolon: 

103 return src, False 

104 return tokens_to_src(tokens), True 

105 

106 

107def put_trailing_semicolon_back(src: str, has_trailing_semicolon: bool) -> str: 

108 """Put trailing semicolon back if cell originally had it. 

109 

110 Mirrors the logic in `quiet` from `IPython.core.displayhook`, but uses 

111 ``tokenize_rt`` so that round-tripping works fine. 

112 """ 

113 if not has_trailing_semicolon: 

114 return src 

115 from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src 

116 

117 tokens = src_to_tokens(src) 

118 for idx, token in reversed_enumerate(tokens): 

119 if token.name in TOKENS_TO_IGNORE: 

120 continue 

121 tokens[idx] = token._replace(src=token.src + ";") 

122 break 

123 else: # pragma: nocover 

124 raise AssertionError( 

125 "INTERNAL ERROR: Was not able to reinstate trailing semicolon. " 

126 "Please report a bug on https://github.com/psf/black/issues. " 

127 ) from None 

128 return str(tokens_to_src(tokens)) 

129 

130 

131def mask_cell(src: str) -> Tuple[str, List[Replacement]]: 

132 """Mask IPython magics so content becomes parseable Python code. 

133 

134 For example, 

135 

136 %matplotlib inline 

137 'foo' 

138 

139 becomes 

140 

141 "25716f358c32750e" 

142 'foo' 

143 

144 The replacements are returned, along with the transformed code. 

145 """ 

146 replacements: List[Replacement] = [] 

147 try: 

148 ast.parse(src) 

149 except SyntaxError: 

150 # Might have IPython magics, will process below. 

151 pass 

152 else: 

153 # Syntax is fine, nothing to mask, early return. 

154 return src, replacements 

155 

156 from IPython.core.inputtransformer2 import TransformerManager 

157 

158 transformer_manager = TransformerManager() 

159 transformed = transformer_manager.transform_cell(src) 

160 transformed, cell_magic_replacements = replace_cell_magics(transformed) 

161 replacements += cell_magic_replacements 

162 transformed = transformer_manager.transform_cell(transformed) 

163 transformed, magic_replacements = replace_magics(transformed) 

164 if len(transformed.splitlines()) != len(src.splitlines()): 

165 # Multi-line magic, not supported. 

166 raise NothingChanged 

167 replacements += magic_replacements 

168 return transformed, replacements 

169 

170 

171def get_token(src: str, magic: str) -> str: 

172 """Return randomly generated token to mask IPython magic with. 

173 

174 For example, if 'magic' was `%matplotlib inline`, then a possible 

175 token to mask it with would be `"43fdd17f7e5ddc83"`. The token 

176 will be the same length as the magic, and we make sure that it was 

177 not already present anywhere else in the cell. 

178 """ 

179 assert magic 

180 nbytes = max(len(magic) // 2 - 1, 1) 

181 token = TOKEN_HEX(nbytes) 

182 counter = 0 

183 while token in src: 

184 token = TOKEN_HEX(nbytes) 

185 counter += 1 

186 if counter > 100: 

187 raise AssertionError( 

188 "INTERNAL ERROR: Black was not able to replace IPython magic. " 

189 "Please report a bug on https://github.com/psf/black/issues. " 

190 f"The magic might be helpful: {magic}" 

191 ) from None 

192 if len(token) + 2 < len(magic): 

193 token = f"{token}." 

194 return f'"{token}"' 

195 

196 

197def replace_cell_magics(src: str) -> Tuple[str, List[Replacement]]: 

198 """Replace cell magic with token. 

199 

200 Note that 'src' will already have been processed by IPython's 

201 TransformerManager().transform_cell. 

202 

203 Example, 

204 

205 get_ipython().run_cell_magic('t', '-n1', 'ls =!ls\\n') 

206 

207 becomes 

208 

209 "a794." 

210 ls =!ls 

211 

212 The replacement, along with the transformed code, is returned. 

213 """ 

214 replacements: List[Replacement] = [] 

215 

216 tree = ast.parse(src) 

217 

218 cell_magic_finder = CellMagicFinder() 

219 cell_magic_finder.visit(tree) 

220 if cell_magic_finder.cell_magic is None: 

221 return src, replacements 

222 header = cell_magic_finder.cell_magic.header 

223 mask = get_token(src, header) 

224 replacements.append(Replacement(mask=mask, src=header)) 

225 return f"{mask}\n{cell_magic_finder.cell_magic.body}", replacements 

226 

227 

228def replace_magics(src: str) -> Tuple[str, List[Replacement]]: 

229 """Replace magics within body of cell. 

230 

231 Note that 'src' will already have been processed by IPython's 

232 TransformerManager().transform_cell. 

233 

234 Example, this 

235 

236 get_ipython().run_line_magic('matplotlib', 'inline') 

237 'foo' 

238 

239 becomes 

240 

241 "5e67db56d490fd39" 

242 'foo' 

243 

244 The replacement, along with the transformed code, are returned. 

245 """ 

246 replacements = [] 

247 magic_finder = MagicFinder() 

248 magic_finder.visit(ast.parse(src)) 

249 new_srcs = [] 

250 for i, line in enumerate(src.splitlines(), start=1): 

251 if i in magic_finder.magics: 

252 offsets_and_magics = magic_finder.magics[i] 

253 if len(offsets_and_magics) != 1: # pragma: nocover 

254 raise AssertionError( 

255 f"Expecting one magic per line, got: {offsets_and_magics}\n" 

256 "Please report a bug on https://github.com/psf/black/issues." 

257 ) 

258 col_offset, magic = ( 

259 offsets_and_magics[0].col_offset, 

260 offsets_and_magics[0].magic, 

261 ) 

262 mask = get_token(src, magic) 

263 replacements.append(Replacement(mask=mask, src=magic)) 

264 line = line[:col_offset] + mask 

265 new_srcs.append(line) 

266 return "\n".join(new_srcs), replacements 

267 

268 

269def unmask_cell(src: str, replacements: List[Replacement]) -> str: 

270 """Remove replacements from cell. 

271 

272 For example 

273 

274 "9b20" 

275 foo = bar 

276 

277 becomes 

278 

279 %%time 

280 foo = bar 

281 """ 

282 for replacement in replacements: 

283 src = src.replace(replacement.mask, replacement.src) 

284 return src 

285 

286 

287def _is_ipython_magic(node: ast.expr) -> TypeGuard[ast.Attribute]: 

288 """Check if attribute is IPython magic. 

289 

290 Note that the source of the abstract syntax tree 

291 will already have been processed by IPython's 

292 TransformerManager().transform_cell. 

293 """ 

294 return ( 

295 isinstance(node, ast.Attribute) 

296 and isinstance(node.value, ast.Call) 

297 and isinstance(node.value.func, ast.Name) 

298 and node.value.func.id == "get_ipython" 

299 ) 

300 

301 

302def _get_str_args(args: List[ast.expr]) -> List[str]: 

303 str_args = [] 

304 for arg in args: 

305 assert isinstance(arg, ast.Str) 

306 str_args.append(arg.s) 

307 return str_args 

308 

309 

310@dataclasses.dataclass(frozen=True) 

311class CellMagic: 

312 name: str 

313 params: Optional[str] 

314 body: str 

315 

316 @property 

317 def header(self) -> str: 

318 if self.params: 

319 return f"%%{self.name} {self.params}" 

320 return f"%%{self.name}" 

321 

322 

323# ast.NodeVisitor + dataclass = breakage under mypyc. 

324class CellMagicFinder(ast.NodeVisitor): 

325 """Find cell magics. 

326 

327 Note that the source of the abstract syntax tree 

328 will already have been processed by IPython's 

329 TransformerManager().transform_cell. 

330 

331 For example, 

332 

333 %%time\nfoo() 

334 

335 would have been transformed to 

336 

337 get_ipython().run_cell_magic('time', '', 'foo()\\n') 

338 

339 and we look for instances of the latter. 

340 """ 

341 

342 def __init__(self, cell_magic: Optional[CellMagic] = None) -> None: 

343 self.cell_magic = cell_magic 

344 

345 def visit_Expr(self, node: ast.Expr) -> None: 

346 """Find cell magic, extract header and body.""" 

347 if ( 

348 isinstance(node.value, ast.Call) 

349 and _is_ipython_magic(node.value.func) 

350 and node.value.func.attr == "run_cell_magic" 

351 ): 

352 args = _get_str_args(node.value.args) 

353 self.cell_magic = CellMagic(name=args[0], params=args[1], body=args[2]) 

354 self.generic_visit(node) 

355 

356 

357@dataclasses.dataclass(frozen=True) 

358class OffsetAndMagic: 

359 col_offset: int 

360 magic: str 

361 

362 

363# Unsurprisingly, subclassing ast.NodeVisitor means we can't use dataclasses here 

364# as mypyc will generate broken code. 

365class MagicFinder(ast.NodeVisitor): 

366 """Visit cell to look for get_ipython calls. 

367 

368 Note that the source of the abstract syntax tree 

369 will already have been processed by IPython's 

370 TransformerManager().transform_cell. 

371 

372 For example, 

373 

374 %matplotlib inline 

375 

376 would have been transformed to 

377 

378 get_ipython().run_line_magic('matplotlib', 'inline') 

379 

380 and we look for instances of the latter (and likewise for other 

381 types of magics). 

382 """ 

383 

384 def __init__(self) -> None: 

385 self.magics: Dict[int, List[OffsetAndMagic]] = collections.defaultdict(list) 

386 

387 def visit_Assign(self, node: ast.Assign) -> None: 

388 """Look for system assign magics. 

389 

390 For example, 

391 

392 black_version = !black --version 

393 env = %env var 

394 

395 would have been (respectively) transformed to 

396 

397 black_version = get_ipython().getoutput('black --version') 

398 env = get_ipython().run_line_magic('env', 'var') 

399 

400 and we look for instances of any of the latter. 

401 """ 

402 if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func): 

403 args = _get_str_args(node.value.args) 

404 if node.value.func.attr == "getoutput": 

405 src = f"!{args[0]}" 

406 elif node.value.func.attr == "run_line_magic": 

407 src = f"%{args[0]}" 

408 if args[1]: 

409 src += f" {args[1]}" 

410 else: 

411 raise AssertionError( 

412 f"Unexpected IPython magic {node.value.func.attr!r} found. " 

413 "Please report a bug on https://github.com/psf/black/issues." 

414 ) from None 

415 self.magics[node.value.lineno].append( 

416 OffsetAndMagic(node.value.col_offset, src) 

417 ) 

418 self.generic_visit(node) 

419 

420 def visit_Expr(self, node: ast.Expr) -> None: 

421 """Look for magics in body of cell. 

422 

423 For examples, 

424 

425 !ls 

426 !!ls 

427 ?ls 

428 ??ls 

429 

430 would (respectively) get transformed to 

431 

432 get_ipython().system('ls') 

433 get_ipython().getoutput('ls') 

434 get_ipython().run_line_magic('pinfo', 'ls') 

435 get_ipython().run_line_magic('pinfo2', 'ls') 

436 

437 and we look for instances of any of the latter. 

438 """ 

439 if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func): 

440 args = _get_str_args(node.value.args) 

441 if node.value.func.attr == "run_line_magic": 

442 if args[0] == "pinfo": 

443 src = f"?{args[1]}" 

444 elif args[0] == "pinfo2": 

445 src = f"??{args[1]}" 

446 else: 

447 src = f"%{args[0]}" 

448 if args[1]: 

449 src += f" {args[1]}" 

450 elif node.value.func.attr == "system": 

451 src = f"!{args[0]}" 

452 elif node.value.func.attr == "getoutput": 

453 src = f"!!{args[0]}" 

454 else: 

455 raise NothingChanged # unsupported magic. 

456 self.magics[node.value.lineno].append( 

457 OffsetAndMagic(node.value.col_offset, src) 

458 ) 

459 self.generic_visit(node)