src/black/handle_ipynb_magics.py

   1 """Functions to process IPython magics with."""
   2
   3 from functools import lru_cache
   4 import dataclasses
   5 import ast
   6 from typing import Dict, List, Tuple, Optional
   7
   8 import secrets
   9 import sys
  10 import collections
  11
  12 if sys.version_info >= (3, 10):
  13     from typing import TypeGuard
  14 else:
  15     from typing_extensions import TypeGuard
  16
  17 from black.report import NothingChanged
  18 from black.output import out
  19
  20
  21 TRANSFORMED_MAGICS = frozenset(
  22     (
  23         "get_ipython().run_cell_magic",
  24         "get_ipython().system",
  25         "get_ipython().getoutput",
  26         "get_ipython().run_line_magic",
  27     )
  28 )
  29 TOKENS_TO_IGNORE = frozenset(
  30     (
  31         "ENDMARKER",
  32         "NL",
  33         "NEWLINE",
  34         "COMMENT",
  35         "DEDENT",
  36         "UNIMPORTANT_WS",
  37         "ESCAPED_NL",
  38     )
  39 )
  40 PYTHON_CELL_MAGICS = frozenset(
  41     (
  42         "capture",
  43         "prun",
  44         "pypy",
  45         "python",
  46         "python3",
  47         "time",
  48         "timeit",
  49     )
  50 )
  51 TOKEN_HEX = secrets.token_hex
  52
  53
  54 @dataclasses.dataclass(frozen=True)
  55 class Replacement:
  56     mask: str
  57     src: str
  58
  59
  60 @lru_cache()
  61 def jupyter_dependencies_are_installed(*, verbose: bool, quiet: bool) -> bool:
  62     try:
  63         import IPython  # noqa:F401
  64         import tokenize_rt  # noqa:F401
  65     except ModuleNotFoundError:
  66         if verbose or not quiet:
  67             msg = (
  68                 "Skipping .ipynb files as Jupyter dependencies are not installed.\n"
  69                 "You can fix this by running ``pip install black[jupyter]``"
  70             )
  71             out(msg)
  72         return False
  73     else:
  74         return True
  75
  76
  77 def remove_trailing_semicolon(src: str) -> Tuple[str, bool]:
  78     """Remove trailing semicolon from Jupyter notebook cell.
  79
  80     For example,
  81
  82         fig, ax = plt.subplots()
  83         ax.plot(x_data, y_data);  # plot data
  84
  85     would become
  86
  87         fig, ax = plt.subplots()
  88         ax.plot(x_data, y_data)  # plot data
  89
  90     Mirrors the logic in `quiet` from `IPython.core.displayhook`, but uses
  91     ``tokenize_rt`` so that round-tripping works fine.
  92     """
  93     from tokenize_rt import (
  94         src_to_tokens,
  95         tokens_to_src,
  96         reversed_enumerate,
  97     )
  98
  99     tokens = src_to_tokens(src)
 100     trailing_semicolon = False
 101     for idx, token in reversed_enumerate(tokens):
 102         if token.name in TOKENS_TO_IGNORE:
 103             continue
 104         if token.name == "OP" and token.src == ";":
 105             del tokens[idx]
 106             trailing_semicolon = True
 107         break
 108     if not trailing_semicolon:
 109         return src, False
 110     return tokens_to_src(tokens), True
 111
 112
 113 def put_trailing_semicolon_back(src: str, has_trailing_semicolon: bool) -> str:
 114     """Put trailing semicolon back if cell originally had it.
 115
 116     Mirrors the logic in `quiet` from `IPython.core.displayhook`, but uses
 117     ``tokenize_rt`` so that round-tripping works fine.
 118     """
 119     if not has_trailing_semicolon:
 120         return src
 121     from tokenize_rt import src_to_tokens, tokens_to_src, reversed_enumerate
 122
 123     tokens = src_to_tokens(src)
 124     for idx, token in reversed_enumerate(tokens):
 125         if token.name in TOKENS_TO_IGNORE:
 126             continue
 127         tokens[idx] = token._replace(src=token.src + ";")
 128         break
 129     else:  # pragma: nocover
 130         raise AssertionError(
 131             "INTERNAL ERROR: Was not able to reinstate trailing semicolon. "
 132             "Please report a bug on https://github.com/psf/black/issues.  "
 133         ) from None
 134     return str(tokens_to_src(tokens))
 135
 136
 137 def mask_cell(src: str) -> Tuple[str, List[Replacement]]:
 138     """Mask IPython magics so content becomes parseable Python code.
 139
 140     For example,
 141
 142         %matplotlib inline
 143         'foo'
 144
 145     becomes
 146
 147         "25716f358c32750e"
 148         'foo'
 149
 150     The replacements are returned, along with the transformed code.
 151     """
 152     replacements: List[Replacement] = []
 153     try:
 154         ast.parse(src)
 155     except SyntaxError:
 156         # Might have IPython magics, will process below.
 157         pass
 158     else:
 159         # Syntax is fine, nothing to mask, early return.
 160         return src, replacements
 161
 162     from IPython.core.inputtransformer2 import TransformerManager
 163
 164     transformer_manager = TransformerManager()
 165     transformed = transformer_manager.transform_cell(src)
 166     transformed, cell_magic_replacements = replace_cell_magics(transformed)
 167     replacements += cell_magic_replacements
 168     transformed = transformer_manager.transform_cell(transformed)
 169     transformed, magic_replacements = replace_magics(transformed)
 170     if len(transformed.splitlines()) != len(src.splitlines()):
 171         # Multi-line magic, not supported.
 172         raise NothingChanged
 173     replacements += magic_replacements
 174     return transformed, replacements
 175
 176
 177 def get_token(src: str, magic: str) -> str:
 178     """Return randomly generated token to mask IPython magic with.
 179
 180     For example, if 'magic' was `%matplotlib inline`, then a possible
 181     token to mask it with would be `"43fdd17f7e5ddc83"`. The token
 182     will be the same length as the magic, and we make sure that it was
 183     not already present anywhere else in the cell.
 184     """
 185     assert magic
 186     nbytes = max(len(magic) // 2 - 1, 1)
 187     token = TOKEN_HEX(nbytes)
 188     counter = 0
 189     while token in src:
 190         token = TOKEN_HEX(nbytes)
 191         counter += 1
 192         if counter > 100:
 193             raise AssertionError(
 194                 "INTERNAL ERROR: Black was not able to replace IPython magic. "
 195                 "Please report a bug on https://github.com/psf/black/issues.  "
 196                 f"The magic might be helpful: {magic}"
 197             ) from None
 198     if len(token) + 2 < len(magic):
 199         token = f"{token}."
 200     return f'"{token}"'
 201
 202
 203 def replace_cell_magics(src: str) -> Tuple[str, List[Replacement]]:
 204     """Replace cell magic with token.
 205
 206     Note that 'src' will already have been processed by IPython's
 207     TransformerManager().transform_cell.
 208
 209     Example,
 210
 211         get_ipython().run_cell_magic('t', '-n1', 'ls =!ls\\n')
 212
 213     becomes
 214
 215         "a794."
 216         ls =!ls
 217
 218     The replacement, along with the transformed code, is returned.
 219     """
 220     replacements: List[Replacement] = []
 221
 222     tree = ast.parse(src)
 223
 224     cell_magic_finder = CellMagicFinder()
 225     cell_magic_finder.visit(tree)
 226     if cell_magic_finder.cell_magic is None:
 227         return src, replacements
 228     header = cell_magic_finder.cell_magic.header
 229     mask = get_token(src, header)
 230     replacements.append(Replacement(mask=mask, src=header))
 231     return f"{mask}\n{cell_magic_finder.cell_magic.body}", replacements
 232
 233
 234 def replace_magics(src: str) -> Tuple[str, List[Replacement]]:
 235     """Replace magics within body of cell.
 236
 237     Note that 'src' will already have been processed by IPython's
 238     TransformerManager().transform_cell.
 239
 240     Example, this
 241
 242         get_ipython().run_line_magic('matplotlib', 'inline')
 243         'foo'
 244
 245     becomes
 246
 247         "5e67db56d490fd39"
 248         'foo'
 249
 250     The replacement, along with the transformed code, are returned.
 251     """
 252     replacements = []
 253     magic_finder = MagicFinder()
 254     magic_finder.visit(ast.parse(src))
 255     new_srcs = []
 256     for i, line in enumerate(src.splitlines(), start=1):
 257         if i in magic_finder.magics:
 258             offsets_and_magics = magic_finder.magics[i]
 259             if len(offsets_and_magics) != 1:  # pragma: nocover
 260                 raise AssertionError(
 261                     f"Expecting one magic per line, got: {offsets_and_magics}\n"
 262                     "Please report a bug on https://github.com/psf/black/issues."
 263                 )
 264             col_offset, magic = (
 265                 offsets_and_magics[0].col_offset,
 266                 offsets_and_magics[0].magic,
 267             )
 268             mask = get_token(src, magic)
 269             replacements.append(Replacement(mask=mask, src=magic))
 270             line = line[:col_offset] + mask
 271         new_srcs.append(line)
 272     return "\n".join(new_srcs), replacements
 273
 274
 275 def unmask_cell(src: str, replacements: List[Replacement]) -> str:
 276     """Remove replacements from cell.
 277
 278     For example
 279
 280         "9b20"
 281         foo = bar
 282
 283     becomes
 284
 285         %%time
 286         foo = bar
 287     """
 288     for replacement in replacements:
 289         src = src.replace(replacement.mask, replacement.src)
 290     return src
 291
 292
 293 def _is_ipython_magic(node: ast.expr) -> TypeGuard[ast.Attribute]:
 294     """Check if attribute is IPython magic.
 295
 296     Note that the source of the abstract syntax tree
 297     will already have been processed by IPython's
 298     TransformerManager().transform_cell.
 299     """
 300     return (
 301         isinstance(node, ast.Attribute)
 302         and isinstance(node.value, ast.Call)
 303         and isinstance(node.value.func, ast.Name)
 304         and node.value.func.id == "get_ipython"
 305     )
 306
 307
 308 def _get_str_args(args: List[ast.expr]) -> List[str]:
 309     str_args = []
 310     for arg in args:
 311         assert isinstance(arg, ast.Str)
 312         str_args.append(arg.s)
 313     return str_args
 314
 315
 316 @dataclasses.dataclass(frozen=True)
 317 class CellMagic:
 318     name: str
 319     params: Optional[str]
 320     body: str
 321
 322     @property
 323     def header(self) -> str:
 324         if self.params:
 325             return f"%%{self.name} {self.params}"
 326         return f"%%{self.name}"
 327
 328
 329 # ast.NodeVisitor + dataclass = breakage under mypyc.
 330 class CellMagicFinder(ast.NodeVisitor):
 331     """Find cell magics.
 332
 333     Note that the source of the abstract syntax tree
 334     will already have been processed by IPython's
 335     TransformerManager().transform_cell.
 336
 337     For example,
 338
 339         %%time\nfoo()
 340
 341     would have been transformed to
 342
 343         get_ipython().run_cell_magic('time', '', 'foo()\\n')
 344
 345     and we look for instances of the latter.
 346     """
 347
 348     def __init__(self, cell_magic: Optional[CellMagic] = None) -> None:
 349         self.cell_magic = cell_magic
 350
 351     def visit_Expr(self, node: ast.Expr) -> None:
 352         """Find cell magic, extract header and body."""
 353         if (
 354             isinstance(node.value, ast.Call)
 355             and _is_ipython_magic(node.value.func)
 356             and node.value.func.attr == "run_cell_magic"
 357         ):
 358             args = _get_str_args(node.value.args)
 359             self.cell_magic = CellMagic(name=args[0], params=args[1], body=args[2])
 360         self.generic_visit(node)
 361
 362
 363 @dataclasses.dataclass(frozen=True)
 364 class OffsetAndMagic:
 365     col_offset: int
 366     magic: str
 367
 368
 369 # Unsurprisingly, subclassing ast.NodeVisitor means we can't use dataclasses here
 370 # as mypyc will generate broken code.
 371 class MagicFinder(ast.NodeVisitor):
 372     """Visit cell to look for get_ipython calls.
 373
 374     Note that the source of the abstract syntax tree
 375     will already have been processed by IPython's
 376     TransformerManager().transform_cell.
 377
 378     For example,
 379
 380         %matplotlib inline
 381
 382     would have been transformed to
 383
 384         get_ipython().run_line_magic('matplotlib', 'inline')
 385
 386     and we look for instances of the latter (and likewise for other
 387     types of magics).
 388     """
 389
 390     def __init__(self) -> None:
 391         self.magics: Dict[int, List[OffsetAndMagic]] = collections.defaultdict(list)
 392
 393     def visit_Assign(self, node: ast.Assign) -> None:
 394         """Look for system assign magics.
 395
 396         For example,
 397
 398             black_version = !black --version
 399             env = %env var
 400
 401         would have been (respectively) transformed to
 402
 403             black_version = get_ipython().getoutput('black --version')
 404             env = get_ipython().run_line_magic('env', 'var')
 405
 406         and we look for instances of any of the latter.
 407         """
 408         if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func):
 409             args = _get_str_args(node.value.args)
 410             if node.value.func.attr == "getoutput":
 411                 src = f"!{args[0]}"
 412             elif node.value.func.attr == "run_line_magic":
 413                 src = f"%{args[0]}"
 414                 if args[1]:
 415                     src += f" {args[1]}"
 416             else:
 417                 raise AssertionError(
 418                     f"Unexpected IPython magic {node.value.func.attr!r} found. "
 419                     "Please report a bug on https://github.com/psf/black/issues."
 420                 ) from None
 421             self.magics[node.value.lineno].append(
 422                 OffsetAndMagic(node.value.col_offset, src)
 423             )
 424         self.generic_visit(node)
 425
 426     def visit_Expr(self, node: ast.Expr) -> None:
 427         """Look for magics in body of cell.
 428
 429         For examples,
 430
 431             !ls
 432             !!ls
 433             ?ls
 434             ??ls
 435
 436         would (respectively) get transformed to
 437
 438             get_ipython().system('ls')
 439             get_ipython().getoutput('ls')
 440             get_ipython().run_line_magic('pinfo', 'ls')
 441             get_ipython().run_line_magic('pinfo2', 'ls')
 442
 443         and we look for instances of any of the latter.
 444         """
 445         if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func):
 446             args = _get_str_args(node.value.args)
 447             if node.value.func.attr == "run_line_magic":
 448                 if args[0] == "pinfo":
 449                     src = f"?{args[1]}"
 450                 elif args[0] == "pinfo2":
 451                     src = f"??{args[1]}"
 452                 else:
 453                     src = f"%{args[0]}"
 454                     if args[1]:
 455                         src += f" {args[1]}"
 456             elif node.value.func.attr == "system":
 457                 src = f"!{args[0]}"
 458             elif node.value.func.attr == "getoutput":
 459                 src = f"!!{args[0]}"
 460             else:
 461                 raise NothingChanged  # unsupported magic.
 462             self.magics[node.value.lineno].append(
 463                 OffsetAndMagic(node.value.col_offset, src)
 464             )
 465         self.generic_visit(node)