From 903e86089737204553403c04364730e12f8e95c1 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 29 Nov 2023 15:22:11 -0600
Subject: [PATCH] Allow transparent cell magics

---
 .../fixtures/jupyter/cell/cell_magic.json     |  7 ++-
 .../jupyter/cell/valid_cell_magic.json        | 11 +++++
 crates/ruff_notebook/src/cell.rs              | 45 ++++++++++++++++++-
 crates/ruff_notebook/src/notebook.rs          |  1 +
 4 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 crates/ruff_notebook/resources/test/fixtures/jupyter/cell/valid_cell_magic.json

diff --git a/crates/ruff_notebook/resources/test/fixtures/jupyter/cell/cell_magic.json b/crates/ruff_notebook/resources/test/fixtures/jupyter/cell/cell_magic.json
index ef68b202e6811..e0de8c0241141 100644
--- a/crates/ruff_notebook/resources/test/fixtures/jupyter/cell/cell_magic.json
+++ b/crates/ruff_notebook/resources/test/fixtures/jupyter/cell/cell_magic.json
@@ -4,5 +4,10 @@
   "id": "1",
   "metadata": {},
   "outputs": [],
-  "source": ["%%timeit\n", "print('hello world')"]
+  "source": [
+    "%%script bash\n",
+    "for i in 1 2 3; do\n",
+    "  echo $i\n",
+    "done"
+  ]
 }
diff --git a/crates/ruff_notebook/resources/test/fixtures/jupyter/cell/valid_cell_magic.json b/crates/ruff_notebook/resources/test/fixtures/jupyter/cell/valid_cell_magic.json
new file mode 100644
index 0000000000000..2cb89fa63b9bd
--- /dev/null
+++ b/crates/ruff_notebook/resources/test/fixtures/jupyter/cell/valid_cell_magic.json
@@ -0,0 +1,11 @@
+{
+  "execution_count": null,
+  "cell_type": "code",
+  "id": "1",
+  "metadata": {},
+  "outputs": [],
+  "source": [
+    "%%timeit\n",
+    "print('hello world')"
+  ]
+}
diff --git a/crates/ruff_notebook/src/cell.rs b/crates/ruff_notebook/src/cell.rs
index d80ef336de02e..caa4cb3204a03 100644
--- a/crates/ruff_notebook/src/cell.rs
+++ b/crates/ruff_notebook/src/cell.rs
@@ -170,7 +170,50 @@ impl Cell {
         }
 
         // Detect cell magics (which operate on multiple lines).
-        lines.any(|line| line.trim_start().starts_with("%%"))
+        lines.any(|line| {
+            line.split_whitespace().next().is_some_and(|first| {
+                if first.len() < 2 {
+                    return false;
+                }
+                let (token, command) = first.split_at(2);
+                // These cell magics are special in that the lines following them are valid
+                // Python code and the variables defined in that scope are available to the
+                // rest of the notebook.
+                //
+                // For example:
+                //
+                // Cell 1:
+                // ```python
+                // x = 1
+                // ```
+                //
+                // Cell 2:
+                // ```python
+                // %%time
+                // y = x
+                // ```
+                //
+                // Cell 3:
+                // ```python
+                // print(y)  # Here, `y` is available.
+                // ```
+                //
+                // This is to avoid false positives when these variables are referenced
+                // elsewhere in the notebook.
+                token == "%%"
+                    && !matches!(
+                        command,
+                        "capture"
+                            | "debug"
+                            | "prun"
+                            | "pypy"
+                            | "python"
+                            | "python3"
+                            | "time"
+                            | "timeit"
+                    )
+            })
+        })
     }
 }
 
diff --git a/crates/ruff_notebook/src/notebook.rs b/crates/ruff_notebook/src/notebook.rs
index a6714fa12b496..7c9e8356b79d3 100644
--- a/crates/ruff_notebook/src/notebook.rs
+++ b/crates/ruff_notebook/src/notebook.rs
@@ -426,6 +426,7 @@ mod tests {
     #[test_case(Path::new("code_and_magic.json"), true; "code_and_magic")]
     #[test_case(Path::new("only_code.json"), true; "only_code")]
     #[test_case(Path::new("cell_magic.json"), false; "cell_magic")]
+    #[test_case(Path::new("valid_cell_magic.json"), true; "valid_cell_magic")]
     #[test_case(Path::new("automagic.json"), false; "automagic")]
     #[test_case(Path::new("automagics.json"), false; "automagics")]
     #[test_case(Path::new("automagic_before_code.json"), false; "automagic_before_code")]