diff --git a/build_scripts/gen_api_md.py b/build_scripts/gen_api_md.py index b811db2707..8a2509458e 100644 --- a/build_scripts/gen_api_md.py +++ b/build_scripts/gen_api_md.py @@ -16,12 +16,16 @@ python build_scripts/gen_api_md.py """ +import ast import json import re import sys +import textwrap from dataclasses import dataclass from pathlib import Path +import yaml + # Import sibling script for post-generation TOC validation. sys.path.insert(0, str(Path(__file__).parent)) import validate_docs @@ -46,6 +50,31 @@ class SymbolEntry: anchor: str # MyST label, e.g. "api-pyrit_prompt_target-PromptTarget" +@dataclass(frozen=True) +class ExampleReference: + """A user-guide page that demonstrates an API symbol.""" + + title: str + path: str + + +@dataclass(frozen=True) +class _ImportBinding: + """A local import name bound to either a symbol or a PyRIT namespace.""" + + symbol: SymbolEntry | None = None + namespace: str | None = None + + +_EXAMPLE_DOC_EXCLUDED_PREFIXES = ( + "_api/", + "api/", + "blog/", + "contributing/", + "generate_docs/", +) + + # Backtick code spans that look like Python identifiers (with optional # dotted paths) — candidates for symbol cross-reference rewriting. Matches # either `name` or ``name``. The leading negative lookbehind prevents @@ -225,6 +254,268 @@ def _add(key: str, entry: SymbolEntry) -> None: return index +def _unique_public_symbol(entries: list[SymbolEntry]) -> SymbolEntry | None: + """Return one public class/function target when all matches share an anchor.""" + by_anchor = {entry.anchor: entry for entry in entries if entry.kind in ("class", "function")} + return next(iter(by_anchor.values())) if len(by_anchor) == 1 else None + + +def _resolve_imported_symbol( + qualified_name: str, + *, + symbol_index: dict[str, list[SymbolEntry]], +) -> SymbolEntry | None: + """Resolve an imported API name, with a conservative short-name fallback.""" + exact = _unique_public_symbol(symbol_index.get(qualified_name, [])) + if exact: + return exact + short_entries = symbol_index.get(qualified_name.rsplit(".", 1)[-1], []) + module_matches = [entry for entry in short_entries if qualified_name.startswith(f"{entry.module}.")] + return _unique_public_symbol(module_matches) or _unique_public_symbol(short_entries) + + +def _is_pyrit_module(name: str) -> bool: + """Return whether a dotted import path belongs to the PyRIT package.""" + return name == "pyrit" or name.startswith("pyrit.") + + +def _collect_import_bindings( + trees: list[ast.AST], + *, + symbol_index: dict[str, list[SymbolEntry]], +) -> dict[str, _ImportBinding]: + """Collect local names introduced by PyRIT imports.""" + bindings: dict[str, _ImportBinding] = {} + for tree in trees: + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module and _is_pyrit_module(node.module): + _add_from_import_bindings(node, bindings=bindings, symbol_index=symbol_index) + elif isinstance(node, ast.Import): + _add_import_bindings(node, bindings=bindings) + return bindings + + +def _add_from_import_bindings( + node: ast.ImportFrom, + *, + bindings: dict[str, _ImportBinding], + symbol_index: dict[str, list[SymbolEntry]], +) -> None: + """Add bindings from one ``from pyrit... import ...`` statement.""" + if not node.module: + return + for alias in node.names: + if alias.name == "*": + continue + local_name = alias.asname or alias.name + qualified_name = f"{node.module}.{alias.name}" + symbol = _resolve_imported_symbol(qualified_name, symbol_index=symbol_index) + bindings[local_name] = _ImportBinding(symbol=symbol, namespace=None if symbol else qualified_name) + + +def _add_import_bindings( + node: ast.Import, + *, + bindings: dict[str, _ImportBinding], +) -> None: + """Add namespace bindings from one ``import pyrit...`` statement.""" + for alias in node.names: + if not _is_pyrit_module(alias.name): + continue + local_name = alias.asname or alias.name.split(".", 1)[0] + namespace = alias.name if alias.asname else local_name + bindings[local_name] = _ImportBinding(namespace=namespace) + + +def _attribute_parts(node: ast.Attribute) -> list[str] | None: + """Return the dotted-name parts represented by an attribute expression.""" + parts = [node.attr] + value = node.value + while isinstance(value, ast.Attribute): + parts.append(value.attr) + value = value.value + if not isinstance(value, ast.Name): + return None + parts.append(value.id) + return list(reversed(parts)) + + +def _collect_used_anchors( + trees: list[ast.AST], + *, + bindings: dict[str, _ImportBinding], + symbol_index: dict[str, list[SymbolEntry]], +) -> set[str]: + """Find imported API symbols that are actually referenced by the page.""" + anchors: set[str] = set() + for tree in trees: + for node in ast.walk(tree): + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load): + binding = bindings.get(node.id) + if binding and binding.symbol: + anchors.add(binding.symbol.anchor) + elif isinstance(node, ast.Attribute) and isinstance(node.ctx, ast.Load): + symbol = _resolve_attribute_symbol(node, bindings=bindings, symbol_index=symbol_index) + if symbol: + anchors.add(symbol.anchor) + return anchors + + +def _resolve_attribute_symbol( + node: ast.Attribute, + *, + bindings: dict[str, _ImportBinding], + symbol_index: dict[str, list[SymbolEntry]], +) -> SymbolEntry | None: + """Resolve an attribute rooted in an imported PyRIT namespace.""" + parts = _attribute_parts(node) + if not parts: + return None + binding = bindings.get(parts[0]) + if not binding or not binding.namespace: + return None + for end in range(1, len(parts)): + qualified_name = ".".join([binding.namespace, *parts[1 : end + 1]]) + exact = _unique_public_symbol(symbol_index.get(qualified_name, [])) + if exact: + return exact + segment = parts[end] + if segment[:1].isupper(): + class_match = _unique_public_symbol(symbol_index.get(segment, [])) + if class_match and class_match.kind == "class": + return class_match + qualified_name = ".".join([binding.namespace, *parts[1:]]) + return _resolve_imported_symbol(qualified_name, symbol_index=symbol_index) + + +def _parse_python_chunks(chunks: list[str]) -> list[ast.AST]: + """Parse independent code chunks, ignoring notebook magics or invalid snippets.""" + trees: list[ast.AST] = [] + for chunk in chunks: + try: + trees.append(ast.parse(chunk)) + except SyntaxError: + continue + return trees + + +def _jupytext_code_chunks(text: str) -> list[str]: + """Extract Python cells from a percent-format Jupytext file.""" + chunks: list[str] = [] + current: list[str] = [] + is_markdown = False + for line in text.splitlines(): + if line.startswith("# %%"): + if current and not is_markdown: + chunks.append("\n".join(current)) + current = [] + is_markdown = "[markdown]" in line + elif not is_markdown: + current.append(line) + if current and not is_markdown: + chunks.append("\n".join(current)) + return chunks + + +def _markdown_code_chunks(text: str) -> list[str]: + """Extract fenced Python examples from a Markdown page.""" + pattern = re.compile(r"^[ \t]{0,3}```(?:python|py)\s*\n(.*?)^[ \t]{0,3}```\s*$", flags=re.MULTILINE | re.DOTALL) + return [textwrap.dedent(match.group(1)) for match in pattern.finditer(text)] + + +def _notebook_content(path: Path) -> tuple[list[str], str | None]: + """Return code-cell sources and the first H1 from a notebook.""" + try: + notebook = json.loads(path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + return [], None + cells = notebook.get("cells", []) + chunks = ["".join(cell.get("source", [])) for cell in cells if cell.get("cell_type") == "code"] + markdown = "\n".join("".join(cell.get("source", [])) for cell in cells if cell.get("cell_type") == "markdown") + return chunks, _first_markdown_h1(markdown) + + +def _first_markdown_h1(text: str) -> str | None: + """Return the first level-one Markdown heading.""" + match = re.search(r"^#\s+(.+?)\s*$", text, flags=re.MULTILINE) + return match.group(1).strip() if match else None + + +def _markdown_title(text: str) -> str | None: + """Return a frontmatter title when present, otherwise the first H1.""" + frontmatter = re.match(r"^---\s*\n(.*?)^---\s*$", text, flags=re.MULTILINE | re.DOTALL) + if frontmatter: + match = re.search(r"^title:\s*(.+?)\s*$", frontmatter.group(1), flags=re.MULTILINE) + if match: + return match.group(1).strip().strip("\"'") + return _first_markdown_h1(text) + + +def _jupytext_h1(text: str) -> str | None: + """Return the first level-one heading encoded in Jupytext comments.""" + match = re.search(r"^#\s+#\s+(.+?)\s*$", text, flags=re.MULTILINE) + return match.group(1).strip() if match else None + + +def _humanize_stem(path: Path) -> str: + """Create a readable fallback title from a documentation filename.""" + stem = re.sub(r"^\d+(?:_\d+)*_", "", path.stem) + return stem.replace("_", " ").replace("-", " ").title() + + +def _read_example_page(page_path: Path) -> tuple[list[str], str]: + """Read one TOC page and return its Python chunks and display title.""" + if page_path.suffix == ".ipynb" and page_path.with_suffix(".py").exists(): + source_path = page_path.with_suffix(".py") + text = source_path.read_text(encoding="utf-8") + return _jupytext_code_chunks(text), _jupytext_h1(text) or _humanize_stem(page_path) + if page_path.suffix == ".ipynb": + chunks, title = _notebook_content(page_path) + return chunks, title or _humanize_stem(page_path) + text = page_path.read_text(encoding="utf-8") + if page_path.suffix == ".py": + return _jupytext_code_chunks(text), _jupytext_h1(text) or _humanize_stem(page_path) + return _markdown_code_chunks(text), _markdown_title(text) or _humanize_stem(page_path) + + +def _example_toc_paths(*, doc_root: Path, toc_path: Path) -> list[Path]: + """Return existing user-guide pages referenced by the MyST TOC.""" + config = yaml.safe_load(toc_path.read_text(encoding="utf-8")) + entries = config.get("project", {}).get("toc", []) + references = validate_docs.parse_toc_files(entries) + paths = [] + for reference in sorted(references): + normalized = Path(reference).as_posix() + if normalized.startswith(_EXAMPLE_DOC_EXCLUDED_PREFIXES): + continue + path = doc_root / reference + if path.is_file() and path.suffix in (".md", ".ipynb", ".py"): + paths.append(path) + return paths + + +def _build_example_index( + *, + doc_root: Path, + toc_path: Path, + symbol_index: dict[str, list[SymbolEntry]], +) -> dict[str, list[ExampleReference]]: + """Map API anchors to TOC pages that import and use those symbols.""" + examples: dict[str, set[ExampleReference]] = {} + for page_path in _example_toc_paths(doc_root=doc_root, toc_path=toc_path): + chunks, title = _read_example_page(page_path) + trees = _parse_python_chunks(chunks) + bindings = _collect_import_bindings(trees, symbol_index=symbol_index) + anchors = _collect_used_anchors(trees, bindings=bindings, symbol_index=symbol_index) + reference = ExampleReference(title=title, path=page_path.relative_to(doc_root).as_posix()) + for anchor in anchors: + examples.setdefault(anchor, set()).add(reference) + return { + anchor: sorted(references, key=lambda reference: (reference.title.casefold(), reference.path)) + for anchor, references in examples.items() + } + + def _resolve_symbol(raw: str, index: dict[str, list[SymbolEntry]], current_class: str | None) -> SymbolEntry | None: """Return the cross-reference target for a bare backtick-quoted symbol. @@ -382,6 +673,17 @@ def _process_docstring_text( return _rewrite_symbol_refs(escaped, symbol_index, current_class=current_class) +def render_examples(examples: list[ExampleReference] | None) -> str: + """Render links to user-guide pages that exercise an API symbol.""" + if not examples: + return "" + parts = ["**Examples:**\n"] + for example in examples: + title = example.title.replace("[", r"\[").replace("]", r"\]") + parts.append(f"- [{title}](../{example.path})") + return "\n".join(parts) + + def render_function( func: dict, *, @@ -389,6 +691,7 @@ def render_function( module: str, class_name: str | None = None, symbol_index: dict[str, list[SymbolEntry]] | None = None, + examples_by_anchor: dict[str, list[ExampleReference]] | None = None, ) -> str: """Render a function as markdown.""" name = func["name"] @@ -431,6 +734,10 @@ def render_function( if raises_md: parts.append(raises_md + "\n") + examples_md = render_examples(examples_by_anchor.get(anchor) if examples_by_anchor else None) + if examples_md: + parts.append(examples_md + "\n") + return "\n".join(parts) @@ -439,6 +746,7 @@ def render_class( *, module: str, symbol_index: dict[str, list[SymbolEntry]] | None = None, + examples_by_anchor: dict[str, list[ExampleReference]] | None = None, ) -> str: """Render a class as markdown.""" name = cls["name"] @@ -466,6 +774,10 @@ def render_class( parts.append("**Constructor Parameters:**\n") parts.append(render_params(init_params) + "\n") + examples_md = render_examples(examples_by_anchor.get(anchor) if examples_by_anchor else None) + if examples_md: + parts.append(examples_md + "\n") + # Methods methods = cls.get("methods", []) if methods: @@ -477,6 +789,7 @@ def render_class( module=module, class_name=name, symbol_index=symbol_index, + examples_by_anchor=examples_by_anchor, ) for m in methods ) @@ -498,6 +811,7 @@ def render_module( data: dict, *, symbol_index: dict[str, list[SymbolEntry]] | None = None, + examples_by_anchor: dict[str, list[ExampleReference]] | None = None, ) -> str: """Render a full module page.""" mod_name = data["name"] @@ -524,9 +838,25 @@ def render_module( if functions: parts.append("## Functions\n") - parts.extend(render_function(f, module=mod_name, symbol_index=symbol_index) for f in functions) + parts.extend( + render_function( + f, + module=mod_name, + symbol_index=symbol_index, + examples_by_anchor=examples_by_anchor, + ) + for f in functions + ) - parts.extend(render_class(cls, module=mod_name, symbol_index=symbol_index) for cls in classes) + parts.extend( + render_class( + cls, + module=mod_name, + symbol_index=symbol_index, + examples_by_anchor=examples_by_anchor, + ) + for cls in classes + ) if aliases: parts.append("## Re-exports\n") @@ -683,13 +1013,18 @@ def main() -> None: # Build a symbol index over the post-resolution module tree so the # docstring rewriter can turn backticked names into MyST cross-references. symbol_index = _build_symbol_index(modules) + example_index = _build_example_index( + doc_root=Path("doc"), + toc_path=Path("doc/myst.yml"), + symbol_index=symbol_index, + ) # Generate per-module pages for data in modules: mod_name = data["name"] slug = mod_name.replace(".", "_") md_path = API_MD_DIR / f"{slug}.md" - content = render_module(data, symbol_index=symbol_index) + content = render_module(data, symbol_index=symbol_index, examples_by_anchor=example_index) members = data.get("members", []) rendered_count = sum(1 for m in members if m.get("kind") in ("class", "function")) md_path.write_text(content, encoding="utf-8") diff --git a/tests/unit/build_scripts/test_gen_api_md.py b/tests/unit/build_scripts/test_gen_api_md.py index db4472d300..5fc7df438f 100644 --- a/tests/unit/build_scripts/test_gen_api_md.py +++ b/tests/unit/build_scripts/test_gen_api_md.py @@ -1,8 +1,12 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from pathlib import Path + from build_scripts.gen_api_md import ( + ExampleReference, SymbolEntry, + _build_example_index, _build_symbol_index, _class_anchor, _format_bases, @@ -541,3 +545,212 @@ def test_render_module_label_uses_module_slug_for_nested_packages() -> None: assert "label: api-pyrit_executor_attack" in out assert "# pyrit.executor.attack" in out + + +def _write_doc(doc_root: Path, relative_path: str, content: str) -> None: + path = doc_root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + +def _write_toc(doc_root: Path, *files: str) -> Path: + toc_path = doc_root / "myst.yml" + entries = "\n".join(f" - file: {file}" for file in files) + toc_path.write_text(f"project:\n toc:\n{entries}\n", encoding="utf-8") + return toc_path + + +def test_build_example_index_scans_toc_user_docs_and_used_imports(tmp_path: Path) -> None: + symbol_index = _build_symbol_index( + [ + _fake_module( + "pyrit.targets", + [ + _fake_class("DirectTarget"), + _fake_class("AliasedTarget"), + _fake_class("UnusedTarget"), + _fake_class("LookalikeTarget"), + _fake_function("used_function"), + _fake_function("unused_function"), + _fake_class("InternalTarget"), + ], + ) + ] + ) + toc_path = _write_toc( + tmp_path, + "guide/targets.md", + "guide/invalid.md", + "contributing/development.md", + ) + _write_doc( + tmp_path, + "guide/targets.md", + """--- +title: Target examples +--- + + ```python + from pyrit.targets import AliasedTarget as Target + from pyrit.targets import DirectTarget, UnusedTarget + from pyrit.targets import unused_function, used_function + from pyrite.targets import LookalikeTarget + + DirectTarget() + Target() + used_function() + LookalikeTarget() + ``` +""", + ) + _write_doc(tmp_path, "guide/invalid.md", "# Invalid\n\n```python\nthis is not valid Python !!!\n```\n") + _write_doc( + tmp_path, + "contributing/development.md", + "# Internal development\n\n```python\nfrom pyrit.targets import InternalTarget\nInternalTarget()\n```\n", + ) + _write_doc( + tmp_path, + "guide/not-in-toc.md", + "# Orphan\n\n```python\nfrom pyrit.targets import UnusedTarget\nUnusedTarget()\n```\n", + ) + + result = _build_example_index(doc_root=tmp_path, toc_path=toc_path, symbol_index=symbol_index) + expected = [ExampleReference(title="Target examples", path="guide/targets.md")] + + assert result[_class_anchor("pyrit.targets", "DirectTarget")] == expected + assert result[_class_anchor("pyrit.targets", "AliasedTarget")] == expected + assert _class_anchor("pyrit.targets", "UnusedTarget") not in result + assert _class_anchor("pyrit.targets", "LookalikeTarget") not in result + assert result[_function_anchor("pyrit.targets", "used_function")] == expected + assert _function_anchor("pyrit.targets", "unused_function") not in result + assert _class_anchor("pyrit.targets", "InternalTarget") not in result + + +def test_build_example_index_uses_jupytext_companions_dedupes_and_sorts(tmp_path: Path) -> None: + symbol_index = _build_symbol_index([_fake_module("pyrit.targets", [_fake_class("PromptTarget")])]) + toc_path = _write_toc(tmp_path, "guide/zulu.ipynb", "guide/alpha.ipynb") + for stem in ("zulu", "alpha"): + _write_doc(tmp_path, f"guide/{stem}.ipynb", "{}") + _write_doc( + tmp_path, + "guide/zulu.py", + """# %% [markdown] +# # Zulu guide + +# %% +from pyrit.targets import PromptTarget as Target + +# %% +Target() +Target() +""", + ) + _write_doc( + tmp_path, + "guide/alpha.py", + """# %% [markdown] +# # Alpha guide + +# %% +from pyrit.targets import PromptTarget + +# %% +PromptTarget() +""", + ) + + result = _build_example_index(doc_root=tmp_path, toc_path=toc_path, symbol_index=symbol_index) + + assert result[_class_anchor("pyrit.targets", "PromptTarget")] == [ + ExampleReference(title="Alpha guide", path="guide/alpha.ipynb"), + ExampleReference(title="Zulu guide", path="guide/zulu.ipynb"), + ] + + +def test_build_example_index_resolves_module_alias_and_unique_short_name(tmp_path: Path) -> None: + symbol_index = _build_symbol_index( + [ + _fake_module( + "pyrit.targets", + [_fake_class("ModuleTarget"), _fake_class("DeepTarget"), _fake_function("execute")], + ), + _fake_module("pyrit.prompt_target", [_fake_class("TargetCapabilities")]), + _fake_module("pyrit.models", [_fake_class("TargetCapabilities")]), + _fake_module("pyrit.first", [_fake_class("AmbiguousTarget")]), + _fake_module("pyrit.second", [_fake_class("AmbiguousTarget")]), + ] + ) + toc_path = _write_toc(tmp_path, "guide/imports.md") + _write_doc( + tmp_path, + "guide/imports.md", + """# Import styles + +```python +import pyrit.targets as targets +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.targets.internal.deep import AmbiguousTarget, DeepTarget + +targets.ModuleTarget.execute() +TargetCapabilities() +DeepTarget() +AmbiguousTarget() +``` +""", + ) + + result = _build_example_index(doc_root=tmp_path, toc_path=toc_path, symbol_index=symbol_index) + expected = [ExampleReference(title="Import styles", path="guide/imports.md")] + + assert result[_class_anchor("pyrit.targets", "ModuleTarget")] == expected + assert result[_class_anchor("pyrit.targets", "DeepTarget")] == expected + assert result[_class_anchor("pyrit.prompt_target", "TargetCapabilities")] == expected + assert _class_anchor("pyrit.models", "TargetCapabilities") not in result + assert _function_anchor("pyrit.targets", "execute") not in result + assert _class_anchor("pyrit.first", "AmbiguousTarget") not in result + assert _class_anchor("pyrit.second", "AmbiguousTarget") not in result + + +def test_build_example_index_reads_standalone_notebook_code_cells(tmp_path: Path) -> None: + symbol_index = _build_symbol_index([_fake_module("pyrit.targets", [_fake_class("NotebookTarget")])]) + toc_path = _write_toc(tmp_path, "guide/notebook.ipynb") + _write_doc( + tmp_path, + "guide/notebook.ipynb", + """{ + "cells": [ + {"cell_type": "markdown", "source": ["# Notebook guide"]}, + {"cell_type": "code", "source": ["from pyrit.targets import NotebookTarget\\n", "NotebookTarget()"]} + ] +}""", + ) + + result = _build_example_index(doc_root=tmp_path, toc_path=toc_path, symbol_index=symbol_index) + + assert result[_class_anchor("pyrit.targets", "NotebookTarget")] == [ + ExampleReference(title="Notebook guide", path="guide/notebook.ipynb") + ] + + +def test_render_module_adds_examples_to_functions_and_classes() -> None: + module = _fake_module( + "pyrit.examples", + [_fake_function("run_example"), _fake_class("ExampleTarget", methods=["run"])], + ) + examples_by_anchor = { + _function_anchor("pyrit.examples", "run_example"): [ + ExampleReference(title="Function guide", path="guide/function.md") + ], + _class_anchor("pyrit.examples", "ExampleTarget"): [ + ExampleReference(title="Class guide", path="guide/class.ipynb") + ], + } + + out = render_module(module, symbol_index={}, examples_by_anchor=examples_by_anchor) + + assert out.count("**Examples:**") == 2 + assert "- [Function guide](../guide/function.md)" in out + assert "- [Class guide](../guide/class.ipynb)" in out + class_section = out.split("## `ExampleTarget`", 1)[1] + assert class_section.index("- [Class guide]") < class_section.index("**Methods:**")