prepare release 2021.11b1 (#2622)

[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 63c8aafe35b70e4d16b7d59f349feeb6f392601a..2fe6739209dd626597a43011b8bc576ae203cc9a 100644 (file)
--- a/src/black/handle_ipynb_magics.py
+++ b/src/black/handle_ipynb_magics.py
@@ -39,18 +39,18 @@ TOKENS_TO_IGNORE = frozenset(
 )
 NON_PYTHON_CELL_MAGICS = 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
     )
 )
 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
     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
         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
 
 
     return f"{mask}\n{cell_magic_finder.cell_magic.body}", replacements
 
 


@@ -311,13 +312,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:
 @dataclasses.dataclass(frozen=True)
 class CellMagic:

-    header: str
+    name: str
+    params: Optional[str]

     body: 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.
 
 class CellMagicFinder(ast.NodeVisitor):
     """Find cell magics.
 


@@ -336,7 +352,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."""


@@ -345,14 +362,8 @@ class CellMagicFinder(ast.NodeVisitor):
             and _is_ipython_magic(node.value.func)
             and node.value.func.attr == "run_cell_magic"
         ):
             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)
 
 
         self.generic_visit(node)
 
 


@@ -362,7 +373,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.
 


@@ -382,9 +394,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.


@@ -404,12 +415,8 @@ class MagicFinder(ast.NodeVisitor):
             and _is_ipython_magic(node.value.func)
             and node.value.func.attr == "getoutput"
         ):
             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)
             )
             self.magics[node.value.lineno].append(
                 OffsetAndMagic(node.value.col_offset, src)
             )


@@ -435,11 +442,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):
         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]}"
             if node.value.func.attr == "run_line_magic":
                 if args[0] == "pinfo":
                     src = f"?{args[1]}"