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 return src 

327 

328 

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

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

331 

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

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

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

335 empty string. 

336 """ 

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

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

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

340 return line 

341 return "" 

342 

343 

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

345 """Check if attribute is IPython magic. 

346 

347 Note that the source of the abstract syntax tree 

348 will already have been processed by IPython's 

349 TransformerManager().transform_cell. 

350 """ 

351 return ( 

352 isinstance(node, ast.Attribute) 

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

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

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

356 ) 

357 

358 

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

360 str_args = [] 

361 for arg in args: 

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

363 str_args.append(arg.value) 

364 return str_args 

365 

366 

367@dataclasses.dataclass(frozen=True) 

368class CellMagic: 

369 name: str 

370 params: str | None 

371 body: str 

372 

373 @property 

374 def header(self) -> str: 

375 if self.params: 

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

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

378 

379 

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

381class CellMagicFinder(ast.NodeVisitor): 

382 r"""Find cell magics. 

383 

384 Note that the source of the abstract syntax tree 

385 will already have been processed by IPython's 

386 TransformerManager().transform_cell. 

387 

388 For example, 

389 

390 %%time\n 

391 foo() 

392 

393 would have been transformed to 

394 

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

396 

397 and we look for instances of the latter. 

398 """ 

399 

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

401 self.cell_magic = cell_magic 

402 

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

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

405 if ( 

406 isinstance(node.value, ast.Call) 

407 and _is_ipython_magic(node.value.func) 

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

409 ): 

410 args = _get_str_args(node.value.args) 

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

412 self.generic_visit(node) 

413 

414 

415@dataclasses.dataclass(frozen=True) 

416class OffsetAndMagic: 

417 col_offset: int 

418 magic: str 

419 

420 

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

422# as mypyc will generate broken code. 

423class MagicFinder(ast.NodeVisitor): 

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

425 

426 Note that the source of the abstract syntax tree 

427 will already have been processed by IPython's 

428 TransformerManager().transform_cell. 

429 

430 For example, 

431 

432 %matplotlib inline 

433 

434 would have been transformed to 

435 

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

437 

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

439 types of magics). 

440 """ 

441 

442 def __init__(self) -> None: 

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

444 

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

446 """Look for system assign magics. 

447 

448 For example, 

449 

450 black_version = !black --version 

451 env = %env var 

452 

453 would have been (respectively) transformed to 

454 

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

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

457 

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

459 """ 

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

461 args = _get_str_args(node.value.args) 

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

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

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

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

466 if args[1]: 

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

468 else: 

469 raise AssertionError( 

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

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

472 ) from None 

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

474 OffsetAndMagic(node.value.col_offset, src) 

475 ) 

476 self.generic_visit(node) 

477 

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

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

480 

481 For examples, 

482 

483 !ls 

484 !!ls 

485 ?ls 

486 ??ls 

487 

488 would (respectively) get transformed to 

489 

490 get_ipython().system('ls') 

491 get_ipython().getoutput('ls') 

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

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

494 

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

496 """ 

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

498 args = _get_str_args(node.value.args) 

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

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

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

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

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

504 else: 

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

506 if args[1]: 

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

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

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

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

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

512 else: 

513 raise NothingChanged # unsupported magic. 

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

515 OffsetAndMagic(node.value.col_offset, src) 

516 ) 

517 self.generic_visit(node)