From 6b59ab410ddad1b6eab21b42e34f1051217165a5 Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Mon, 16 Jun 2025 10:53:58 -0400
Subject: [PATCH] add testing

---
 docs/_scripts/generate_llms_text.py           | 16 +++++----
 docs/_scripts/notebook_hooks.py               | 33 +++++++++++++++++++
 .../unit_tests/test_conditional_rendering.py  | 22 +++++++++++++
 3 files changed, 65 insertions(+), 6 deletions(-)
 create mode 100644 docs/tests/unit_tests/test_conditional_rendering.py

diff --git a/docs/_scripts/generate_llms_text.py b/docs/_scripts/generate_llms_text.py
index d71f28589..ff981834e 100644
--- a/docs/_scripts/generate_llms_text.py
+++ b/docs/_scripts/generate_llms_text.py
@@ -3,19 +3,21 @@
 import asyncio
 import glob
 import os
-from typing import TypedDict, List, Optional
-import pydantic
 import re
-from pydantic import BaseModel, Field
-from langchain_core.rate_limiters import InMemoryRateLimiter
+from typing import TypedDict, List, Optional
 
 import yaml
 from langchain.chat_models import init_chat_model
+from langchain_core.rate_limiters import InMemoryRateLimiter
 from mkdocs.structure.files import File
 from mkdocs.structure.pages import Page
+from pydantic import BaseModel, Field
 from yaml import SafeLoader
 
-from _scripts.notebook_hooks import _on_page_markdown_with_config
+from _scripts.notebook_hooks import (
+    _on_page_markdown_with_config,
+    _apply_conditional_rendering,
+)
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 # Get source directory (parent of HERE / docs)
@@ -211,7 +213,9 @@ async def process_nav_items(nav_items: list[NavItem]) -> list[NavItem]:
     # Remove any items that start with http:// or https:// looking only for
     # local file at this stages.
     nav_items = [
-        item for item in nav_items if not item["url"].startswith(("http://", "https://"))
+        item
+        for item in nav_items
+        if not item["url"].startswith(("http://", "https://"))
     ]
     # Process items in parallel
     tasks = [process_single_item(item) for item in nav_items]
diff --git a/docs/_scripts/notebook_hooks.py b/docs/_scripts/notebook_hooks.py
index 30545408e..a50b12e62 100644
--- a/docs/_scripts/notebook_hooks.py
+++ b/docs/_scripts/notebook_hooks.py
@@ -158,6 +158,35 @@ def _add_path_to_code_blocks(markdown: str, page: Page) -> str:
     return code_block_pattern.sub(replace_code_block_header, markdown)
 
 
+def _apply_conditional_rendering(md_text: str, target_language: str) -> str:
+    if target_language not in {"python", "js"}:
+        raise ValueError("target_language must be 'python' or 'js'")
+
+    pattern = re.compile(
+        r"(?P<indent>[ \t]*):::(?P<language>\w+)\s*\n"
+        r"(?P<content>((?:.*\n)*?))"  # Capture the content inside the block
+        r"(?P=indent):::"  # Match closing with the same indentation
+    )
+
+    def replace_conditional_blocks(match: re.Match) -> str:
+        """Keep active conditionals."""
+        language = match.group("language")
+        content = match.group("content")
+
+        if language not in {"python", "js"}:
+            # If the language is not supported, return the original block
+            return match.group(0)
+
+        if language == target_language:
+            return content
+
+        # If the language does not match, return an empty string
+        return ""
+
+    processed = pattern.sub(replace_conditional_blocks, md_text)
+    return processed
+
+
 def _highlight_code_blocks(markdown: str) -> str:
     """Find code blocks with highlight comments and add hl_lines attribute.
 
@@ -257,6 +286,10 @@ def _on_page_markdown_with_config(
     # Apply highlight comments to code blocks
     markdown = _highlight_code_blocks(markdown)
 
+    # Apply conditional rendering for code blocks
+    target_language = kwargs.get("target_language", "js")
+    markdown = _apply_conditional_rendering(markdown, target_language)
+
     # Add file path as an attribute to code blocks that are executable.
     # This file path is used to associate fixtures with the executable code
     # which can be used in CI to test the docs without making network requests.
diff --git a/docs/tests/unit_tests/test_conditional_rendering.py b/docs/tests/unit_tests/test_conditional_rendering.py
new file mode 100644
index 000000000..2bfcedd3b
--- /dev/null
+++ b/docs/tests/unit_tests/test_conditional_rendering.py
@@ -0,0 +1,22 @@
+from _scripts.notebook_hooks import _apply_conditional_rendering
+
+
+CONDITIONAL_RENDERING = """
+above
+:::js
+js-content
+:::
+between
+:::python
+python-content
+:::
+below
+"""
+
+
+def test_conditional_rendering() -> None:
+    """Test logic for conditional rendering of content."""
+    output = _apply_conditional_rendering(CONDITIONAL_RENDERING, "js")
+    assert output.strip() == "above\njs-content\n\nbetween\n\nbelow"
+    output = _apply_conditional_rendering(CONDITIONAL_RENDERING, "python")
+    assert output.strip() == "above\n\nbetween\npython-content\n\nbelow"