Document shebang comment behaviour (#3787)

[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 ad93c444efc98a7dfd2c12b54dc4cd834bb08b2f..2a2d62220e2797c63af1cbeb6dbd8b2e41ac9ba7 100644 (file)
--- a/src/black/handle_ipynb_magics.py
+++ b/src/black/handle_ipynb_magics.py
@@ -1,18 +1,20 @@
 """Functions to process IPython magics with."""
 """Functions to process IPython magics with."""

-from functools import lru_cache
-import dataclasses
-import ast
-from typing import Dict

 
 

-import secrets
-from typing import List, Tuple
+import ast

 import collections
 import collections

+import dataclasses
+import secrets
+import sys
+from functools import lru_cache
+from typing import Dict, List, Optional, Tuple

 
 

-from typing import Optional
-from typing_extensions import TypeGuard
-from black.report import NothingChanged
-from black.output import out
+if sys.version_info >= (3, 10):
+    from typing import TypeGuard
+else:
+    from typing_extensions import TypeGuard

 
 

+from black.output import out
+from black.report import NothingChanged

 
 TRANSFORMED_MAGICS = frozenset(
     (
 
 TRANSFORMED_MAGICS = frozenset(
     (


@@ -33,22 +35,18 @@ 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

 
 
 @dataclasses.dataclass(frozen=True)
 
 
 @dataclasses.dataclass(frozen=True)


@@ -57,16 +55,21 @@ class Replacement:
     src: str
 
 
     src: str
 
 

-@lru_cache()
+@lru_cache

 def jupyter_dependencies_are_installed(*, verbose: bool, quiet: bool) -> bool:
     try:
 def jupyter_dependencies_are_installed(*, verbose: bool, quiet: bool) -> bool:
     try:

-        import IPython  # noqa:F401
+        # isort: off
+        # tokenize_rt is less commonly installed than IPython
+        # and IPython is expensive to import

         import tokenize_rt  # noqa:F401
         import tokenize_rt  # noqa:F401

+        import IPython  # noqa:F401
+
+        # isort: on

     except ModuleNotFoundError:
         if verbose or not quiet:
             msg = (
                 "Skipping .ipynb files as Jupyter dependencies are not installed.\n"
     except ModuleNotFoundError:
         if verbose or not quiet:
             msg = (
                 "Skipping .ipynb files as Jupyter dependencies are not installed.\n"

-                "You can fix this by running ``pip install black[jupyter]``"
+                'You can fix this by running ``pip install "black[jupyter]"``'

             )
             out(msg)
         return False
             )
             out(msg)
         return False


@@ -90,11 +93,7 @@ def remove_trailing_semicolon(src: str) -> Tuple[str, bool]:
     Mirrors the logic in `quiet` from `IPython.core.displayhook`, but uses
     ``tokenize_rt`` so that round-tripping works fine.
     """
     Mirrors the logic in `quiet` from `IPython.core.displayhook`, but uses
     ``tokenize_rt`` so that round-tripping works fine.
     """

-    from tokenize_rt import (
-        src_to_tokens,
-        tokens_to_src,
-        reversed_enumerate,
-    )
+    from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src

 
     tokens = src_to_tokens(src)
     trailing_semicolon = False
 
     tokens = src_to_tokens(src)
     trailing_semicolon = False


@@ -118,7 +117,7 @@ def put_trailing_semicolon_back(src: str, has_trailing_semicolon: bool) -> str:
     """
     if not has_trailing_semicolon:
         return src
     """
     if not has_trailing_semicolon:
         return src

-    from tokenize_rt import src_to_tokens, tokens_to_src, reversed_enumerate
+    from tokenize_rt import reversed_enumerate, src_to_tokens, tokens_to_src

 
     tokens = src_to_tokens(src)
     for idx, token in reversed_enumerate(tokens):
 
     tokens = src_to_tokens(src)
     for idx, token in reversed_enumerate(tokens):


@@ -184,10 +183,10 @@ def get_token(src: str, magic: str) -> str:
     """
     assert magic
     nbytes = max(len(magic) // 2 - 1, 1)
     """
     assert magic
     nbytes = max(len(magic) // 2 - 1, 1)

-    token = secrets.token_hex(nbytes)
+    token = TOKEN_HEX(nbytes)

     counter = 0
     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(
         counter += 1
         if counter > 100:
             raise AssertionError(


@@ -225,10 +224,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
     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
 
 
     return f"{mask}\n{cell_magic_finder.cell_magic.body}", replacements
 
 


@@ -306,13 +304,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.
 


@@ -322,7 +335,8 @@ class CellMagicFinder(ast.NodeVisitor):
 
     For example,
 
 
     For example,
 

-        %%time\nfoo()
+        %%time\n
+        foo()

 
     would have been transformed to
 
 
     would have been transformed to
 


@@ -331,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."""


@@ -340,14 +355,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)
 
 


@@ -357,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.
 


@@ -377,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.


@@ -387,24 +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"
-        ):
-            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)
             )
             self.magics[node.value.lineno].append(
                 OffsetAndMagic(node.value.col_offset, src)
             )


@@ -430,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):
         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]}"


@@ -443,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]}"