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

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

2 

3import ast 

4import collections 

5import dataclasses 

6import re 

7import secrets 

8import string 

9from collections.abc import Collection 

10from functools import lru_cache 

11from importlib.util import find_spec 

12from typing import TypeGuard 

13 

14from black.mode import Mode 

15from black.output import out 

16from black.report import NothingChanged 

17 

18TRANSFORMED_MAGICS = frozenset(( 

19 "get_ipython().run_cell_magic", 

20 "get_ipython().system", 

21 "get_ipython().getoutput", 

22 "get_ipython().run_line_magic", 

23)) 

24TOKENS_TO_IGNORE = frozenset(( 

25 "ENDMARKER", 

26 "NL", 

27 "NEWLINE", 

28 "COMMENT", 

29 "DEDENT", 

30 "UNIMPORTANT_WS", 

31 "ESCAPED_NL", 

32)) 

33PYTHON_CELL_MAGICS = frozenset(( 

34 "capture", 

35 "prun", 

36 "pypy", 

37 "python", 

38 "python3", 

39 "time", 

40 "timeit", 

41)) 

42 

43 

44@dataclasses.dataclass(frozen=True) 

45class Replacement: 

46 mask: str 

47 src: str 

48 

49 

50@lru_cache 

51def jupyter_dependencies_are_installed(*, warn: bool) -> bool: 

52 installed = ( 

53 find_spec("tokenize_rt") is not None and find_spec("IPython") is not None 

54 ) 

55 if not installed and warn: 

56 msg = ( 

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

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

59 ) 

60 out(msg) 

61 return installed 

62 

63 

64def validate_cell(src: str, mode: Mode) -> None: 

65 r"""Check that cell does not already contain TransformerManager transformations, 

66 or non-Python cell magics, which might cause tokenizer_rt to break because of 

67 indentations. 

68 

69 If a cell contains ``!ls``, then it'll be transformed to 

70 ``get_ipython().system('ls')``. However, if the cell originally contained 

71 ``get_ipython().system('ls')``, then it would get transformed in the same way: 

72 

73 >>> TransformerManager().transform_cell("get_ipython().system('ls')") 

74 "get_ipython().system('ls')\n" 

75 >>> TransformerManager().transform_cell("!ls") 

76 "get_ipython().system('ls')\n" 

77 

78 Due to the impossibility of safely roundtripping in such situations, cells 

79 containing transformed magics will be ignored. 

80 """ 

81 if any(transformed_magic in src for transformed_magic in TRANSFORMED_MAGICS): 

82 raise NothingChanged 

83 

84 line = _get_code_start(src) 

85 if line.startswith("%%") and ( 

86 line.split(maxsplit=1)[0][2:] 

87 not in PYTHON_CELL_MAGICS | mode.python_cell_magics 

88 ): 

89 raise NothingChanged 

90 

91 

92def remove_trailing_semicolon(src: str) -> tuple[str, bool]: 

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

94 

95 For example, 

96 

97 fig, ax = plt.subplots() 

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

99 

100 would become 

101 

102 fig, ax = plt.subplots() 

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

104 

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

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

107 """ 

108 from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src 

109 

110 tokens = src_to_tokens(src) 

111 trailing_semicolon = False 

112 for idx, token in reversed_enumerate(tokens): 

113 if token.name in TOKENS_TO_IGNORE: 

114 continue 

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

116 del tokens[idx] 

117 trailing_semicolon = True 

118 break 

119 if not trailing_semicolon: 

120 return src, False 

121 return tokens_to_src(tokens), True 

122 

123 

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

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

126 

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

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

129 """ 

130 if not has_trailing_semicolon: 

131 return src 

132 from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src 

133 

134 tokens = src_to_tokens(src) 

135 for idx, token in reversed_enumerate(tokens): 

136 if token.name in TOKENS_TO_IGNORE: 

137 continue 

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

139 break 

140 else: # pragma: nocover 

141 raise AssertionError( 

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

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

144 ) from None 

145 return str(tokens_to_src(tokens)) 

146 

147 

148def mask_cell(src: str) -> tuple[str, list[Replacement]]: 

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

150 

151 For example, 

152 

153 %matplotlib inline 

154 'foo' 

155 

156 becomes 

157 

158 b"25716f358c32750" 

159 'foo' 

160 

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

162 """ 

163 replacements: list[Replacement] = [] 

164 try: 

165 ast.parse(src) 

166 except SyntaxError: 

167 # Might have IPython magics, will process below. 

168 pass 

169 else: 

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

171 return src, replacements 

172 

173 from IPython.core.inputtransformer2 import TransformerManager 

174 

175 transformer_manager = TransformerManager() 

176 # A side effect of the following transformation is that it also removes any 

177 # empty lines at the beginning of the cell. 

178 transformed = transformer_manager.transform_cell(src) 

179 transformed, cell_magic_replacements = replace_cell_magics(transformed) 

180 replacements += cell_magic_replacements 

181 transformed = transformer_manager.transform_cell(transformed) 

182 transformed, magic_replacements = replace_magics(transformed) 

183 if len(transformed.strip().splitlines()) != len(src.strip().splitlines()): 

184 # Multi-line magic, not supported. 

185 raise NothingChanged 

186 replacements += magic_replacements 

187 return transformed, replacements 

188 

189 

190def create_token(n_chars: int) -> str: 

191 """Create a randomly generated token that is n_chars characters long.""" 

192 assert n_chars > 0 

193 if n_chars == 1: 

194 return secrets.choice(string.ascii_letters) 

195 if n_chars < 4: 

196 return "_" + "".join( 

197 secrets.choice(string.ascii_letters + string.digits + "_") 

198 for _ in range(n_chars - 1) 

199 ) 

200 n_bytes = max(n_chars // 2 - 1, 1) 

201 token = secrets.token_hex(n_bytes) 

202 if len(token) + 3 > n_chars: 

203 token = token[:-1] 

204 # We use a bytestring so that the string does not get interpreted 

205 # as a docstring. 

206 return f'b"{token}"' 

207 

208 

209def get_token(src: str, magic: str, existing_tokens: Collection[str] = ()) -> str: 

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

211 

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

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

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

215 not already present anywhere else in the cell. 

216 """ 

217 assert magic 

218 n_chars = len(magic) 

219 token = create_token(n_chars) 

220 counter = 0 

221 while token in src or token in existing_tokens: 

222 token = create_token(n_chars) 

223 counter += 1 

224 if counter > 100: 

225 raise AssertionError( 

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

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

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

229 ) from None 

230 return token 

231 

232 

233def replace_cell_magics(src: str) -> tuple[str, list[Replacement]]: 

234 r"""Replace cell magic with token. 

235 

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

237 TransformerManager().transform_cell. 

238 

239 Example, 

240 

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

242 

243 becomes 

244 

245 "a794." 

246 ls =!ls 

247 

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

249 """ 

250 replacements: list[Replacement] = [] 

251 

252 tree = ast.parse(src) 

253 

254 cell_magic_finder = CellMagicFinder() 

255 cell_magic_finder.visit(tree) 

256 if cell_magic_finder.cell_magic is None: 

257 return src, replacements 

258 header = cell_magic_finder.cell_magic.header 

259 mask = get_token(src, header) 

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

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

262 

263 

264def replace_magics(src: str) -> tuple[str, list[Replacement]]: 

265 """Replace magics within body of cell. 

266 

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

268 TransformerManager().transform_cell. 

269 

270 Example, this 

271 

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

273 'foo' 

274 

275 becomes 

276 

277 "5e67db56d490fd39" 

278 'foo' 

279 

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

281 """ 

282 replacements = [] 

283 existing_tokens: set[str] = set() 

284 magic_finder = MagicFinder() 

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

286 new_srcs = [] 

287 for i, line in enumerate(src.split("\n"), start=1): 

288 if i in magic_finder.magics: 

289 offsets_and_magics = magic_finder.magics[i] 

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

291 raise AssertionError( 

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

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

294 ) 

295 col_offset, magic = ( 

296 offsets_and_magics[0].col_offset, 

297 offsets_and_magics[0].magic, 

298 ) 

299 mask = get_token(src, magic, existing_tokens) 

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

301 existing_tokens.add(mask) 

302 # AST column offsets are UTF-8 byte offsets, not character indices. 

303 prefix = line.encode("utf-8")[:col_offset].decode("utf-8") 

304 line = prefix + mask 

305 new_srcs.append(line) 

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

307 

308 

309def unmask_cell(src: str, replacements: list[Replacement]) -> str: 

310 """Remove replacements from cell. 

311 

312 For example 

313 

314 "9b20" 

315 foo = bar 

316 

317 becomes 

318 

319 %%time 

320 foo = bar 

321 """ 

322 for replacement in replacements: 

323 if src.count(replacement.mask) != 1: 

324 raise NothingChanged 

325 src = src.replace(replacement.mask, replacement.src, 1) 

326 if replacements: 

327 # Formatting can move a mask to where IPython no longer recognises the 

328 # magic, e.g. into the parentheses around a long right-hand side. 

329 try: 

330 _, unmasked_replacements = mask_cell(src) 

331 except SyntaxError: 

332 raise NothingChanged from None 

333 if [r.src for r in unmasked_replacements] != [r.src for r in replacements]: 

334 raise NothingChanged 

335 return src 

336 

337 

338def _get_code_start(src: str) -> str: 

339 """Provides the first line where the code starts. 

340 

341 Iterates over lines of code until it finds the first line that doesn't 

342 contain only empty spaces and comments. It removes any empty spaces at the 

343 start of the line and returns it. If such line doesn't exist, it returns an 

344 empty string. 

345 """ 

346 for match in re.finditer(".+", src): 

347 line = match.group(0).lstrip() 

348 if line and not line.startswith("#"): 

349 return line 

350 return "" 

351 

352 

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

354 """Check if attribute is IPython magic. 

355 

356 Note that the source of the abstract syntax tree 

357 will already have been processed by IPython's 

358 TransformerManager().transform_cell. 

359 """ 

360 return ( 

361 isinstance(node, ast.Attribute) 

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

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

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

365 ) 

366 

367 

368def _get_str_args(args: list[ast.expr]) -> list[str]: 

369 str_args = [] 

370 for arg in args: 

371 assert isinstance(arg, ast.Constant) and isinstance(arg.value, str) 

372 str_args.append(arg.value) 

373 return str_args 

374 

375 

376@dataclasses.dataclass(frozen=True) 

377class CellMagic: 

378 name: str 

379 params: str | None 

380 body: str 

381 

382 @property 

383 def header(self) -> str: 

384 if self.params: 

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

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

387 

388 

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

390class CellMagicFinder(ast.NodeVisitor): 

391 r"""Find cell magics. 

392 

393 Note that the source of the abstract syntax tree 

394 will already have been processed by IPython's 

395 TransformerManager().transform_cell. 

396 

397 For example, 

398 

399 %%time\n 

400 foo() 

401 

402 would have been transformed to 

403 

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

405 

406 and we look for instances of the latter. 

407 """ 

408 

409 def __init__(self, cell_magic: CellMagic | None = None) -> None: 

410 self.cell_magic = cell_magic 

411 

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

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

414 if ( 

415 isinstance(node.value, ast.Call) 

416 and _is_ipython_magic(node.value.func) 

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

418 ): 

419 args = _get_str_args(node.value.args) 

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

421 self.generic_visit(node) 

422 

423 

424@dataclasses.dataclass(frozen=True) 

425class OffsetAndMagic: 

426 col_offset: int 

427 magic: str 

428 

429 

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

431# as mypyc will generate broken code. 

432class MagicFinder(ast.NodeVisitor): 

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

434 

435 Note that the source of the abstract syntax tree 

436 will already have been processed by IPython's 

437 TransformerManager().transform_cell. 

438 

439 For example, 

440 

441 %matplotlib inline 

442 

443 would have been transformed to 

444 

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

446 

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

448 types of magics). 

449 """ 

450 

451 def __init__(self) -> None: 

452 self.magics: dict[int, list[OffsetAndMagic]] = collections.defaultdict(list) 

453 

454 def visit_Assign(self, node: ast.Assign | ast.AnnAssign) -> None: 

455 """Look for system assign magics. 

456 

457 For example, 

458 

459 black_version = !black --version 

460 env = %env var 

461 files: list[str] = !ls 

462 

463 would have been (respectively) transformed to 

464 

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

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

467 files: list[str] = get_ipython().getoutput('ls') 

468 

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

470 """ 

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

472 args = _get_str_args(node.value.args) 

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

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

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

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

477 if args[1]: 

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

479 else: 

480 raise AssertionError( 

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

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

483 ) from None 

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

485 OffsetAndMagic(node.value.col_offset, src) 

486 ) 

487 self.generic_visit(node) 

488 

489 def visit_AnnAssign(self, node: ast.AnnAssign) -> None: 

490 """Look for system assign magics with a type annotation.""" 

491 self.visit_Assign(node) 

492 

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

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

495 

496 For examples, 

497 

498 !ls 

499 !!ls 

500 ?ls 

501 ??ls 

502 

503 would (respectively) get transformed to 

504 

505 get_ipython().system('ls') 

506 get_ipython().getoutput('ls') 

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

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

509 

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

511 """ 

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

513 args = _get_str_args(node.value.args) 

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

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

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

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

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

519 else: 

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

521 if args[1]: 

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

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

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

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

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

527 else: 

528 raise NothingChanged # unsupported magic. 

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

530 OffsetAndMagic(node.value.col_offset, src) 

531 ) 

532 self.generic_visit(node)