]> git.madduck.net Git - etc/vim.git/blobdiff - src/black/handle_ipynb_magics.py

madduck's git repository

Every one of the projects in this repository is available at the canonical URL git://git.madduck.net/madduck/pub/<projectpath> — see each project's metadata for the exact URL.

All patches and comments are welcome. Please squash your changes to logical commits before using git-format-patch and git-send-email to patches@git.madduck.net. If you'd read over the Git project's submission guidelines and adhered to them, I'd be especially grateful.

SSH access, as well as push access can be individually arranged.

If you use my repositories frequently, consider adding the following snippet to ~/.gitconfig and using the third clone URL listed for each project:

[url "git://git.madduck.net/madduck/"]
  insteadOf = madduck:

Add a missing space in Python 2 deprecation (GH-2590)
[etc/vim.git] / src / black / handle_ipynb_magics.py
index 63c8aafe35b70e4d16b7d59f349feeb6f392601a..f10eaed4f3ef950e6671c59e169ee967e78ca86e 100644 (file)
@@ -39,18 +39,18 @@ TOKENS_TO_IGNORE = frozenset(
 )
 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 @@ 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)
 
 
@@ -404,12 +414,8 @@ class MagicFinder(ast.NodeVisitor):
             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 @@ 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]}"