X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/b1d060101626aa1c332f52e4bdf0ae5e4cc07990..ea5293b0360b6e2e56c004d1e27cbe33c87dcb94:/src/black/handle_ipynb_magics.py?ds=sidebyside diff --git a/src/black/handle_ipynb_magics.py b/src/black/handle_ipynb_magics.py index ad93c44..9e1af75 100644 --- a/src/black/handle_ipynb_magics.py +++ b/src/black/handle_ipynb_magics.py @@ -1,18 +1,20 @@ """Functions to process IPython magics with.""" -from functools import lru_cache -import dataclasses -import ast -from typing import Dict -import secrets -from typing import List, Tuple +import ast import collections +import dataclasses +import secrets +import sys +from functools import lru_cache +from typing import Dict, List, Optional, Tuple -from typing import Optional -from typing_extensions import TypeGuard -from black.report import NothingChanged -from black.output import out +if sys.version_info >= (3, 10): + from typing import TypeGuard +else: + from typing_extensions import TypeGuard +from black.output import out +from black.report import NothingChanged TRANSFORMED_MAGICS = frozenset( ( @@ -33,22 +35,18 @@ TOKENS_TO_IGNORE = frozenset( "ESCAPED_NL", ) ) -NON_PYTHON_CELL_MAGICS = frozenset( +PYTHON_CELL_MAGICS = frozenset( ( - "%%bash", - "%%html", - "%%javascript", - "%%js", - "%%latex", - "%%markdown", - "%%perl", - "%%ruby", - "%%script", - "%%sh", - "%%svg", - "%%writefile", + "capture", + "prun", + "pypy", + "python", + "python3", + "time", + "timeit", ) ) +TOKEN_HEX = secrets.token_hex @dataclasses.dataclass(frozen=True) @@ -66,7 +64,7 @@ def jupyter_dependencies_are_installed(*, verbose: bool, quiet: bool) -> bool: if verbose or not quiet: msg = ( "Skipping .ipynb files as Jupyter dependencies are not installed.\n" - "You can fix this by running ``pip install black[jupyter]``" + 'You can fix this by running ``pip install "black[jupyter]"``' ) out(msg) return False @@ -90,11 +88,7 @@ def remove_trailing_semicolon(src: str) -> Tuple[str, bool]: Mirrors the logic in `quiet` from `IPython.core.displayhook`, but uses ``tokenize_rt`` so that round-tripping works fine. """ - from tokenize_rt import ( - src_to_tokens, - tokens_to_src, - reversed_enumerate, - ) + from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src tokens = src_to_tokens(src) trailing_semicolon = False @@ -118,7 +112,7 @@ def put_trailing_semicolon_back(src: str, has_trailing_semicolon: bool) -> str: """ if not has_trailing_semicolon: return src - from tokenize_rt import src_to_tokens, tokens_to_src, reversed_enumerate + from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src tokens = src_to_tokens(src) for idx, token in reversed_enumerate(tokens): @@ -184,10 +178,10 @@ def get_token(src: str, magic: str) -> str: """ assert magic nbytes = max(len(magic) // 2 - 1, 1) - token = secrets.token_hex(nbytes) + token = TOKEN_HEX(nbytes) counter = 0 - while token in src: # pragma: nocover - token = secrets.token_hex(nbytes) + while token in src: + token = TOKEN_HEX(nbytes) counter += 1 if counter > 100: raise AssertionError( @@ -225,10 +219,9 @@ def replace_cell_magics(src: str) -> Tuple[str, List[Replacement]]: cell_magic_finder.visit(tree) if cell_magic_finder.cell_magic is None: return src, replacements - if cell_magic_finder.cell_magic.header.split()[0] in NON_PYTHON_CELL_MAGICS: - raise NothingChanged - mask = get_token(src, cell_magic_finder.cell_magic.header) - replacements.append(Replacement(mask=mask, src=cell_magic_finder.cell_magic.header)) + header = cell_magic_finder.cell_magic.header + mask = get_token(src, header) + replacements.append(Replacement(mask=mask, src=header)) return f"{mask}\n{cell_magic_finder.cell_magic.body}", replacements @@ -306,13 +299,28 @@ def _is_ipython_magic(node: ast.expr) -> TypeGuard[ast.Attribute]: ) +def _get_str_args(args: List[ast.expr]) -> List[str]: + str_args = [] + for arg in args: + assert isinstance(arg, ast.Str) + str_args.append(arg.s) + return str_args + + @dataclasses.dataclass(frozen=True) class CellMagic: - header: str + name: str + params: Optional[str] body: str + @property + def header(self) -> str: + if self.params: + return f"%%{self.name} {self.params}" + return f"%%{self.name}" -@dataclasses.dataclass + +# ast.NodeVisitor + dataclass = breakage under mypyc. class CellMagicFinder(ast.NodeVisitor): """Find cell magics. @@ -331,7 +339,8 @@ class CellMagicFinder(ast.NodeVisitor): and we look for instances of the latter. """ - cell_magic: Optional[CellMagic] = None + def __init__(self, cell_magic: Optional[CellMagic] = None) -> None: + self.cell_magic = cell_magic def visit_Expr(self, node: ast.Expr) -> None: """Find cell magic, extract header and body.""" @@ -340,14 +349,8 @@ class CellMagicFinder(ast.NodeVisitor): and _is_ipython_magic(node.value.func) and node.value.func.attr == "run_cell_magic" ): - args = [] - for arg in node.value.args: - assert isinstance(arg, ast.Str) - args.append(arg.s) - header = f"%%{args[0]}" - if args[1]: - header += f" {args[1]}" - self.cell_magic = CellMagic(header=header, body=args[2]) + args = _get_str_args(node.value.args) + self.cell_magic = CellMagic(name=args[0], params=args[1], body=args[2]) self.generic_visit(node) @@ -357,7 +360,8 @@ class OffsetAndMagic: magic: str -@dataclasses.dataclass +# Unsurprisingly, subclassing ast.NodeVisitor means we can't use dataclasses here +# as mypyc will generate broken code. class MagicFinder(ast.NodeVisitor): """Visit cell to look for get_ipython calls. @@ -377,9 +381,8 @@ class MagicFinder(ast.NodeVisitor): types of magics). """ - magics: Dict[int, List[OffsetAndMagic]] = dataclasses.field( - default_factory=lambda: collections.defaultdict(list) - ) + def __init__(self) -> None: + self.magics: Dict[int, List[OffsetAndMagic]] = collections.defaultdict(list) def visit_Assign(self, node: ast.Assign) -> None: """Look for system assign magics. @@ -387,24 +390,28 @@ class MagicFinder(ast.NodeVisitor): For example, black_version = !black --version + env = %env var - would have been transformed to + would have been (respectively) transformed to black_version = get_ipython().getoutput('black --version') + env = get_ipython().run_line_magic('env', 'var') - and we look for instances of the latter. + and we look for instances of any of the latter. """ - if ( - isinstance(node.value, ast.Call) - and _is_ipython_magic(node.value.func) - and node.value.func.attr == "getoutput" - ): - args = [] - for arg in node.value.args: - assert isinstance(arg, ast.Str) - args.append(arg.s) - assert args - src = f"!{args[0]}" + if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func): + args = _get_str_args(node.value.args) + if node.value.func.attr == "getoutput": + src = f"!{args[0]}" + elif node.value.func.attr == "run_line_magic": + src = f"%{args[0]}" + if args[1]: + src += f" {args[1]}" + else: + raise AssertionError( + f"Unexpected IPython magic {node.value.func.attr!r} found. " + "Please report a bug on https://github.com/psf/black/issues." + ) from None self.magics[node.value.lineno].append( OffsetAndMagic(node.value.col_offset, src) ) @@ -430,11 +437,7 @@ class MagicFinder(ast.NodeVisitor): and we look for instances of any of the latter. """ if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func): - args = [] - for arg in node.value.args: - assert isinstance(arg, ast.Str) - args.append(arg.s) - assert args + args = _get_str_args(node.value.args) if node.value.func.attr == "run_line_magic": if args[0] == "pinfo": src = f"?{args[1]}" @@ -443,7 +446,6 @@ class MagicFinder(ast.NodeVisitor): else: src = f"%{args[0]}" if args[1]: - assert src is not None src += f" {args[1]}" elif node.value.func.attr == "system": src = f"!{args[0]}"