docs: add unit tests to build pipeline (#3427)

* Add testing step to to docs build pipeline
* Requires updating import structure in some place
* Add simple unit test to cover some logic with highlights

.github/workflows/deploy_docs.yml

@@ -92,6 +92,9 @@ jobs:
           poetry run tslab install --python=python3
           poetry run jupyter kernelspec list
 
+      - name: Run unit tests
+        # Run unit tests on the docs build pipeline
+        run: make tests
       - name: Lint Docs
         # This step lints the docs using the existing linting set up.
         # It should be very fast and should not require any external services.

docs/Makefile

@@ -1,4 +1,4 @@
-.PHONY: lint-docs format-docs build-docs serve-docs serve-clean-docs clean-docs codespell build-typedoc llms-text build-prebuilt
+.PHONY: lint-docs format-docs build-docs serve-docs serve-clean-docs clean-docs codespell build-typedoc llms-text build-prebuilt tests
 
 build-typedoc:
 	cd ../libs/sdk-js && yarn install --include-dev && yarn typedoc
@@ -17,7 +17,7 @@ build-docs: build-typedoc build-prebuilt
 	poetry run python -m mkdocs build --clean -f mkdocs.yml --strict
 
 llms-text:
-	poetry run python _scripts/generate_llms_text.py docs/llms-full.txt
+	poetry run python -m _scripts.generate_llms_text docs/llms-full.txt
 
 install-vercel-deps:
 	dnf install -y python3.11
@@ -33,6 +33,11 @@ install-vercel-deps:
 	poetry run jupyter kernelspec list
 
 
+tests:
+	# RUn unit tests
+	poetry run pytest tests/unit_tests
+
+
 vercel-build-docs: install-vercel-deps
 	make build-docs
 

docs/_scripts/generate_llms_text.py

@@ -2,12 +2,11 @@
 
 import glob
 import os
-import pathlib
 
 from mkdocs.structure.files import File
 from mkdocs.structure.pages import Page
 
-from notebook_hooks import _on_page_markdown_with_config
+from _scripts.notebook_hooks import _on_page_markdown_with_config
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 # Get source directory (parent of HERE / docs)

docs/_scripts/notebook_hooks.py

@@ -1,21 +1,19 @@
 import logging
 import os
+import posixpath
 import re
 import traceback
 from typing import Any, Callable, Dict
 
 from markdown import Markdown
-from pymdownx.superfences import SuperFencesException
+from markdown_exec.hooks import SessionHistoryEntry
 from mkdocs.structure.files import Files, File
 from mkdocs.structure.pages import Page
-import posixpath
-
-from markdown_exec.hooks import SessionHistoryEntry
-
-from generate_api_reference_links import update_markdown_with_imports
-from notebook_convert import convert_notebook
-from setup_vcr import load_postamble, load_preamble, _hash_string
+from pymdownx.superfences import SuperFencesException
 
+from _scripts.generate_api_reference_links import update_markdown_with_imports
+from _scripts.notebook_convert import convert_notebook
+from _scripts.setup_vcr import load_postamble, load_preamble, _hash_string
 
 logger = logging.getLogger(__name__)
 logging.basicConfig()
@@ -101,7 +99,7 @@ def _highlight_code_blocks(markdown: str) -> str:
     # existing hl_lines for Python and JavaScript
     # Pattern to find code blocks with highlight comments, handling optional indentation
     code_block_pattern = re.compile(
-        r"(?P<indent>[ \t]*)```(?P<language>py|python|js|javascript)(?!\s+hl_lines=)\n"
+        r"(?P<indent>[ \t]*)```(?P<language>\w+)[ ]*(?P<attributes>[^\n]*)\n"
         r"(?P<code>((?:.*\n)*?))"  # Capture the code inside the block using named group
         r"(?P=indent)```"  # Match closing backticks with the same indentation
     )
@@ -110,6 +108,13 @@ def _highlight_code_blocks(markdown: str) -> str:
         indent = match.group("indent")
         language = match.group("language")
         code_block = match.group("code")
+        attributes = match.group("attributes").rstrip()
+
+        # Account for a case where hl_lines is manually specified
+        if "hl_lines" in attributes:
+            # Return original code block
+            return match.group(0)
+
         lines = code_block.split("\n")
         highlighted_lines = []
 
@@ -135,20 +140,23 @@ def _highlight_code_blocks(markdown: str) -> str:
         # Reconstruct the new code block
         new_code_block = "\n".join(lines_to_keep)
 
+        # Construct the full code block that also includes
+        # the fenced code block syntax.
+        opening_fence = f"```{language}"
+
+        if attributes:
+            opening_fence += f" {attributes}"
+
         if highlighted_lines:
-            return (
-                f'{indent}```{language} hl_lines="{" ".join(highlighted_lines)}"\n'
-                # The indent and terminating \n is already included in the code block
-                f"{new_code_block}"
-                f"{indent}```"
-            )
-        else:
-            return (
-                f"{indent}```{language}\n"
-                # The indent and terminating \n is already included in the code block
-                f"{new_code_block}"
-                f"{indent}```"
-            )
+            opening_fence += f" hl_lines=\"{' '.join(highlighted_lines)}\""
+
+        return (
+            # The indent and opening fence
+            f"{indent}{opening_fence}\n"
+            # The indent and terminating \n is already included in the code block
+            f"{new_code_block}"
+            f"{indent}```"
+        )
 
     # Replace all code blocks in the markdown
     markdown = code_block_pattern.sub(replace_highlight_comments, markdown)
@@ -212,10 +220,21 @@ def handle_vcr_setup(
             wrapped_lines.append(load_postamble(language))
 
         transformed_source = "\n".join(wrapped_lines)
+
+        # Propagate extras
+        keep_extras = {
+            key: value
+            for key, value in kwargs["extra"].items()
+            if key
+            in {
+                "hl_lines",
+            }
+        }
+
         return dict(
             transform_source=lambda code: (transformed_source, code),
             id=id,
-            extra={},
+            extra=keep_extras,
         )
     except Exception as e:
         raise SuperFencesException(traceback.format_exc()) from e

docs/docs/how-tos/state-reducers.md

@@ -11,7 +11,6 @@ We will use [messages](../concepts/low_level.md/#messagesstate) in our examples.
 
 First, let's install langgraph:
 
-
 ```python
 %%capture --no-stderr
 %pip install -U langgraph

docs/tests/unit_tests/test_highlights.py

@@ -0,0 +1,130 @@
+from mkdocs.config.defaults import MkDocsConfig
+from mkdocs.structure.files import File
+from mkdocs.structure.pages import Page
+
+from _scripts.notebook_hooks import _highlight_code_blocks, on_page_markdown
+
+NO_OP_INPUT_1 = """\
+This is a plain text without any code blocks.
+
+```python
+print("Hello, World!")
+```
+"""
+
+NO_OP_INPUT_2 = """\
+
+=== "Python"
+
+    ```python
+        def foo():
+            pass
+    print("Hello, World!")
+    ```
+"""
+
+
+def test_highlight_code_blocks_no_op() -> None:
+    assert _highlight_code_blocks(NO_OP_INPUT_1) == NO_OP_INPUT_1
+    assert _highlight_code_blocks(NO_OP_INPUT_2) == NO_OP_INPUT_2
+
+
+# Examples are written in multiline style to make sure that whitespace
+# is easy to interpret.
+INPUT_HIGHLIGHT_1 = """\
+This is a plain text without any code blocks.
+
+```python
+# highlight-next-line
+print("Hello, World!")
+```
+"""
+
+EXPECTED_HIGHLIGHT_1 = """\
+This is a plain text without any code blocks.
+
+```python hl_lines="1"
+print("Hello, World!")
+```
+"""
+
+INPUT_HIGHLIGHT_2 = """\
+This is a plain text without any code blocks.
+
+```python
+# highlight-next-line
+print("Hello, World!")
+
+x = 5
+
+# highlight-next-line
+print("Hello, World!")
+
+```
+"""
+
+EXPECTED_HIGHLIGHT_2 = """\
+This is a plain text without any code blocks.
+
+```python hl_lines="1 5"
+print("Hello, World!")
+
+x = 5
+
+print("Hello, World!")
+
+```
+"""
+
+
+# Test end-to-end behavior of on_page_markdown
+INPUT_HIGHLIGHT_3 = """\
+```python exec="on" source="below"
+print("Hello, World!")
+# highlight-next-line
+print("Hello, World!")
+```
+"""
+
+EXPECTED_HIGHLIGHT_3 = """\
+```python exec="on" source="below" hl_lines="2"
+print("Hello, World!")
+print("Hello, World!")
+```
+"""
+
+
+def test_highlight_code_blocks() -> None:
+    """Test that code blocks are highlighted correctly."""
+    assert _highlight_code_blocks(INPUT_HIGHLIGHT_1) == EXPECTED_HIGHLIGHT_1
+    assert _highlight_code_blocks(INPUT_HIGHLIGHT_2) == EXPECTED_HIGHLIGHT_2
+    assert _highlight_code_blocks(INPUT_HIGHLIGHT_3) == EXPECTED_HIGHLIGHT_3
+
+
+END_TO_END_INPUT_HIGHLIGHT_1 = """\
+```python exec="on" source="below"
+print("Hello, World!")
+# highlight-next-line
+print("Hello, World!")
+```
+"""
+
+
+END_TO_END_INPUT_HIGHLIGHT_1_EXPECT = """\
+```python exec="on" source="below" hl_lines="2" path="dummy.md"
+print("Hello, World!")
+print("Hello, World!")
+```
+"""
+
+
+def test_on_page_markdown_highlights() -> None:
+    """Test that on page markdown behaves correctly."""
+    # Create a dummy MkDocs File and Page object.
+    dummy_file = File("dummy.md", "dummy.md", "placeholder", use_directory_urls=False)
+    dummy_page = Page("Test Page", dummy_file, config=MkDocsConfig())
+
+    assert (
+        on_page_markdown(END_TO_END_INPUT_HIGHLIGHT_1, dummy_page)
+        == END_TO_END_INPUT_HIGHLIGHT_1_EXPECT
+    )