x

docs/_scripts/notebook_hooks.py

@@ -1,9 +1,11 @@
 import logging
 import os
+import re
 from typing import Any, Dict
 
-from mkdocs.structure.pages import Page
 from mkdocs.structure.files import Files, File
+from mkdocs.structure.pages import Page
+
 from notebook_convert import convert_notebook
 
 logger = logging.getLogger(__name__)
@@ -35,12 +37,69 @@ def on_files(files: Files, **kwargs: Dict[str, Any]):
     return new_files
 
 
+def _highlight_code_blocks(markdown: str) -> str:
+    """Find code blocks with highlight comments and add hl_lines attribute.
+
+    Args:
+        markdown: The markdown content to process.
+
+    Returns:
+        updated markdown code with code blocks containing highlight comments
+        updated to use the hl_lines attribute.
+    """
+    # Pattern to find code blocks with highlight comments and without
+    # existing hl_lines for Python and JavaScript
+    code_block_pattern = re.compile(
+        r"```(?P<language>py|python|js)(?!\s+hl_lines=)\n"
+        r"(?P<code>((?:.*\n)*?))"  # Capture the code inside the block using named group
+        r"```"
+    )
+
+    def replace_highlight_comments(match: re.Match) -> str:
+        language = match.group("language")
+        code_block = match.group("code")
+        lines = code_block.split("\n")
+        highlighted_lines = []
+
+        # Skip initial empty lines
+        while lines and not lines[0].strip():
+            lines.pop(0)
+
+        lines_to_keep = []
+
+        comment_syntax = (
+            "# highlight-next-line"
+            if language in ["py", "python"]
+            else "// highlight-next-line"
+        )
+
+        for line in lines:
+            if comment_syntax in line:
+                count = len(lines_to_keep) + 1
+                highlighted_lines.append(str(count))
+            else:
+                lines_to_keep.append(line)
+
+        # Reconstruct the new code block
+        new_code_block = "\n".join(lines_to_keep)
+
+        if highlighted_lines:
+            return f'```{language} hl_lines="{" ".join(highlighted_lines)}"\n{new_code_block}\n```'
+        else:
+            return f"```{language}\n{new_code_block}\n```"
+
+    # Replace all code blocks in the markdown
+    markdown = code_block_pattern.sub(replace_highlight_comments, markdown)
+    return markdown
+
+
 def on_page_markdown(markdown: str, page: Page, **kwargs: Dict[str, Any]):
     if DISABLED:
         return markdown
     if page.file.src_path.endswith(".ipynb"):
         logger.info("Processing Jupyter notebook: %s", page.file.src_path)
-        body = convert_notebook(page.file.abs_src_path)
-        return body
+        markdown = convert_notebook(page.file.abs_src_path)
+    # Apply highlight comments to code blocks
+    markdown = _highlight_code_blocks(markdown)
 
     return markdown