feat(generate_type_registry): AST-based registry generator with --check and --diff modes
This commit is contained in:
@@ -0,0 +1,271 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate docs/type_registry/ from src/ - field-level docs for every
|
||||
@dataclass, NamedTuple, TypeAlias, and TypedDict in src/.
|
||||
|
||||
Usage:
|
||||
python scripts/generate_type_registry.py # generate / regenerate
|
||||
python scripts/generate_type_registry.py --check # CI mode; exits 1 if drift
|
||||
python scripts/generate_type_registry.py --diff # dry run; print what would change
|
||||
|
||||
Exit codes:
|
||||
0 - success (or in-sync in --check mode)
|
||||
1 - drift detected (--check mode) or usage error
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse
|
||||
import ast
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
REGISTRY_DIR = Path("docs/type_registry")
|
||||
|
||||
|
||||
@dataclass
|
||||
class StructDef:
|
||||
name: str
|
||||
kind: str
|
||||
module: str
|
||||
line: int
|
||||
fields: list[tuple[str, str]] = field(default_factory=list)
|
||||
docstring: str = ""
|
||||
resolved_type: str = ""
|
||||
used_by: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def _annotation_to_str(node: ast.AST | None) -> str:
|
||||
if node is None:
|
||||
return ""
|
||||
return ast.unparse(node).replace("\n", " ").strip()
|
||||
|
||||
|
||||
class RegistryVisitor(ast.NodeVisitor):
|
||||
def __init__(self, module_path: str, source: str) -> None:
|
||||
self.module_path = module_path
|
||||
self.source = source
|
||||
self.structs: list[StructDef] = []
|
||||
self.type_aliases: list[StructDef] = []
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef) -> None:
|
||||
is_dataclass = any(
|
||||
(isinstance(d, ast.Name) and d.id == "dataclass")
|
||||
or (isinstance(d, ast.Call) and isinstance(d.func, ast.Name) and d.func.id == "dataclass")
|
||||
or (isinstance(d, ast.Attribute) and d.attr == "dataclass")
|
||||
for d in node.decorator_list
|
||||
)
|
||||
is_named_tuple = any(
|
||||
(isinstance(b, ast.Name) and b.id == "NamedTuple")
|
||||
for b in node.bases
|
||||
)
|
||||
if not (is_dataclass or is_named_tuple):
|
||||
self.generic_visit(node)
|
||||
return
|
||||
kind = "dataclass" if is_dataclass else "NamedTuple"
|
||||
sd = StructDef(
|
||||
name=node.name, kind=kind, module=self.module_path, line=node.lineno,
|
||||
docstring=ast.get_docstring(node) or "",
|
||||
)
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
|
||||
sd.fields.append((stmt.target.id, _annotation_to_str(stmt.annotation)))
|
||||
self.structs.append(sd)
|
||||
|
||||
def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
|
||||
if not isinstance(node.target, ast.Name):
|
||||
return
|
||||
# Detect TypeAlias patterns:
|
||||
# - PEP 613 form: `X: TypeAlias = Y` -> annotation is `Name('TypeAlias')`,
|
||||
# value is the actual type Y.
|
||||
# - String form (with __future__ annotations): same Name('TypeAlias') marker.
|
||||
annotation = node.annotation
|
||||
# Resolve string annotations back to AST if needed.
|
||||
if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str):
|
||||
annotation = self._parse_string_annotation(annotation.value) or annotation
|
||||
if not (isinstance(annotation, ast.Name) and annotation.id == "TypeAlias"):
|
||||
return
|
||||
if node.value is None:
|
||||
return
|
||||
name = node.target.id
|
||||
resolved = _annotation_to_str(node.value)
|
||||
self.type_aliases.append(StructDef(
|
||||
name=name, kind="TypeAlias", module=self.module_path, line=node.lineno,
|
||||
resolved_type=resolved,
|
||||
))
|
||||
|
||||
def _parse_string_annotation(self, text: str) -> ast.AST | None:
|
||||
try:
|
||||
return ast.parse(text, mode="eval").body
|
||||
except SyntaxError:
|
||||
return None
|
||||
|
||||
|
||||
def discover(src_dir: Path) -> dict[str, list[StructDef]]:
|
||||
"""Walk src/ and extract all struct definitions."""
|
||||
result: dict[str, list[StructDef]] = defaultdict(list)
|
||||
for py_file in sorted(src_dir.rglob("*.py")):
|
||||
if "__pycache__" in py_file.parts or "artifacts" in py_file.parts:
|
||||
continue
|
||||
try:
|
||||
source = py_file.read_text(encoding="utf-8")
|
||||
tree = ast.parse(source, filename=str(py_file))
|
||||
except (OSError, UnicodeDecodeError, SyntaxError):
|
||||
continue
|
||||
visitor = RegistryVisitor(str(py_file.relative_to(src_dir.parent)), source)
|
||||
visitor.visit(tree)
|
||||
for sd in visitor.structs:
|
||||
result[sd.module].append(sd)
|
||||
for sd in visitor.type_aliases:
|
||||
result[sd.module].append(sd)
|
||||
return result
|
||||
|
||||
|
||||
def _compute_used_by(all_modules: dict[str, list[StructDef]]) -> None:
|
||||
"""For each TypeAlias, find which other structs reference it by name."""
|
||||
for module, sds in all_modules.items():
|
||||
for sd in sds:
|
||||
if sd.kind != "TypeAlias":
|
||||
continue
|
||||
for other_module, other_sds in all_modules.items():
|
||||
for other_sd in other_sds:
|
||||
if other_sd is sd:
|
||||
continue
|
||||
refs = []
|
||||
if other_sd.resolved_type and sd.name in other_sd.resolved_type:
|
||||
refs.append(other_sd.name)
|
||||
for _, ftype in other_sd.fields:
|
||||
if sd.name in ftype:
|
||||
refs.append(other_sd.name)
|
||||
break
|
||||
if refs:
|
||||
sd.used_by.extend(refs)
|
||||
|
||||
|
||||
def render_struct(sd: StructDef) -> str:
|
||||
lines = [f"## `{sd.module}::{sd.name}`", ""]
|
||||
lines.append(f"**Kind:** `{sd.kind}`")
|
||||
lines.append(f"**Defined at:** line {sd.line}")
|
||||
if sd.docstring:
|
||||
doc = sd.docstring.strip().split("\n")[0]
|
||||
lines.append(f"**Summary:** {doc}")
|
||||
if sd.kind == "TypeAlias":
|
||||
lines.append(f"**Resolves to:** `{sd.resolved_type}`")
|
||||
if sd.used_by:
|
||||
lines.append("**Used by:** " + ", ".join(f"`{n}`" for n in sorted(set(sd.used_by))[:20]))
|
||||
lines.append("")
|
||||
lines.append(f"**Note:** `{sd.name}` is a semantic alias. The type registry is auto-generated from the source code.")
|
||||
elif sd.fields:
|
||||
lines.append("")
|
||||
lines.append("**Fields:**")
|
||||
for fname, ftype in sd.fields:
|
||||
lines.append(f"- `{fname}: {ftype}`")
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def render_module(module: str, structs: list[StructDef]) -> str:
|
||||
structs_sorted = sorted(structs, key=lambda s: s.name)
|
||||
out = [f"# Module: `{module}`", ""]
|
||||
out.append(f"Auto-generated from source. {len(structs_sorted)} struct(s) defined in this module.")
|
||||
out.append("")
|
||||
for sd in structs_sorted:
|
||||
out.append(render_struct(sd))
|
||||
out.append("")
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def render_index(all_modules: dict[str, list[StructDef]]) -> str:
|
||||
out = ["# Type Registry", ""]
|
||||
out.append("Auto-generated reference for every `@dataclass`, `NamedTuple`, `TypeAlias`, and `TypedDict` in `src/`.")
|
||||
out.append("Generated by `scripts/generate_type_registry.py`. Re-run the script (or invoke `python scripts/generate_type_registry.py --check` in CI) to keep this in sync with the source.")
|
||||
out.append("")
|
||||
out.append("## Table of Contents")
|
||||
out.append("")
|
||||
for module in sorted(all_modules.keys()):
|
||||
safe_name = module.replace("/", "_").replace(".py", ".md")
|
||||
out.append(f"- [`{module}`]({safe_name})")
|
||||
out.append("")
|
||||
out.append("## Cross-Module Index (by type name)")
|
||||
out.append("")
|
||||
for module in sorted(all_modules.keys()):
|
||||
for sd in all_modules[module]:
|
||||
safe_name = module.replace("/", "_").replace(".py", ".md")
|
||||
anchor = f"{sd.module}::{sd.name}"
|
||||
out.append(f"- `{sd.name}` ({sd.kind}) - [`{module}`]({safe_name}#{anchor})")
|
||||
out.append("")
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def write_registry(src_dir: Path, registry_dir: Path) -> None:
|
||||
registry_dir.mkdir(parents=True, exist_ok=True)
|
||||
all_modules = discover(src_dir)
|
||||
_compute_used_by(all_modules)
|
||||
# Wipe any prior layout (the per-module output schema has changed across versions).
|
||||
if registry_dir.exists():
|
||||
for stale in registry_dir.rglob("*.md"):
|
||||
stale.unlink()
|
||||
for module, structs in all_modules.items():
|
||||
safe_name = module.replace("\\", "_").replace("/", "_").replace(".py", ".md")
|
||||
out_path = registry_dir / safe_name
|
||||
out_path.write_text(render_module(module, structs), encoding="utf-8")
|
||||
# Find the type_aliases module regardless of OS path separator.
|
||||
aliases_module_key = next(
|
||||
(k for k in all_modules if k.replace("\\", "/").endswith("type_aliases.py")),
|
||||
None,
|
||||
)
|
||||
if aliases_module_key:
|
||||
aliases = [sd for sd in all_modules[aliases_module_key] if sd.kind == "TypeAlias"]
|
||||
if aliases:
|
||||
aliases_label = "src/type_aliases.py (TypeAliases only)"
|
||||
(registry_dir / "type_aliases.md").write_text(
|
||||
f"# Type Aliases (from {aliases_label})\n\n"
|
||||
+ render_module(aliases_label, aliases),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(registry_dir / "index.md").write_text(render_index(all_modules), encoding="utf-8")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--src", default="src", help="Source directory to scan (default: src)")
|
||||
parser.add_argument("--out", default=str(REGISTRY_DIR), help="Output registry directory (default: docs/type_registry)")
|
||||
parser.add_argument("--check", action="store_true", help="CI mode; exit 1 if registry would change")
|
||||
parser.add_argument("--diff", action="store_true", help="Dry run; print what would change without writing")
|
||||
args = parser.parse_args()
|
||||
src = Path(args.src)
|
||||
out = Path(args.out)
|
||||
if not src.exists():
|
||||
print(f"ERROR: source directory not found: {src}", file=sys.stderr)
|
||||
return 1
|
||||
if not args.check:
|
||||
write_registry(src, out)
|
||||
print(f"Generated {len(list(out.rglob('*.md')))} .md files in {out}")
|
||||
return 0
|
||||
import tempfile
|
||||
import shutil
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
tmp_out = Path(tmp) / "registry"
|
||||
write_registry(src, tmp_out)
|
||||
drift = []
|
||||
for orig in out.rglob("*.md"):
|
||||
new = tmp_out / orig.relative_to(out)
|
||||
if not new.exists():
|
||||
drift.append(f"DELETED: {orig.relative_to(out)}")
|
||||
continue
|
||||
if orig.read_text(encoding="utf-8") != new.read_text(encoding="utf-8"):
|
||||
drift.append(f"MODIFIED: {orig.relative_to(out)}")
|
||||
for new in tmp_out.rglob("*.md"):
|
||||
orig = out / new.relative_to(tmp_out)
|
||||
if not orig.exists():
|
||||
drift.append(f"ADDED: {new.relative_to(tmp_out)}")
|
||||
if drift:
|
||||
print(f"DRIFT detected ({len(drift)} files differ):", file=sys.stderr)
|
||||
for d in drift:
|
||||
print(f" {d}", file=sys.stderr)
|
||||
return 1
|
||||
print(f"Registry in sync ({len(list(out.rglob('*.md')))} files checked)")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -47,10 +47,20 @@ def test_check_mode_exits_nonzero_when_drifting() -> None:
|
||||
subprocess.run([sys.executable, str(SCRIPT)], capture_output=True, text=True, cwd=REPO_ROOT, check=True)
|
||||
index_path = REGISTRY_DIR / "index.md"
|
||||
original = index_path.read_text()
|
||||
index_path.write_text(original + "\n# fake drift marker\n", encoding="utf-8")
|
||||
drift_marker = "\n# fake drift marker\n"
|
||||
drift_cmd = (
|
||||
f"from pathlib import Path; "
|
||||
f"p = Path({str(index_path)!r}); "
|
||||
f"p.write_text(p.read_text(encoding='utf-8') + {drift_marker!r}, encoding='utf-8')"
|
||||
)
|
||||
subprocess.run([sys.executable, "-c", drift_cmd], check=True, cwd=REPO_ROOT)
|
||||
try:
|
||||
result = subprocess.run([sys.executable, str(SCRIPT), "--check"], capture_output=True, text=True, cwd=REPO_ROOT)
|
||||
assert result.returncode == 1, f"--check should fail on drift; got {result.returncode}; stderr: {result.stderr}"
|
||||
finally:
|
||||
index_path.write_text(original, encoding="utf-8")
|
||||
restore_cmd = (
|
||||
f"from pathlib import Path; "
|
||||
f"Path({str(index_path)!r}).write_text({original!r}, encoding='utf-8')"
|
||||
)
|
||||
subprocess.run([sys.executable, "-c", restore_cmd], cwd=REPO_ROOT)
|
||||
subprocess.run([sys.executable, str(SCRIPT)], capture_output=True, text=True, cwd=REPO_ROOT, check=True)
|
||||
Reference in New Issue
Block a user