ctx-gen
Version:
AI-Enhanced Documentation Generator for Code Understanding
615 lines • 23.9 kB
JavaScript
import fs from 'fs-extra';
import path from 'path';
import { glob } from 'glob';
import { writeFileWithDir, getFileLanguage } from '../utils/fileUtils.js';
/**
* Generate diagrams for code visualization
* @param options - CtxGen options
*/
export async function generateDiagrams(options) {
// Create diagrams directory if it doesn't exist
const diagramDir = path.join(options.docsDir, 'diagrams');
await fs.ensureDir(diagramDir);
// Get diagram types from options
const diagramTypes = options.diagrams.split(',');
// Generate requested diagrams
for (const type of diagramTypes) {
switch (type.trim()) {
case 'call-graph':
await generateCallGraph(options);
break;
case 'data-flow':
await generateDataFlow(options);
break;
case 'module-deps':
await generateModuleDependencies(options);
break;
case 'class-diagram':
await generateClassDiagram(options);
break;
default:
console.warn(`Unknown diagram type: ${type}`);
}
}
}
/**
* Generate a call graph diagram
* @param options - CtxGen options
*/
async function generateCallGraph(options) {
const diagramPath = path.join(options.docsDir, 'diagrams', 'call_graph.mmd');
// Start with the Mermaid diagram header
let diagram = '%%{init: {\'theme\': \'base\', \'themeVariables\': { \'primaryColor\': \'#f8f9fa\', \'primaryBorderColor\': \'#adb5bd\', \'primaryTextColor\': \'#343a40\'}}}%%\n';
diagram += 'flowchart TD\n';
// Get all source files
const excludePatterns = options.exclude.split(',');
const languages = options.languages.split(',');
// Focus on TypeScript/JavaScript files for call graphs
// as they're easier to analyze without a full compiler
const jstsLanguages = languages.filter(l => l === 'typescript' || l === 'javascript' || l === 'ts' || l === 'js');
if (jstsLanguages.length === 0) {
// If no TS/JS files, create a simplified diagram
diagram += ' Note["No TypeScript/JavaScript files found for call graph"] \n';
await writeFileWithDir(diagramPath, diagram);
return;
}
// Create glob patterns for each language
const patterns = jstsLanguages.map(lang => `**/*.${lang}`);
// Get all files matching the patterns
const files = await glob(patterns, {
ignore: excludePatterns,
cwd: process.cwd(),
absolute: true,
});
// Simple exported function detection
// This is a simplified approach - a real implementation would use AST parsing
const functions = {};
const calls = {};
for (const file of files) {
try {
const content = await fs.readFile(file, 'utf-8');
const moduleName = path.basename(file, path.extname(file));
// Extract exported function names (simplified)
const exportedFunctions = extractExportedFunctions(content);
if (exportedFunctions.length > 0) {
functions[moduleName] = exportedFunctions;
// Extract function calls (simplified)
const functionCalls = extractFunctionCalls(content, exportedFunctions, Object.values(functions).flat());
calls[moduleName] = functionCalls;
}
}
catch (_error) {
console.error(`Error analyzing ${file}`);
}
}
// Create nodes for modules
Object.keys(functions).forEach(moduleName => {
diagram += ` ${sanitizeId(moduleName)}[${moduleName}]\n`;
});
// Create nodes for functions
Object.entries(functions).forEach(([moduleName, functionNames]) => {
functionNames.forEach(functionName => {
diagram += ` ${sanitizeId(moduleName)}_${sanitizeId(functionName)}["${functionName}"]\n`;
diagram += ` ${sanitizeId(moduleName)} --> ${sanitizeId(moduleName)}_${sanitizeId(functionName)}\n`;
});
});
// Create edges for function calls
Object.entries(calls).forEach(([moduleName, calledFunctions]) => {
calledFunctions.forEach(called => {
// Find which module contains the called function
let targetModule = '';
Object.entries(functions).forEach(([module, funcs]) => {
if (funcs.includes(called)) {
targetModule = module;
}
});
if (targetModule && targetModule !== moduleName) {
diagram += ` ${sanitizeId(moduleName)} -.-> ${sanitizeId(targetModule)}\n`;
}
});
});
await writeFileWithDir(diagramPath, diagram);
}
/**
* Generate a data flow diagram
* @param options - CtxGen options
*/
async function generateDataFlow(options) {
const diagramPath = path.join(options.docsDir, 'diagrams', 'data_flow.mmd');
// Start with the Mermaid diagram header
let diagram = '%%{init: {\'theme\': \'base\', \'themeVariables\': { \'primaryColor\': \'#e9f5db\', \'primaryBorderColor\': \'#87986a\', \'primaryTextColor\': \'#344e41\'}}}%%\n';
diagram += 'flowchart LR\n';
// Get package info if available
let hasCli = false;
let hasCore = false;
let hasUtils = false;
try {
// Check for common directories
hasCli = (await fs.pathExists('src/cli')) || (await fs.pathExists('lib/cli'));
hasCore = (await fs.pathExists('src/core')) || (await fs.pathExists('lib/core'));
hasUtils = (await fs.pathExists('src/utils')) || (await fs.pathExists('lib/utils'));
}
catch (_error) {
// Ignore errors
}
// Create a simplified data flow diagram based on common patterns
diagram += ' Input[/"Input Data"/]\n';
if (hasCli) {
diagram += ' CLI["CLI Interface"]\n';
diagram += ' Input --> CLI\n';
if (hasCore) {
diagram += ' Core["Core Processing"]\n';
diagram += ' CLI --> Core\n';
}
else {
diagram += ' Process["Processing"]\n';
diagram += ' CLI --> Process\n';
}
}
else if (hasCore) {
diagram += ' Core["Core Processing"]\n';
diagram += ' Input --> Core\n';
}
else {
diagram += ' Process["Processing"]\n';
diagram += ' Input --> Process\n';
}
if (hasUtils) {
diagram += ' Utils["Utilities"]\n';
if (hasCore) {
diagram += ' Core --> Utils\n';
}
else {
diagram += ' Process --> Utils\n';
}
}
diagram += ' Output[/"Output Data"/]\n';
if (hasCore) {
diagram += ' Core --> Output\n';
}
else {
diagram += ' Process --> Output\n';
}
await writeFileWithDir(diagramPath, diagram);
}
/**
* Generate a module dependencies diagram
* @param options - CtxGen options
*/
async function generateModuleDependencies(options) {
const diagramPath = path.join(options.docsDir, 'diagrams', 'module_deps.mmd');
// Start with the Mermaid diagram header
let diagram = '%%{init: {\'theme\': \'base\', \'themeVariables\': { \'primaryColor\': \'#f1faee\', \'primaryBorderColor\': \'#a8dadc\', \'primaryTextColor\': \'#1d3557\'}}}%%\n';
diagram += 'flowchart TD\n';
// Get all source files
const excludePatterns = options.exclude.split(',');
const languages = options.languages.split(',');
// Create glob patterns for each language
const patterns = languages.map(lang => `**/*.${lang}`);
// Get all files matching the patterns
const files = await glob(patterns, {
ignore: excludePatterns,
cwd: process.cwd(),
absolute: true,
});
// Group files by directory (module)
const modules = {};
for (const file of files) {
const dir = path.dirname(path.relative(process.cwd(), file));
if (!modules[dir]) {
modules[dir] = [];
}
modules[dir].push(file);
}
// Extract imports between modules
const imports = {};
for (const [module, moduleFiles] of Object.entries(modules)) {
imports[module] = new Set();
for (const file of moduleFiles) {
try {
const content = await fs.readFile(file, 'utf-8');
// Extract imports (simplified)
const moduleImports = extractImportedModules(content, Object.keys(modules));
for (const importedModule of moduleImports) {
if (importedModule !== module) {
imports[module].add(importedModule);
}
}
}
catch (error) {
console.error(`Error analyzing imports in ${file}:`, error);
}
}
}
// Create nodes for modules
Object.keys(modules).forEach(module => {
if (module === '.')
return;
// Clean up module name
const moduleName = module.replace(/^src\/|^lib\//, '');
diagram += ` ${sanitizeId(moduleName)}["${moduleName}"]\n`;
});
// Create edges for imports
Object.entries(imports).forEach(([module, moduleImports]) => {
if (module === '.')
return;
// Clean up module name
const moduleName = module.replace(/^src\/|^lib\//, '');
moduleImports.forEach(importedModule => {
if (importedModule === '.')
return;
// Clean up imported module name
const importedName = importedModule.replace(/^src\/|^lib\//, '');
diagram += ` ${sanitizeId(moduleName)} --> ${sanitizeId(importedName)}\n`;
});
});
await writeFileWithDir(diagramPath, diagram);
}
/**
* Generate a class diagram showing classes and their relationships
* @param options - CtxGen options
*/
async function generateClassDiagram(options) {
const diagramPath = path.join(options.docsDir, 'diagrams', 'class_diagram.mmd');
// Start with the Mermaid diagram header
let diagram = '%%{init: {\'theme\': \'base\', \'themeVariables\': { \'primaryColor\': \'#f0f4fc\', \'primaryBorderColor\': \'#6c8ebf\', \'primaryTextColor\': \'#0d2644\'}}}%%\n';
diagram += 'classDiagram\n';
// Get all source files
const excludePatterns = options.exclude.split(',');
const languages = options.languages.split(',');
// Focus on languages that typically use classes
const classLanguages = languages.filter(l => ['typescript', 'javascript', 'java', 'python', 'csharp', 'cpp'].includes(l));
if (classLanguages.length === 0) {
// If no languages with classes, create a simplified diagram
diagram += ' class NoteClass{\n';
diagram += ' No class-based languages found\n';
diagram += ' for generating class diagram\n';
diagram += ' }\n';
await writeFileWithDir(diagramPath, diagram);
return;
}
// Create glob patterns for each language
const patterns = classLanguages
.map(lang => {
if (lang === 'typescript')
return ['**/*.ts', '**/*.tsx'];
if (lang === 'javascript')
return ['**/*.js', '**/*.jsx'];
if (lang === 'python')
return ['**/*.py'];
if (lang === 'java')
return ['**/*.java'];
if (lang === 'csharp')
return ['**/*.cs'];
if (lang === 'cpp')
return ['**/*.cpp', '**/*.hpp'];
return [`**/*.${lang}`];
})
.flat();
// Get all files matching the patterns
const files = await glob(patterns, {
ignore: excludePatterns,
cwd: process.cwd(),
absolute: true,
});
// Extract class information
const classes = {};
const relationships = [];
// Process each file to extract class information
for (const file of files) {
try {
const content = await fs.readFile(file, 'utf-8');
const language = getFileLanguage(file);
if (!language)
continue;
// Extract classes based on language
if (language === 'typescript' || language === 'javascript') {
extractTypeScriptClasses(content, classes, relationships);
}
else if (language === 'python') {
extractPythonClasses(content, classes, relationships);
}
else if (language === 'java') {
extractJavaClasses(content, classes, relationships);
}
// Add more language extractors as needed
}
catch (error) {
console.error(`Error analyzing ${file} for class diagram:`, error);
}
}
// Generate class definitions for mermaid
Object.entries(classes).forEach(([className, classInfo]) => {
diagram += ` class ${sanitizeId(className)} {\n`;
// Add properties
classInfo.properties.forEach(prop => {
diagram += ` ${prop}\n`;
});
// Add methods
classInfo.methods.forEach(method => {
diagram += ` ${method}()\n`;
});
diagram += ' }\n';
});
// Generate class relationships
relationships.forEach(rel => {
if (rel.type === 'extends') {
diagram += ` ${sanitizeId(rel.from)} --|> ${sanitizeId(rel.to)}\n`;
}
else if (rel.type === 'implements') {
diagram += ` ${sanitizeId(rel.from)} ..|> ${sanitizeId(rel.to)}\n`;
}
else if (rel.type === 'association') {
diagram += ` ${sanitizeId(rel.from)} --> ${sanitizeId(rel.to)}`;
if (rel.label) {
diagram += ` : ${rel.label}`;
}
diagram += '\n';
}
else if (rel.type === 'composition') {
diagram += ` ${sanitizeId(rel.from)} *-- ${sanitizeId(rel.to)}\n`;
}
else if (rel.type === 'aggregation') {
diagram += ` ${sanitizeId(rel.from)} o-- ${sanitizeId(rel.to)}\n`;
}
});
await writeFileWithDir(diagramPath, diagram);
}
/**
* Extract exported function names from file content
* @param content - File content
* @returns Array of exported function names
*/
function extractExportedFunctions(content) {
const functions = [];
// Match exported functions (simplified)
const exportRegex = /export\s+(async\s+)?function\s+(\w+)/g;
let match;
while ((match = exportRegex.exec(content)) !== null) {
functions.push(match[2]);
}
// Also match export const name = function
const exportConstRegex = /export\s+const\s+(\w+)\s*=\s*(async\s+)?function/g;
while ((match = exportConstRegex.exec(content)) !== null) {
functions.push(match[1]);
}
// Also match export const name = () =>
const exportArrowRegex = /export\s+const\s+(\w+)\s*=\s*(async\s+)?\([^)]*\)\s*=>/g;
while ((match = exportArrowRegex.exec(content)) !== null) {
functions.push(match[1]);
}
return functions;
}
/**
* Extract function calls from file content
* @param content - File content
* @param localFunctions - Functions in the current file
* @param allFunctions - All functions across files
* @returns Array of called function names
*/
function extractFunctionCalls(content, localFunctions, allFunctions) {
const calls = new Set();
// Check for calls to functions that are exported from other files
for (const func of allFunctions) {
if (!localFunctions.includes(func)) {
const callPattern = new RegExp(`[^\\w]${func}\\s*\\(`, 'g');
if (callPattern.test(content)) {
calls.add(func);
}
}
}
return Array.from(calls);
}
/**
* Extract imported modules from file content
* @param content - File content
* @param allModules - All module directory paths
* @returns Array of imported module paths
*/
function extractImportedModules(content, allModules) {
const modules = new Set();
// Match import statements
const importRegex = /from\s+['"]([^'"]+)['"]/g;
let match;
while ((match = importRegex.exec(content)) !== null) {
const importPath = match[1];
// Convert relative import to absolute path
if (importPath.startsWith('.')) {
// This is simplified and doesn't handle complex cases
// A real implementation would resolve paths properly
continue;
}
// Find matching module
for (const module of allModules) {
if (importPath.includes(module)) {
modules.add(module);
break;
}
}
}
return Array.from(modules);
}
/**
* Sanitize a string for use as a Mermaid flowchart ID
* @param id - String to sanitize
* @returns Sanitized ID
*/
function sanitizeId(id) {
return id.replace(/[^a-zA-Z0-9]/g, '_');
}
/**
* Extract classes from TypeScript/JavaScript code
* @param content - File content
* @param classes - Classes collection to update
* @param relationships - Relationships collection to update
*/
function extractTypeScriptClasses(content, classes, relationships) {
// Match class declarations
const classRegex = /class\s+(\w+)(?:\s+extends\s+(\w+))?(?:\s+implements\s+([^{]+))?/g;
let match;
while ((match = classRegex.exec(content)) !== null) {
const className = match[1];
const extendsClass = match[2];
const implementsList = match[3]?.split(',').map(i => i.trim());
// Initialize class if not exists
if (!classes[className]) {
classes[className] = {
properties: [],
methods: [],
};
}
if (extendsClass) {
classes[className].extends = extendsClass;
// Add inheritance relationship
relationships.push({
from: className,
to: extendsClass,
type: 'extends',
});
}
if (implementsList && implementsList.length > 0) {
classes[className].implements = implementsList;
// Add implementation relationships
implementsList.forEach(impl => {
relationships.push({
from: className,
to: impl,
type: 'implements',
});
});
}
// Extract class body
const classBodyMatch = content.substring(match.index).match(/{([^{}]*(?:{[^{}]*}[^{}]*)*)}/);
if (classBodyMatch) {
const classBody = classBodyMatch[1];
// Extract properties
const propertyRegex = /(?:public|private|protected)?\s+(\w+)\s*:/g;
let propMatch;
while ((propMatch = propertyRegex.exec(classBody)) !== null) {
classes[className].properties.push(propMatch[1]);
}
// Extract methods
const methodRegex = /(?:public|private|protected)?\s+(\w+)\s*\([^)]*\)/g;
let methodMatch;
while ((methodMatch = methodRegex.exec(classBody)) !== null) {
classes[className].methods.push(methodMatch[1]);
}
}
}
}
/**
* Extract classes from Python code
* @param content - File content
* @param classes - Classes collection to update
* @param relationships - Relationships collection to update
*/
function extractPythonClasses(content, classes, relationships) {
// Match class declarations with inheritance
const classRegex = /class\s+(\w+)(?:\s*\(([^)]+)\))?:/g;
let match;
while ((match = classRegex.exec(content)) !== null) {
const className = match[1];
const parentClasses = match[2]?.split(',').map(c => c.trim());
// Initialize class if not exists
if (!classes[className]) {
classes[className] = {
properties: [],
methods: [],
};
}
// Handle inheritance
if (parentClasses && parentClasses.length > 0) {
// In Python, first base class is considered the main one
classes[className].extends = parentClasses[0];
// Add relationships
parentClasses.forEach(parent => {
relationships.push({
from: className,
to: parent,
type: 'extends',
});
});
}
// Extract methods (simplified)
const methodRegex = /def\s+(\w+)\s*\(/g;
let methodMatch;
while ((methodMatch = methodRegex.exec(content)) !== null) {
const methodName = methodMatch[1];
// Skip dunder methods
if (!methodName.startsWith('__') || methodName === '__init__') {
classes[className].methods.push(methodName);
}
}
// Extract properties (simplified - from __init__ method)
const initMatch = content.match(new RegExp('def\\s+__init__\\s*\\([^)]*\\):\\s*([^]*?)(?:def|$)', 's'));
if (initMatch) {
const initBody = initMatch[1];
const propRegex = /self\.(\w+)\s*=/g;
let propMatch;
while ((propMatch = propRegex.exec(initBody)) !== null) {
classes[className].properties.push(propMatch[1]);
}
}
}
}
/**
* Extract classes from Java code
* @param content - File content
* @param classes - Classes collection to update
* @param relationships - Relationships collection to update
*/
function extractJavaClasses(content, classes, relationships) {
// Match class declarations
const classRegex = /class\s+(\w+)(?:\s+extends\s+(\w+))?(?:\s+implements\s+([^{]+))?/g;
let match;
while ((match = classRegex.exec(content)) !== null) {
const className = match[1];
const extendsClass = match[2];
const implementsList = match[3]?.split(',').map(i => i.trim());
// Initialize class if not exists
if (!classes[className]) {
classes[className] = {
properties: [],
methods: [],
};
}
if (extendsClass) {
classes[className].extends = extendsClass;
// Add inheritance relationship
relationships.push({
from: className,
to: extendsClass,
type: 'extends',
});
}
if (implementsList && implementsList.length > 0) {
classes[className].implements = implementsList;
// Add implementation relationships
implementsList.forEach(impl => {
relationships.push({
from: className,
to: impl,
type: 'implements',
});
});
}
// Extract class body (simplified)
const classBodyMatch = content.substring(match.index).match(/{([^{}]*(?:{[^{}]*}[^{}]*)*)}/);
if (classBodyMatch) {
const classBody = classBodyMatch[1];
// Extract properties (simplified)
const propertyRegex = /(?:private|public|protected)\s+(?:static\s+)?(?:final\s+)?(?:\w+)(?:<[^>]+>)?\s+(\w+)\s*[;=]/g;
let propMatch;
while ((propMatch = propertyRegex.exec(classBody)) !== null) {
classes[className].properties.push(propMatch[1]);
}
// Extract methods (simplified)
const methodRegex = /(?:private|public|protected)\s+(?:static\s+)?(?:\w+)(?:<[^>]+>)?\s+(\w+)\s*\([^)]*\)/g;
let methodMatch;
while ((methodMatch = methodRegex.exec(classBody)) !== null) {
classes[className].methods.push(methodMatch[1]);
}
}
}
}
//# sourceMappingURL=diagramGenerator.js.map