scripts/docs/hooks.py
Rama master · réplica de lectura del repositorio original.
import importlib.util
import json
import logging
import mimetypes
import os
import posixpath
import re
import shutil
import sys
from pathlib import Path
from xml.sax.saxutils import escape
import yaml
from mkdocs.structure.files import File
mimetypes.add_type("text/plain", ".md")
log = logging.getLogger("mkdocs")
WELL_KNOWN_SKILLS_DIR = ".well-known/skills"
SKILL_PATH = ("skills", "dstack", "SKILL.md")
DISABLE_LLM_TXT_ENV = "DSTACK_DOCS_DISABLE_LLM_TXT"
DISABLE_YAML_SCHEMAS_ENV = "DSTACK_DOCS_DISABLE_YAML_SCHEMAS"
SCHEMA_REFERENCE_PREFIX = "docs/reference/"
SWAGGER_TAG_ARG = r"(?:\s+tag=(?P<tag_quote>[\"'])(?P<tag>.*?)(?P=tag_quote))?"
SWAGGER_TOKEN = re.compile(rf"!!swagger(?:\s+(?P<path>[^\s<>&:!]+){SWAGGER_TAG_ARG})?!!")
SWAGGER_HTTP_TOKEN = re.compile(
rf"!!swagger-http(?:\s+(?P<path>https?://[^\s!]+){SWAGGER_TAG_ARG})?!!"
)
SWAGGER_USAGE_MSG = (
"Usage: '!!swagger <filename> [tag=\"tag name\"]!!' or "
"'!!swagger-http <url> [tag=\"tag name\"]!!'. "
"File must either exist locally and be placed next to the .md that contains "
"the swagger statement, or be an http(s) URL."
)
HTTP_METHODS = {"get", "put", "post", "delete", "options", "head", "patch", "trace"}
UNTAGGED_OPENAPI_TAG = "default"
def _expand_schema_references(text: str) -> str:
"""Lazy load gen_schema_reference by file path so it works regardless of sys.path."""
hooks_dir = os.path.dirname(os.path.abspath(__file__))
gen_path = os.path.join(hooks_dir, "gen_schema_reference.py")
spec = importlib.util.spec_from_file_location("gen_schema_reference", gen_path)
if spec is None or spec.loader is None:
raise ImportError(f"Cannot load {gen_path}")
module = importlib.util.module_from_spec(spec)
sys.modules["gen_schema_reference"] = module
spec.loader.exec_module(module)
return module.expand_schema_references(text)
def _get_schema_expanded_content(rel_path, config, src_path=None):
"""Return expanded markdown for reference/**/*.md that contain #SCHEMA#, else None.
If src_path is given (e.g. from on_post_build loop), read from it; else build path from config.
"""
if os.environ.get(DISABLE_YAML_SCHEMAS_ENV):
return None
if not rel_path.startswith(SCHEMA_REFERENCE_PREFIX) or not rel_path.endswith(".md"):
log.debug(f"Skipping {rel_path}: not in {SCHEMA_REFERENCE_PREFIX} or not .md")
return None
if src_path is None:
repo_root = os.path.dirname(config["config_file_path"])
docs_dir = config["docs_dir"]
if not os.path.isabs(docs_dir):
docs_dir = os.path.join(repo_root, docs_dir)
src_path = os.path.join(docs_dir, rel_path.replace("/", os.sep))
if not os.path.isfile(src_path):
log.debug(f"Skipping {rel_path}: source file not found at {src_path}")
return None
try:
with open(src_path, "r", encoding="utf-8") as f:
text = f.read()
except OSError as e:
log.debug(f"Skipping {rel_path}: error reading file: {e}")
return None
if "#SCHEMA#" not in text:
log.debug(f"Skipping {rel_path}: no #SCHEMA# placeholders found")
return None
log.debug(f"Expanding schema references in {rel_path}")
return _expand_schema_references(text)
def on_page_read_source(page, config):
"""Use expanded schema content for reference docs when rendering HTML."""
rel_path = page.file.src_uri
content = _get_schema_expanded_content(rel_path, config)
if content is not None:
return content
return None
def on_page_markdown(markdown, page, config, files):
"""Render Swagger UI tokens with the project's preferred defaults."""
while True:
match = SWAGGER_TOKEN.search(markdown)
is_http = False
if match is None:
match = SWAGGER_HTTP_TOKEN.search(markdown)
is_http = True
if match is None:
return markdown
markdown = _replace_swagger_token(markdown, match, is_http, page, files)
def _replace_swagger_token(markdown, match, is_http, page, files):
pre_token = markdown[: match.start()]
post_token = markdown[match.end() :]
path = match.group("path")
tag = match.groupdict().get("tag")
operation_headings = ""
if path is None:
return _swagger_error(pre_token, post_token, SWAGGER_USAGE_MSG)
if is_http:
url = path
else:
try:
api_file = Path(page.file.abs_src_path).parent / path
except ValueError as exc: # pragma: no cover
return _swagger_error(pre_token, post_token, f"Invalid path. {exc.args[0]}")
if not api_file.exists():
return _swagger_error(pre_token, post_token, f"File {path} not found.")
try:
src_uri = api_file.relative_to(page.file.src_dir).as_posix()
except ValueError as exc:
return _swagger_error(
pre_token,
post_token,
f"File {path} must be inside the docs directory. {exc.args[0]}",
)
new_file = File(src_uri, page.file.src_dir, page.file.dest_dir, False)
url = _relative_url(page.file.dest_uri, new_file.dest_uri)
for file in files:
if file.dest_uri != new_file.dest_uri:
continue
if file.abs_src_path == new_file.abs_src_path:
break
return _swagger_error(
pre_token,
post_token,
"Cannot use 2 different swagger files with same filename in same page.",
)
else:
files.append(new_file)
operation_headings = _openapi_operation_headings(api_file, tag)
return pre_token + operation_headings + _swagger_html(url, tag) + post_token
def _relative_url(page_dest_uri: str, asset_dest_uri: str) -> str:
page_dir = posixpath.dirname(page_dest_uri)
return posixpath.relpath(asset_dest_uri, page_dir)
def _swagger_html(url: str, tag: str | None) -> str:
tag_attr = ""
if tag is not None:
tag_attr = f' data-openapi-tag="{_escape_html_attr(tag)}"'
return f"""
<div class="dstack-swagger-ui" data-openapi-url="{_escape_html_attr(url)}"{tag_attr}></div>
"""
def _openapi_operation_headings(api_file: Path, tag: str | None) -> str:
try:
schema = json.loads(api_file.read_text())
except (OSError, json.JSONDecodeError) as exc:
log.warning(f"Cannot generate Swagger operation headings from {api_file}: {exc}")
return ""
operations = _get_openapi_operations(schema, tag)
if not operations:
return ""
used_ids: set[str] = set()
headings = [_openapi_operation_heading(operation, used_ids) for operation in operations]
return "\n".join(headings) + "\n\n"
def _get_openapi_operations(
schema: dict,
tag: str | None,
) -> list[dict[str, str]]:
operations = []
for path, path_item in schema.get("paths", {}).items():
if not isinstance(path_item, dict):
continue
for method, operation in path_item.items():
method = method.lower()
if method not in HTTP_METHODS or not isinstance(operation, dict):
continue
operation_tags = operation.get("tags") or [UNTAGGED_OPENAPI_TAG]
if tag is not None and tag not in operation_tags:
continue
operations.append(
{
"method": method,
"path": path,
"summary": str(operation.get("summary") or ""),
}
)
return operations
def _openapi_operation_heading(operation: dict[str, str], used_ids: set[str]) -> str:
method = operation["method"]
path = operation["path"]
label = _openapi_operation_label(operation)
anchor_id = _openapi_operation_anchor_id(method, path, used_ids)
attrs = [
f"#{anchor_id}",
".dstack-swagger-operation-anchor",
f"data-toc-label={json.dumps(label)}",
f"data-openapi-method={json.dumps(method)}",
f"data-openapi-path={json.dumps(path)}",
]
return f"## {label} {{ {' '.join(attrs)} }}"
def _openapi_operation_label(operation: dict[str, str]) -> str:
summary = operation.get("summary", "").strip()
if summary:
return summary
return operation["path"]
def _openapi_operation_anchor_id(method: str, path: str, used_ids: set[str]) -> str:
base = re.sub(r"[^a-z0-9]+", "-", f"{method}-{path}".lower()).strip("-") or method
anchor_id = base
index = 2
while anchor_id in used_ids:
anchor_id = f"{base}-{index}"
index += 1
used_ids.add(anchor_id)
return anchor_id
def _escape_html_attr(value: str) -> str:
return escape(value, {'"': """})
def _swagger_error(pre_token: str, post_token: str, message: str) -> str:
return pre_token + escape(f"!! SWAGGER ERROR: {message} !!") + post_token
def on_config(config):
if os.environ.get(DISABLE_YAML_SCHEMAS_ENV):
log.warning("YAML schema reference generation is disabled")
if os.environ.get(DISABLE_LLM_TXT_ENV):
log.warning("llms.txt generation is disabled")
return config
def on_post_build(config):
"""Copy .md files to site (raw) and write .well-known/skills index."""
site_dir = config["site_dir"]
docs_dir = config["docs_dir"]
# Create .nojekyll to prevent GitHub Pages from ignoring .well-known directory
nojekyll_path = os.path.join(site_dir, ".nojekyll")
with open(nojekyll_path, "w") as f:
f.write("")
# Create _config.yml to explicitly include .well-known directory
# This ensures Jekyll (if it runs) includes the .well-known directory
config_yml_path = os.path.join(site_dir, "_config.yml")
with open(config_yml_path, "w") as f:
f.write('include: [".well-known"]\n')
for root, _, files in os.walk(docs_dir):
for file in files:
if not file.endswith(".md"):
continue
src_path = os.path.join(root, file)
rel_path = os.path.relpath(src_path, docs_dir).replace(os.sep, "/")
content = _get_schema_expanded_content(rel_path, config, src_path=src_path)
dest_path = os.path.join(site_dir, rel_path)
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
if content is not None:
# Write expanded schema content
log.info(f"Expanding schema references in {rel_path}")
with open(dest_path, "w", encoding="utf-8") as f:
f.write(content)
else:
# Just copy the file as-is
shutil.copy2(src_path, dest_path)
_write_well_known_skills(config, site_dir)
_generate_llms_files(config, site_dir)
def _write_well_known_skills(config, site_dir):
"""Parse skills/dstack/SKILL.md and write .well-known/skills/index.json. name and description come from frontmatter only."""
repo_root = os.path.dirname(config["config_file_path"])
skill_src = os.path.join(repo_root, *SKILL_PATH)
if not os.path.isfile(skill_src):
return
name = None
description = None
try:
with open(skill_src, "r", encoding="utf-8") as f:
text = f.read()
if text.startswith("---"):
parts = text.split("---", 2)
if len(parts) >= 3:
data = yaml.safe_load(parts[1])
if isinstance(data, dict):
name = data.get("name")
description = data.get("description")
except Exception as e:
log.error(f"Skill parsing error: {e}")
if not name or not description:
log.warning(
"skills/dstack/SKILL.md missing name or description in frontmatter; skipping .well-known/skills"
)
return
out_dir = os.path.join(site_dir, WELL_KNOWN_SKILLS_DIR, name)
os.makedirs(out_dir, exist_ok=True)
shutil.copy2(skill_src, os.path.join(out_dir, "SKILL.md"))
# Serve skill at site root (both skill.md and SKILL.md) from skills/dstack/SKILL.md
shutil.copy2(skill_src, os.path.join(site_dir, "skill.md"))
shutil.copy2(skill_src, os.path.join(site_dir, "SKILL.md"))
index_path = os.path.join(site_dir, WELL_KNOWN_SKILLS_DIR, "index.json")
index = {
"skills": [
{"name": name, "description": description.strip()[:1024], "files": ["SKILL.md"]}
]
}
with open(index_path, "w", encoding="utf-8") as f:
json.dump(index, f, indent=2)
log.info(f"Published skill: {name}")
def _generate_llms_files(config, site_dir):
"""Generate llms.txt and llms-full.txt using external script."""
if os.environ.get(DISABLE_LLM_TXT_ENV):
return
repo_root = os.path.dirname(config["config_file_path"])
# Import and run the generator
hooks_dir = os.path.dirname(os.path.abspath(__file__))
gen_path = os.path.join(hooks_dir, "gen_llms_files.py")
spec = importlib.util.spec_from_file_location("gen_llms_files", gen_path)
if spec is None or spec.loader is None:
log.error(f"Cannot load {gen_path}")
return
module = importlib.util.module_from_spec(spec)
sys.modules["gen_llms_files"] = module
spec.loader.exec_module(module)
try:
# Pass mkdocs config to generator
module.generate_llms_files(repo_root, site_dir, config)
log.info("Generated llms.txt and llms-full.txt")
except Exception as e:
log.error(f"Failed to generate llms files: {e}")