Refactor Jupyter magic handling (#2545)

This commit is contained in:
Marco Edward Gorelli 2021-10-27 15:36:10 +01:00 committed by GitHub
parent 62ed5389fc
commit 26970742b7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -39,18 +39,18 @@
)
NON_PYTHON_CELL_MAGICS = frozenset(
(
"%%bash",
"%%html",
"%%javascript",
"%%js",
"%%latex",
"%%markdown",
"%%perl",
"%%ruby",
"%%script",
"%%sh",
"%%svg",
"%%writefile",
"bash",
"html",
"javascript",
"js",
"latex",
"markdown",
"perl",
"ruby",
"script",
"sh",
"svg",
"writefile",
)
)
TOKEN_HEX = secrets.token_hex
@ -230,10 +230,11 @@ 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:
if cell_magic_finder.cell_magic.name 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
@ -311,11 +312,26 @@ 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
class CellMagicFinder(ast.NodeVisitor):
@ -345,14 +361,8 @@ def visit_Expr(self, node: ast.Expr) -> None:
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)
@ -404,12 +414,8 @@ def visit_Assign(self, node: ast.Assign) -> None:
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]}"
(arg,) = _get_str_args(node.value.args)
src = f"!{arg}"
self.magics[node.value.lineno].append(
OffsetAndMagic(node.value.col_offset, src)
)
@ -435,11 +441,7 @@ def visit_Expr(self, node: ast.Expr) -> None:
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]}"