X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/a8b4665e7d6eb945c47820adb1a3f8b006adce0c..abdc31cd4fb6cf88339f8ade0adf9c5356a95aa3:/src/black/handle_ipynb_magics.py?ds=inline diff --git a/src/black/handle_ipynb_magics.py b/src/black/handle_ipynb_magics.py index b18f862..a0ed56b 100644 --- a/src/black/handle_ipynb_magics.py +++ b/src/black/handle_ipynb_magics.py @@ -37,22 +37,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) @@ -188,10 +184,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( @@ -229,10 +225,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 @@ -310,13 +305,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. @@ -335,7 +345,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.""" @@ -344,14 +355,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) @@ -361,7 +366,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. @@ -381,9 +387,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. @@ -391,24 +396,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) ) @@ -434,11 +443,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]}" @@ -447,7 +452,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]}"