Don't let TokenError bubble up from lib2to3_parse (GH-2343)

[etc/vim.git] / src / black / handle_ipynb_magics.py

diff --git a/src/black/handle_ipynb_magics.py b/src/black/handle_ipynb_magics.py
index f10eaed4f3ef950e6671c59e169ee967e78ca86e..a0ed56baafcf86781b8c06da06922a52d2b5308a 100644 (file)
--- a/src/black/handle_ipynb_magics.py
+++ b/src/black/handle_ipynb_magics.py
@@ -37,20 +37,15 @@ TOKENS_TO_IGNORE = frozenset(
         "ESCAPED_NL",
     )
 )
         "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
     )
 )
 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
     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))
     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}"
 
 
         return f"%%{self.name}"
 
 

-@dataclasses.dataclass
+# ast.NodeVisitor + dataclass = breakage under mypyc.

 class CellMagicFinder(ast.NodeVisitor):
     """Find cell magics.
 
 class CellMagicFinder(ast.NodeVisitor):
     """Find cell magics.
 


@@ -352,7 +345,8 @@ class CellMagicFinder(ast.NodeVisitor):
     and we look for instances of the latter.
     """
 
     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."""
 
     def visit_Expr(self, node: ast.Expr) -> None:
         """Find cell magic, extract header and body."""


@@ -372,7 +366,8 @@ class OffsetAndMagic:
     magic: str
 
 
     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.
 
 class MagicFinder(ast.NodeVisitor):
     """Visit cell to look for get_ipython calls.
 


@@ -392,9 +387,8 @@ class MagicFinder(ast.NodeVisitor):
     types of magics).
     """
 
     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.
 
     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
         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')
 
             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)
             )
             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]:
                 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]}"
                         src += f" {args[1]}"
             elif node.value.func.attr == "system":
                 src = f"!{args[0]}"