""" Auto-Documentation Generator for Atomizer This module automatically generates documentation from code, ensuring that skills and protocols stay in sync with the implementation. Usage: python -m optimization_engine.auto_doc extractors python -m optimization_engine.auto_doc templates python -m optimization_engine.auto_doc all """ import inspect import importlib import json from pathlib import Path from datetime import datetime from typing import Dict, List, Any, Optional def get_extractor_info() -> List[Dict[str, Any]]: """Extract information about all registered extractors.""" from optimization_engine import extractors extractor_info = [] # Get all exported functions for name in extractors.__all__: obj = getattr(extractors, name) if callable(obj): # Get function signature try: sig = inspect.signature(obj) params = [ { 'name': p.name, 'default': str(p.default) if p.default != inspect.Parameter.empty else None, 'annotation': str(p.annotation) if p.annotation != inspect.Parameter.empty else None } for p in sig.parameters.values() ] except (ValueError, TypeError): params = [] # Get docstring doc = inspect.getdoc(obj) or "No documentation available" # Determine category category = "general" if "stress" in name.lower(): category = "stress" elif "temperature" in name.lower() or "thermal" in name.lower() or "heat" in name.lower(): category = "thermal" elif "modal" in name.lower() or "frequency" in name.lower(): category = "modal" elif "zernike" in name.lower(): category = "optical" elif "mass" in name.lower(): category = "mass" elif "strain" in name.lower(): category = "strain" elif "spc" in name.lower() or "reaction" in name.lower() or "force" in name.lower(): category = "forces" # Determine phase phase = "Phase 1" if name in ['extract_principal_stress', 'extract_max_principal_stress', 'extract_min_principal_stress', 'extract_strain_energy', 'extract_total_strain_energy', 'extract_strain_energy_density', 'extract_spc_forces', 'extract_total_reaction_force', 'extract_reaction_component', 'check_force_equilibrium']: phase = "Phase 2" elif name in ['extract_temperature', 'extract_temperature_gradient', 'extract_heat_flux', 'get_max_temperature', 'extract_modal_mass', 'extract_frequencies', 'get_first_frequency', 'get_modal_mass_ratio']: phase = "Phase 3" extractor_info.append({ 'name': name, 'module': obj.__module__, 'category': category, 'phase': phase, 'parameters': params, 'docstring': doc, 'is_class': inspect.isclass(obj) }) return extractor_info def get_template_info() -> List[Dict[str, Any]]: """Extract information about available study templates.""" templates_file = Path(__file__).parent / 'templates' / 'registry.json' if not templates_file.exists(): return [] with open(templates_file) as f: data = json.load(f) return data.get('templates', []) def generate_extractor_markdown(extractors: List[Dict[str, Any]]) -> str: """Generate markdown documentation for extractors.""" lines = [ "# Atomizer Extractor Library", "", f"*Auto-generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}*", "", "This document is automatically generated from the extractor source code.", "", "---", "", "## Quick Reference", "", "| Extractor | Category | Phase | Description |", "|-----------|----------|-------|-------------|", ] for ext in sorted(extractors, key=lambda x: (x['category'], x['name'])): doc_first_line = ext['docstring'].split('\n')[0][:60] lines.append(f"| `{ext['name']}` | {ext['category']} | {ext['phase']} | {doc_first_line} |") lines.extend(["", "---", ""]) # Group by category categories = {} for ext in extractors: cat = ext['category'] if cat not in categories: categories[cat] = [] categories[cat].append(ext) for cat_name, cat_extractors in sorted(categories.items()): lines.append(f"## {cat_name.title()} Extractors") lines.append("") for ext in sorted(cat_extractors, key=lambda x: x['name']): lines.append(f"### `{ext['name']}`") lines.append("") lines.append(f"**Module**: `{ext['module']}`") lines.append(f"**Phase**: {ext['phase']}") lines.append("") # Parameters if ext['parameters']: lines.append("**Parameters**:") lines.append("") for param in ext['parameters']: default_str = f" = `{param['default']}`" if param['default'] else "" lines.append(f"- `{param['name']}`{default_str}") lines.append("") # Docstring lines.append("**Description**:") lines.append("") lines.append("```") lines.append(ext['docstring']) lines.append("```") lines.append("") lines.append("---") lines.append("") return '\n'.join(lines) def generate_template_markdown(templates: List[Dict[str, Any]]) -> str: """Generate markdown documentation for templates.""" lines = [ "# Atomizer Study Templates", "", f"*Auto-generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}*", "", "Available templates for quick study creation.", "", "---", "", "## Template Reference", "", "| Template | Objectives | Extractors |", "|----------|------------|------------|", ] for tmpl in templates: # Handle objectives that might be dicts or strings obj_list = tmpl.get('objectives', []) if obj_list and isinstance(obj_list[0], dict): objectives = ', '.join([o.get('name', str(o)) for o in obj_list]) else: objectives = ', '.join(obj_list) extractors = ', '.join(tmpl.get('extractors', [])) lines.append(f"| `{tmpl['name']}` | {objectives} | {extractors} |") lines.extend(["", "---", ""]) for tmpl in templates: lines.append(f"## {tmpl['name']}") lines.append("") lines.append(f"**Description**: {tmpl.get('description', 'N/A')}") lines.append("") lines.append(f"**Category**: {tmpl.get('category', 'N/A')}") lines.append(f"**Solver**: {tmpl.get('solver', 'N/A')}") lines.append(f"**Sampler**: {tmpl.get('sampler', 'N/A')}") lines.append(f"**Turbo Suitable**: {'Yes' if tmpl.get('turbo_suitable') else 'No'}") lines.append("") lines.append(f"**Example Study**: `{tmpl.get('example_study', 'N/A')}`") lines.append("") if tmpl.get('objectives'): lines.append("**Objectives**:") for obj in tmpl['objectives']: if isinstance(obj, dict): lines.append(f"- {obj.get('name', '?')} ({obj.get('direction', '?')}) - Extractor: {obj.get('extractor', '?')}") else: lines.append(f"- {obj}") lines.append("") if tmpl.get('extractors'): lines.append("**Extractors Used**:") for ext in tmpl['extractors']: lines.append(f"- {ext}") lines.append("") if tmpl.get('recommended_trials'): lines.append("**Recommended Trials**:") for key, val in tmpl['recommended_trials'].items(): lines.append(f"- {key}: {val}") lines.append("") lines.append("---") lines.append("") return '\n'.join(lines) def generate_cheatsheet_update(extractors: List[Dict[str, Any]]) -> str: """Generate the extractor quick reference for 01_CHEATSHEET.md.""" lines = [ "## Extractor Quick Reference", "", "| Physics | Extractor | Function Call |", "|---------|-----------|---------------|", ] # Map categories to physics names physics_map = { 'stress': 'Von Mises stress', 'thermal': 'Temperature', 'modal': 'Natural frequency', 'optical': 'Zernike WFE', 'mass': 'Mass', 'strain': 'Strain energy', 'forces': 'Reaction forces', 'general': 'Displacement', } for ext in sorted(extractors, key=lambda x: x['category']): if ext['is_class']: continue physics = physics_map.get(ext['category'], ext['category']) # Build function call example params = ext['parameters'][:2] if ext['parameters'] else [] param_str = ', '.join([p['name'] for p in params]) lines.append(f"| {physics} | {ext['name']} | `{ext['name']}({param_str})` |") return '\n'.join(lines) def update_atomizer_context(extractors: List[Dict[str, Any]], templates: List[Dict[str, Any]]): """Update ATOMIZER_CONTEXT.md with current extractor count.""" context_file = Path(__file__).parent.parent / '.claude' / 'ATOMIZER_CONTEXT.md' if not context_file.exists(): print(f"Warning: {context_file} not found") return content = context_file.read_text() # Update extractor library version based on count extractor_count = len(extractors) template_count = len(templates) print(f"Found {extractor_count} extractors and {template_count} templates") # Could add logic here to update version info based on changes def main(): import sys if len(sys.argv) < 2: print("Usage: python -m optimization_engine.auto_doc [extractors|templates|all]") sys.exit(1) command = sys.argv[1] output_dir = Path(__file__).parent.parent / 'docs' / 'generated' output_dir.mkdir(parents=True, exist_ok=True) if command in ['extractors', 'all']: print("Generating extractor documentation...") extractors = get_extractor_info() # Write full documentation doc_content = generate_extractor_markdown(extractors) (output_dir / 'EXTRACTORS.md').write_text(doc_content) print(f" Written: {output_dir / 'EXTRACTORS.md'}") # Write cheatsheet update cheatsheet = generate_cheatsheet_update(extractors) (output_dir / 'EXTRACTOR_CHEATSHEET.md').write_text(cheatsheet) print(f" Written: {output_dir / 'EXTRACTOR_CHEATSHEET.md'}") print(f" Found {len(extractors)} extractors") if command in ['templates', 'all']: print("Generating template documentation...") templates = get_template_info() if templates: doc_content = generate_template_markdown(templates) (output_dir / 'TEMPLATES.md').write_text(doc_content) print(f" Written: {output_dir / 'TEMPLATES.md'}") print(f" Found {len(templates)} templates") else: print(" No templates found") if command == 'all': print("\nUpdating ATOMIZER_CONTEXT.md...") extractors = get_extractor_info() templates = get_template_info() update_atomizer_context(extractors, templates) print("\nDone!") if __name__ == '__main__': main()