X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/26970742b76ddfd921fa36e4005738048477c952..d2c938eb02c414057aa2186c7ae695d5d0d14377:/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 f10eaed..a0ed56b 100644 --- a/src/black/handle_ipynb_magics.py +++ b/src/black/handle_ipynb_magics.py @@ -37,20 +37,15 @@ 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 @@ -230,8 +225,6 @@ 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.name in NON_PYTHON_CELL_MAGICS: - raise NothingChanged header = cell_magic_finder.cell_magic.header mask = get_token(src, header) replacements.append(Replacement(mask=mask, src=header)) @@ -333,7 +326,7 @@ class CellMagic: return f"%%{self.name}" -@dataclasses.dataclass +# ast.NodeVisitor + dataclass = breakage under mypyc. class CellMagicFinder(ast.NodeVisitor): """Find cell magics. @@ -352,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.""" @@ -372,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. @@ -392,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. @@ -402,20 +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" - ): - (arg,) = _get_str_args(node.value.args) - src = f"!{arg}" + 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) ) @@ -450,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]}"