Highlight shell escapes inside timing and profiling cell magics

ipython_pygments_lexers.py

@@ -53,6 +53,7 @@
     do_insertions,
     bygroups,
     using,
+    this,
 )
 from pygments.token import (
     Generic,
@@ -78,10 +79,17 @@
 ]
 
 
+# Cell magics whose body is IPython input delegate to this lexer's own root
+# state, so shell escapes (`!cmd`), line magics and the rest of the IPython
+# grammar keep working inside them. The remaining Python-running magics keep
+# delegating to `Python3Lexer`, because their bodies are not IPython input:
+# `%%debug` runs the body under pdb, `%%python`/`%%python2`/`%%python3`/`%%pypy`
+# run it in a separate interpreter, and `%%writefile`/`%%file` treat it as
+# literal file content.
 ipython_tokens = [
     (
         r"(?s)(\s*)(%%capture)([^\n]*\n)(.*)",
-        bygroups(Text, Operator, Text, using(Python3Lexer)),
+        bygroups(Text, Operator, Text, using(this)),
     ),
     (
         r"(?s)(\s*)(%%debug)([^\n]*\n)(.*)",
@@ -109,7 +117,7 @@
     ),
     (
         r"(?s)(\s*)(%%prun)([^\n]*\n)(.*)",
-        bygroups(Text, Operator, Text, using(Python3Lexer)),
+        bygroups(Text, Operator, Text, using(this)),
     ),
     (
         r"(?s)(\s*)(%%pypy)([^\n]*\n)(.*)",
@@ -133,11 +141,11 @@
     ),
     (
         r"(?s)(\s*)(%%timeit)([^\n]*\n)(.*)",
-        bygroups(Text, Operator, Text, using(Python3Lexer)),
+        bygroups(Text, Operator, Text, using(this)),
     ),
     (
         r"(?s)(\s*)(%%time)([^\n]*\n)(.*)",
-        bygroups(Text, Operator, Text, using(Python3Lexer)),
+        bygroups(Text, Operator, Text, using(this)),
     ),
     (
         r"(?s)(\s*)(%%writefile)([^\n]*\n)(.*)",

test_ipython_pygments_lexers.py

@@ -202,3 +202,54 @@ def test_cell_magics():
         (Token.Text, "\n$foo\n"),
     ]
     assert tokens_2 == list(lexer.get_tokens(fragment_2))
+
+
+@pytest.mark.parametrize("magic", ["%%time", "%%timeit", "%%capture", "%%prun"])
+def test_shell_commands_inside_ipython_input_cell_magics(magic):
+    """Shell escapes keep working inside cell magics that run IPython input.
+
+    ``%%time``, ``%%timeit``, ``%%capture`` and ``%%prun`` all pass their body
+    through the IPython input transformer, so ``!cmd`` reaches the system shell.
+    Lexing the body as plain Python instead flags the ``!`` as an error.
+    """
+    lexer = lexers.IPythonLexer()
+    fragment = f"{magic}\n!cmd\n"
+    tokens = [
+        (Token.Operator, magic),
+        (Token.Text, "\n"),
+        (Token.Operator, "!"),
+        (Token.Text, "cmd"),
+        (TOKEN_WS, "\n"),
+    ]
+    # The trailing newline's token type varies with the pygments version, so
+    # compare the body tokens, as ``test_shell_commands`` does.
+    assert tokens[:-1] == list(lexer.get_tokens(fragment))[:-1]
+
+
+@pytest.mark.parametrize("magic", ["%%debug", "%%python", "%%writefile foo.py"])
+def test_no_shell_commands_where_the_body_is_not_ipython_input(magic):
+    """A ``!`` is not a shell escape in cell magics that do not run IPython input.
+
+    ``%%debug`` runs the body under pdb, ``%%python`` runs it in a separate
+    interpreter, and ``%%writefile`` writes it verbatim, so these bodies stay
+    plain Python and must not gain the IPython shell-escape rule.
+    """
+    lexer = lexers.IPythonLexer()
+    tokens = list(lexer.get_tokens(f"{magic}\n!cmd\n"))
+    assert (Token.Error, "!") in tokens
+    assert (Token.Operator, "!") not in tokens
+
+
+def test_python_still_lexes_inside_a_timing_magic():
+    """Delegating to this lexer's root keeps ordinary Python highlighting."""
+    lexer = lexers.IPythonLexer()
+    fragment = "%%time\nimport math\n"
+    tokens = [
+        (Token.Operator, "%%time"),
+        (Token.Text, "\n"),
+        (Token.Keyword.Namespace, "import"),
+        (TOKEN_WS, " "),
+        (Token.Name.Namespace, "math"),
+        (TOKEN_WS, "\n"),
+    ]
+    assert tokens[:-1] == list(lexer.get_tokens(fragment))[:-1]