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]}"