mirror of
https://github.com/thinking-machines-lab/tinker.git
synced 2026-04-19 12:58:01 +00:00
227 lines
7.9 KiB
Python
Executable file
227 lines
7.9 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
# /// script
|
|
# dependencies = [
|
|
# "pydoc-markdown>=4.8.0",
|
|
# "pyyaml>=6.0",
|
|
# ]
|
|
# ///
|
|
|
|
|
|
import ast
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
|
|
def cd_to_project_root():
|
|
"""Change to the project root (parent of the scripts directory)."""
|
|
script_dir = Path(__file__).resolve().parent
|
|
project_root = script_dir.parent
|
|
os.chdir(project_root)
|
|
print(f"Changed to project root: {project_root}")
|
|
|
|
|
|
class ModuleAnalyzer:
|
|
"""Analyze Python modules to extract public API information."""
|
|
|
|
def __init__(self, src_path: Path):
|
|
self.src_path = src_path
|
|
|
|
def get_module_exports(self, module_path: Path) -> List[str]:
|
|
"""Extract __all__ exports from a module."""
|
|
try:
|
|
content = module_path.read_text()
|
|
tree = ast.parse(content)
|
|
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Assign):
|
|
for target in node.targets:
|
|
if (
|
|
isinstance(target, ast.Name)
|
|
and target.id == "__all__"
|
|
and isinstance(node.value, ast.List)
|
|
):
|
|
return [
|
|
elt.s for elt in node.value.elts if isinstance(elt, ast.Str)
|
|
] or [
|
|
elt.value
|
|
for elt in node.value.elts
|
|
if isinstance(elt, ast.Constant) and isinstance(elt.value, str)
|
|
]
|
|
except Exception as e:
|
|
print(f"Warning: Could not parse {module_path}: {e}")
|
|
return []
|
|
|
|
def find_all_modules(self) -> Dict[str, Path]:
|
|
"""Find all Python modules in the package."""
|
|
modules = {}
|
|
tinker_path = self.src_path / "tinker"
|
|
|
|
for py_file in tinker_path.rglob("*.py"):
|
|
# Skip test files and private modules
|
|
if any(part.startswith(("test", "_test")) for part in py_file.parts):
|
|
continue
|
|
if "__pycache__" in py_file.parts:
|
|
continue
|
|
|
|
# Calculate module name
|
|
relative_path = py_file.relative_to(self.src_path)
|
|
module_parts = list(relative_path.parts[:-1]) # Remove .py file
|
|
module_parts.append(relative_path.stem)
|
|
|
|
# Skip __init__ files in module name
|
|
if module_parts[-1] == "__init__":
|
|
module_parts = module_parts[:-1]
|
|
|
|
module_name = ".".join(module_parts)
|
|
if module_name: # Skip empty module names
|
|
modules[module_name] = py_file
|
|
|
|
return modules
|
|
|
|
|
|
class DocumentationGenerator:
|
|
"""Generate documentation using pydoc-markdown."""
|
|
|
|
def __init__(self, config_path: Path, output_dir: Path):
|
|
self.config_path = config_path
|
|
self.output_dir = output_dir
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
self.analyzer = ModuleAnalyzer(Path("src"))
|
|
|
|
def run_pydoc_markdown(self, modules: List[str], output_file: Path) -> bool:
|
|
"""Run pydoc-markdown for specific modules."""
|
|
try:
|
|
# Build the command
|
|
cmd = ["pydoc-markdown", "pydoc-markdown.yml", "-I", "src"]
|
|
|
|
# Add modules
|
|
for module in modules:
|
|
cmd.extend(["-m", module])
|
|
|
|
# Run pydoc-markdown
|
|
result = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
|
|
|
if result.returncode == 0:
|
|
# Write output to file
|
|
output_file.parent.mkdir(parents=True, exist_ok=True)
|
|
output_file.write_text(result.stdout)
|
|
print(f"Generated: {output_file}")
|
|
return True
|
|
else:
|
|
print(f"Error generating {output_file}: {result.stderr}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"Exception generating {output_file}: {e}")
|
|
return False
|
|
|
|
def generate_public_interfaces(self):
|
|
"""Generate documentation for public interface classes."""
|
|
print("\n=== Generating Public Interfaces Documentation ===")
|
|
|
|
# Generate individual pages for each client
|
|
client_modules = [
|
|
("ServiceClient", "tinker.lib.public_interfaces.service_client"),
|
|
("TrainingClient", "tinker.lib.public_interfaces.training_client"),
|
|
("SamplingClient", "tinker.lib.public_interfaces.sampling_client"),
|
|
("RestClient", "tinker.lib.public_interfaces.rest_client"),
|
|
("APIFuture", "tinker.lib.public_interfaces.api_future"),
|
|
]
|
|
|
|
for class_name, module in client_modules:
|
|
output_file = self.output_dir / f"{class_name.lower().replace('_', '-')}.md"
|
|
self.run_pydoc_markdown([module], output_file)
|
|
|
|
def generate_all_types(self):
|
|
"""Generate complete types reference."""
|
|
print("\n=== Generating Complete Types Reference ===")
|
|
|
|
# Get all type modules
|
|
all_modules = self.analyzer.find_all_modules()
|
|
type_modules = [m for m in all_modules if m.startswith("tinker.types")]
|
|
|
|
if type_modules:
|
|
output_file = self.output_dir / "types.md"
|
|
self.run_pydoc_markdown(type_modules, output_file)
|
|
|
|
def generate_exceptions(self):
|
|
"""Generate exception hierarchy documentation."""
|
|
print("\n=== Generating Exception Documentation ===")
|
|
|
|
output_file = self.output_dir / "exceptions.md"
|
|
self.run_pydoc_markdown(["tinker._exceptions"], output_file)
|
|
|
|
def generate_nextra_meta(self):
|
|
"""Generate _meta.json for Nextra navigation."""
|
|
print("\n=== Generating Nextra Navigation Metadata ===")
|
|
|
|
meta = {
|
|
"serviceclient": "ServiceClient",
|
|
"trainingclient": "TrainingClient",
|
|
"samplingclient": "SamplingClient",
|
|
"restclient": "RestClient",
|
|
"apifuture": "APIFuture",
|
|
"types": "Parameters",
|
|
"exceptions": "Exceptions",
|
|
}
|
|
|
|
meta_file = self.output_dir / "_meta.json"
|
|
meta_file.write_text(json.dumps(meta, indent=2))
|
|
print(f"Generated: {meta_file}")
|
|
|
|
def generate_all(self):
|
|
"""Generate all documentation."""
|
|
print("Starting documentation generation...")
|
|
print(f"Output directory: {self.output_dir}")
|
|
|
|
# Generate documentation for each category
|
|
self.generate_public_interfaces()
|
|
self.generate_all_types()
|
|
self.generate_exceptions()
|
|
|
|
# Generate Nextra metadata
|
|
self.generate_nextra_meta()
|
|
|
|
print("\n=== Documentation Generation Complete ===")
|
|
print(f"Markdown files generated in: {self.output_dir}")
|
|
print("\nGenerated files:")
|
|
for file in sorted(self.output_dir.rglob("*.md")):
|
|
print(f" - {file.relative_to(self.output_dir)}")
|
|
|
|
|
|
def main():
|
|
"""Main entry point."""
|
|
# Change to project root first
|
|
cd_to_project_root()
|
|
|
|
# Paths
|
|
project_root = Path.cwd()
|
|
config_path = project_root / "pydoc-markdown.yml"
|
|
output_dir = project_root / "docs" / "api"
|
|
|
|
# Check if config exists
|
|
if not config_path.exists():
|
|
print(f"Error: Configuration file not found at {config_path}")
|
|
print("Please run this script from the project root directory")
|
|
sys.exit(1)
|
|
|
|
# Create generator and run
|
|
generator = DocumentationGenerator(config_path, output_dir)
|
|
generator.generate_all()
|
|
|
|
# Print usage instructions
|
|
print("\n" + "=" * 50)
|
|
print("To use these docs in your Nextra project:")
|
|
print("1. Copy the docs/api directory to your Nextra project")
|
|
print("2. The markdown files are ready to use with Nextra")
|
|
print("3. Navigation structure is defined in _meta.json")
|
|
print("\nTo regenerate docs after code changes:")
|
|
print(" uv run scripts/generate_docs.py")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|