From 0bfb818e87141a6ef93b72ef52068cc80389e71d Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Tue, 29 Apr 2025 17:13:46 -0400
Subject: [PATCH] docs: process cell magics (#4462)

Handle a small thing that can be fairly confusing to new python users.

Before:


![image](https://github.com/user-attachments/assets/39011f0c-0a7e-4f32-94d7-a40f0b14f2ab)


After:


![image](https://github.com/user-attachments/assets/16aa4429-eb4e-44b5-8714-29279d93c7ea)
---
 docs/_scripts/notebook_convert.py             | 28 ++++++------
 .../mdoutput/index.md.j2                      | 13 ++++++
 .../unit_tests/test_notebook_conversion.py    | 43 +++++++++++++++++++
 3 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/docs/_scripts/notebook_convert.py b/docs/_scripts/notebook_convert.py
index 06badf25b..60f4dd708 100644
--- a/docs/_scripts/notebook_convert.py
+++ b/docs/_scripts/notebook_convert.py
@@ -1,7 +1,6 @@
 import ast
 import os
 import re
-from pathlib import Path
 from typing import Literal
 
 import nbformat
@@ -26,7 +25,7 @@ def _uses_input(source: str) -> bool:
 
 
 def _rewrite_cell_magic(code: str) -> str:
-    """Process a code block that uses cell magic.:w
+    """Process a code block that uses cell magic.
 
     - Lines starting with "%%capture" are ignored.
     - Lines starting with "%pip" are rewritten by removing the leading "%" character.
@@ -52,10 +51,14 @@ def _rewrite_cell_magic(code: str) -> str:
         if stripped.startswith("%%capture"):
             continue
         # Rewrite %pip lines by dropping the '%'
-        elif stripped.startswith("%pip"):
-            # Drop the leading '%' character
-            rewritten_lines.append(stripped[1:])
-        # Anything else is not supported
+        elif stripped.startswith("%") or stripped.startswith("!"):
+            # Drop the leading '%' character and then drop all leading whitespace
+            stripped = stripped.lstrip("%! \t")
+            # Check if the line starts with "pip"
+            if stripped.startswith("pip"):
+                rewritten_lines.append(stripped)
+            else:
+                raise NotImplementedError(f"Unhandled line: {line}")
         else:
             raise NotImplementedError(f"Unhandled line: {line}")
 
@@ -247,13 +250,10 @@ class EscapePreprocessor(Preprocessor):
             )
             cell.metadata["exec"] = is_exec
 
-            if self.markdown_exec_migration:
-                # For markdown exec migration we'll re-write cell magic as bash commands
-                if source.startswith("%%"):
-                    cell.source = _rewrite_cell_magic(source)
-                    cell.metadata["language"] = "shell"
-
-                cell.metadata["has_output"] = _has_output(source)
+            # For markdown exec migration we'll re-write cell magic as bash commands
+            if source.startswith("%%"):
+                cell.source = _rewrite_cell_magic(source)
+                cell.metadata["language"] = "shell"
 
             # Remove noqa comments
             cell.source = re.sub(r"#\s*noqa.*$", "", cell.source, flags=re.MULTILINE)
@@ -352,7 +352,7 @@ exporter = MarkdownExporter(
 
 
 def convert_notebook(
-    notebook_path: Path,
+    notebook_path: str,
     mode: Literal["markdown", "exec"] = "markdown",
 ) -> str:
     with open(notebook_path) as f:
diff --git a/docs/_scripts/notebook_convert_templates/mdoutput/index.md.j2 b/docs/_scripts/notebook_convert_templates/mdoutput/index.md.j2
index f74d622d0..05cbef99a 100644
--- a/docs/_scripts/notebook_convert_templates/mdoutput/index.md.j2
+++ b/docs/_scripts/notebook_convert_templates/mdoutput/index.md.j2
@@ -1,5 +1,18 @@
 {% extends 'markdown/index.md.j2' %}
 
+{% block input %}{# cell.metadata.language is an addition of our docs pipeline. #}
+```{%- if 'language' in cell.metadata -%}
+{{ cell.metadata.language }}
+{%- elif 'magics_language' in cell.metadata -%}
+{{ cell.metadata.magics_language }}
+{%- elif 'name' in nb.metadata.get('language_info', {}) -%}
+{{ nb.metadata.language_info.name }}
+{%- endif %}
+{{ cell.source }}
+```
+{% endblock input %}
+
+
 {%- block traceback_line -%}
 ```output
 {{ line.rstrip() | strip_ansi }}
diff --git a/docs/tests/unit_tests/test_notebook_conversion.py b/docs/tests/unit_tests/test_notebook_conversion.py
index 46f8ef60c..23233ddf7 100644
--- a/docs/tests/unit_tests/test_notebook_conversion.py
+++ b/docs/tests/unit_tests/test_notebook_conversion.py
@@ -1,8 +1,13 @@
+import os
+import tempfile
+
+import nbformat
 import pytest
 
 from _scripts.notebook_convert import (
     _convert_links_in_markdown,
     _has_output,
+    convert_notebook,
 )
 
 
@@ -32,3 +37,41 @@ def test_has_output() -> None:
 def test_link_conversion(source: str, expected: str) -> None:
     """Test logic to convert links in markdown cells."""
     assert _convert_links_in_markdown(source) == expected
+
+
+EXPECTED_OUTPUT = """\
+```shell
+pip install -U langgraph
+```
+
+
+```python
+print('Hello')
+```\
+
+"""
+
+
+def test_converting_cell_magic() -> None:
+    """Test converting cell magic to code blocks."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        nb_path = os.path.join(tmpdir, "test_notebook.ipynb")
+
+        # Create a minimal notebook object
+        nb = nbformat.v4.new_notebook()
+        nb.cells = [
+            nbformat.v4.new_code_cell(
+                "%%capture --no-stderr\n"
+                "%pip install -U langgraph"
+            ),
+            nbformat.v4.new_code_cell("print('Hello')"),
+        ]
+        nb.metadata["language_info"] = {"name": "python"}
+
+        # Write to file
+        with open(nb_path, "w", encoding="utf-8") as f:
+            nbformat.write(nb, f)
+
+        # Run the conversion
+        converted = convert_notebook(nb_path)
+        assert converted == EXPECTED_OUTPUT