"""Functions to process IPython magics with."""
+
from functools import lru_cache
import dataclasses
import ast
-from typing import Dict
+from typing import Dict, List, Tuple, Optional
import secrets
-from typing import List, Tuple
+import sys
import collections
-from typing import Optional
-from typing_extensions import TypeGuard
+if sys.version_info >= (3, 10):
+ from typing import TypeGuard
+else:
+ from typing_extensions import TypeGuard
+
from black.report import NothingChanged
from black.output import out
"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)
"""
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(
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
)
+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.
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."""
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)
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.
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.
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)
)
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]}"
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]}"