Handle cases where unknown %% magic command is preceded by empty lines

and lines with comments.

PiperOrigin-RevId: 675938311

patches/pyink.patch

@@ -370,6 +370,20 @@
      )
 
      lines: List[Tuple[int, int]] = []
+@@ -1153,9 +677,10 @@
+     """
+     if any(transformed_magic in src for transformed_magic in TRANSFORMED_MAGICS):
+         raise NothingChanged
+-    if (
+-        src[:2] == "%%"
+-        and src.split()[0][2:] not in PYTHON_CELL_MAGICS | mode.python_cell_magics
++    line = ink.get_code_start(src)
++    if line.startswith("%%") and (
++        line.split(maxsplit=1)[0][2:]
++        not in PYTHON_CELL_MAGICS | mode.python_cell_magics
+     ):
+         raise NothingChanged
+
 @@ -1175,7 +1219,6 @@
          raise NothingChanged
 
@@ -1395,6 +1409,29 @@
 
  runner = CliRunner()
 
+@@ -209,6 +209,22 @@
+         assert result == expected_output
+
+
++@pytest.mark.parametrize(
++    "src",
++    (
++        "   %%custom_magic \nx=2",
++        "\n\n%%custom_magic\nx=2",
++        "# comment\n%%custom_magic\nx=2",
++        "\n  \n # comment with %%time\n\t\n %%custom_magic # comment \nx=2",
++    ),
++)
++def test_cell_magic_with_custom_python_magic_after_spaces_and_comments_noop(
++    src: str,
++) -> None:
++    with pytest.raises(NothingChanged):
++        format_cell(src, fast=True, mode=JUPYTER_MODE)
++
++
+ def test_cell_magic_nested() -> None:
+     src = "%%time\n%%time\n2+2"
+     result = format_cell(src, fast=True, mode=JUPYTER_MODE)
 @@ -385,6 +385,45 @@
      assert result == expected
 

src/pyink/__init__.py

@@ -1150,9 +1150,10 @@ def validate_cell(src: str, mode: Mode) -> None:
     """
     if any(transformed_magic in src for transformed_magic in TRANSFORMED_MAGICS):
         raise NothingChanged
-    if (
-        src[:2] == "%%"
-        and src.split()[0][2:] not in PYTHON_CELL_MAGICS | mode.python_cell_magics
+    line = ink.get_code_start(src)
+    if line.startswith("%%") and (
+        line.split(maxsplit=1)[0][2:]
+        not in PYTHON_CELL_MAGICS | mode.python_cell_magics
     ):
         raise NothingChanged
 

src/pyink/ink.py

@@ -3,6 +3,7 @@
 This is a separate module for easier patch management.
 """
 
+import re
 from typing import (
     Collection,
     Iterator,
@@ -62,6 +63,26 @@ def majority_quote(node: LN) -> Quote:
     return Quote.DOUBLE
 
 
+def get_code_start(src: str) -> str:
+    """Provides the first line where the code starts.
+
+    Iterates over lines of code until it finds the first line that doesn't
+    contain only empty spaces and comments. If such line doesn't exist, it
+    returns an empty string.
+
+    Args:
+      src: The multi-line source code.
+
+    Returns:
+      The first line of code without initial spaces or an empty string.
+    """
+    for match in re.finditer(".+", src):
+        line = match.group(0).lstrip()
+        if line and not line.startswith("#"):
+            return line
+    return ""
+
+
 def convert_unchanged_lines(src_node: Node, lines: Collection[Tuple[int, int]]):
     """Converts unchanged lines to STANDALONE_COMMENT.
 

tests/test_ipynb.py

@@ -210,6 +210,22 @@ def test_cell_magic_with_custom_python_magic(
         assert result == expected_output
 
 
+@pytest.mark.parametrize(
+    "src",
+    (
+        "   %%custom_magic \nx=2",
+        "\n\n%%custom_magic\nx=2",
+        "# comment\n%%custom_magic\nx=2",
+        "\n  \n # comment with %%time\n\t\n %%custom_magic # comment \nx=2",
+    ),
+)
+def test_cell_magic_with_custom_python_magic_after_spaces_and_comments_noop(
+    src: str,
+) -> None:
+    with pytest.raises(NothingChanged):
+        format_cell(src, fast=True, mode=JUPYTER_MODE)
+
+
 def test_cell_magic_nested() -> None:
     src = "%%time\n%%time\n2+2"
     result = format_cell(src, fast=True, mode=JUPYTER_MODE)