Skip to content

Commit 1aeea34

Browse files
carlos-almcarlos-alm
authored
Merge pull request #97 from optave/feat/mermaid-subgraphs-styling
feat: enhance Mermaid export with subgraphs, edge labels, node shapes and styling
2 parents 8f31573 + 0c10e23 commit 1aeea34

4 files changed

Lines changed: 287 additions & 18 deletions

File tree

‎src/cli.js‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,13 +280,15 @@ program
280280
.option('-T, --no-tests', 'Exclude test/spec files')
281281
.option('--include-tests', 'Include test/spec files (overrides excludeTests config)')
282282
.option('--min-confidence <score>', 'Minimum edge confidence threshold (default: 0.5)', '0.5')
283+
.option('--direction <dir>', 'Flowchart direction for Mermaid: TB, LR, RL, BT', 'LR')
283284
.option('-o, --output <file>', 'Write to file instead of stdout')
284285
.action((opts) => {
285286
const db = openReadonlyOrFail(opts.db);
286287
const exportOpts = {
287288
fileLevel: !opts.functions,
288289
noTests: resolveNoTests(opts),
289290
minConfidence: parseFloat(opts.minConfidence),
291+
direction: opts.direction,
290292
};
291293

292294
let output;

‎src/export.js‎

Lines changed: 158 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -125,19 +125,63 @@ export function exportDOT(db, opts = {}) {
125125
return lines.join('\n');
126126
}
127127

128+
/** Escape double quotes for Mermaid labels. */
129+
function escapeLabel(label) {
130+
return label.replace(/"/g, '#quot;');
131+
}
132+
133+
/** Map node kind to Mermaid shape wrapper. */
134+
function mermaidShape(kind, label) {
135+
const escaped = escapeLabel(label);
136+
switch (kind) {
137+
case 'function':
138+
case 'method':
139+
return `(["${escaped}"])`;
140+
case 'class':
141+
case 'interface':
142+
case 'type':
143+
case 'struct':
144+
case 'enum':
145+
case 'trait':
146+
case 'record':
147+
return `{{"${escaped}"}}`;
148+
case 'module':
149+
return `[["${escaped}"]]`;
150+
default:
151+
return `["${escaped}"]`;
152+
}
153+
}
154+
155+
/** Map node role to Mermaid style colors. */
156+
const ROLE_STYLES = {
157+
entry: 'fill:#e8f5e9,stroke:#4caf50',
158+
core: 'fill:#e3f2fd,stroke:#2196f3',
159+
utility: 'fill:#f5f5f5,stroke:#9e9e9e',
160+
dead: 'fill:#ffebee,stroke:#f44336',
161+
leaf: 'fill:#fffde7,stroke:#fdd835',
162+
};
163+
128164
/**
129165
* Export the dependency graph in Mermaid format.
130166
*/
131167
export function exportMermaid(db, opts = {}) {
132168
const fileLevel = opts.fileLevel !== false;
133169
const noTests = opts.noTests || false;
134170
const minConf = opts.minConfidence ?? DEFAULT_MIN_CONFIDENCE;
135-
const lines = ['graph LR'];
171+
const direction = opts.direction || 'LR';
172+
const lines = [`flowchart ${direction}`];
173+
174+
let nodeCounter = 0;
175+
const nodeIdMap = new Map();
176+
function nodeId(key) {
177+
if (!nodeIdMap.has(key)) nodeIdMap.set(key, `n${nodeCounter++}`);
178+
return nodeIdMap.get(key);
179+
}
136180

137181
if (fileLevel) {
138182
let edges = db
139183
.prepare(`
140-
SELECT DISTINCT n1.file AS source, n2.file AS target
184+
SELECT DISTINCT n1.file AS source, n2.file AS target, e.kind AS edge_kind
141185
FROM edges e
142186
JOIN nodes n1 ON e.source_id = n1.id
143187
JOIN nodes n2 ON e.target_id = n2.id
@@ -147,32 +191,133 @@ export function exportMermaid(db, opts = {}) {
147191
.all(minConf);
148192
if (noTests) edges = edges.filter((e) => !isTestFile(e.source) && !isTestFile(e.target));
149193

194+
// Collect all files referenced in edges
195+
const allFiles = new Set();
150196
for (const { source, target } of edges) {
151-
const s = source.replace(/[^a-zA-Z0-9]/g, '_');
152-
const t = target.replace(/[^a-zA-Z0-9]/g, '_');
153-
lines.push(` ${s}["${source}"] --> ${t}["${target}"]`);
197+
allFiles.add(source);
198+
allFiles.add(target);
199+
}
200+
201+
// Build directory groupings — try DB directory nodes first, fall back to path.dirname()
202+
const dirs = new Map();
203+
const hasDirectoryNodes =
204+
db.prepare("SELECT COUNT(*) as c FROM nodes WHERE kind = 'directory'").get().c > 0;
205+
206+
if (hasDirectoryNodes) {
207+
const dbDirs = db.prepare("SELECT id, name FROM nodes WHERE kind = 'directory'").all();
208+
for (const d of dbDirs) {
209+
const containedFiles = db
210+
.prepare(`
211+
SELECT n.name FROM edges e
212+
JOIN nodes n ON e.target_id = n.id
213+
WHERE e.source_id = ? AND e.kind = 'contains' AND n.kind = 'file'
214+
`)
215+
.all(d.id)
216+
.map((r) => r.name)
217+
.filter((f) => allFiles.has(f));
218+
if (containedFiles.length > 0) dirs.set(d.name, containedFiles);
219+
}
220+
} else {
221+
for (const file of allFiles) {
222+
const dir = path.dirname(file) || '.';
223+
if (!dirs.has(dir)) dirs.set(dir, []);
224+
dirs.get(dir).push(file);
225+
}
226+
}
227+
228+
// Emit subgraphs
229+
for (const [dir, files] of [...dirs].sort((a, b) => a[0].localeCompare(b[0]))) {
230+
const sgId = dir.replace(/[^a-zA-Z0-9]/g, '_');
231+
lines.push(` subgraph ${sgId}["${escapeLabel(dir)}"]`);
232+
for (const f of files) {
233+
const nId = nodeId(f);
234+
lines.push(` ${nId}["${escapeLabel(path.basename(f))}"]`);
235+
}
236+
lines.push(' end');
237+
}
238+
239+
// Deduplicate edges per source-target pair, collecting all distinct kinds
240+
const edgeMap = new Map();
241+
for (const { source, target, edge_kind } of edges) {
242+
const key = `${source}|${target}`;
243+
const label = edge_kind === 'imports-type' ? 'imports' : edge_kind;
244+
if (!edgeMap.has(key)) edgeMap.set(key, { source, target, labels: new Set() });
245+
edgeMap.get(key).labels.add(label);
246+
}
247+
248+
for (const { source, target, labels } of edgeMap.values()) {
249+
lines.push(` ${nodeId(source)} -->|${[...labels].join(', ')}| ${nodeId(target)}`);
154250
}
155251
} else {
156252
let edges = db
157253
.prepare(`
158-
SELECT n1.name AS source_name, n1.file AS source_file,
159-
n2.name AS target_name, n2.file AS target_file
254+
SELECT n1.name AS source_name, n1.kind AS source_kind, n1.file AS source_file,
255+
n2.name AS target_name, n2.kind AS target_kind, n2.file AS target_file,
256+
e.kind AS edge_kind
160257
FROM edges e
161258
JOIN nodes n1 ON e.source_id = n1.id
162259
JOIN nodes n2 ON e.target_id = n2.id
163-
WHERE n1.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module') AND n2.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module')
164-
AND e.kind = 'calls'
165-
AND e.confidence >= ?
260+
WHERE n1.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module')
261+
AND n2.kind IN ('function', 'method', 'class', 'interface', 'type', 'struct', 'enum', 'trait', 'record', 'module')
262+
AND e.kind = 'calls'
263+
AND e.confidence >= ?
166264
`)
167265
.all(minConf);
168266
if (noTests)
169267
edges = edges.filter((e) => !isTestFile(e.source_file) && !isTestFile(e.target_file));
170268

269+
// Group nodes by file for subgraphs
270+
const fileNodes = new Map();
271+
const nodeKinds = new Map();
272+
for (const e of edges) {
273+
const sKey = `${e.source_file}::${e.source_name}`;
274+
const tKey = `${e.target_file}::${e.target_name}`;
275+
nodeId(sKey);
276+
nodeId(tKey);
277+
nodeKinds.set(sKey, e.source_kind);
278+
nodeKinds.set(tKey, e.target_kind);
279+
280+
if (!fileNodes.has(e.source_file)) fileNodes.set(e.source_file, new Map());
281+
fileNodes.get(e.source_file).set(sKey, e.source_name);
282+
283+
if (!fileNodes.has(e.target_file)) fileNodes.set(e.target_file, new Map());
284+
fileNodes.get(e.target_file).set(tKey, e.target_name);
285+
}
286+
287+
// Emit subgraphs grouped by file
288+
for (const [file, nodes] of [...fileNodes].sort((a, b) => a[0].localeCompare(b[0]))) {
289+
const sgId = file.replace(/[^a-zA-Z0-9]/g, '_');
290+
lines.push(` subgraph ${sgId}["${escapeLabel(file)}"]`);
291+
for (const [key, name] of nodes) {
292+
const kind = nodeKinds.get(key);
293+
lines.push(` ${nodeId(key)}${mermaidShape(kind, name)}`);
294+
}
295+
lines.push(' end');
296+
}
297+
298+
// Emit edges with labels
171299
for (const e of edges) {
172-
const sId = `${e.source_file}_${e.source_name}`.replace(/[^a-zA-Z0-9]/g, '_');
173-
const tId = `${e.target_file}_${e.target_name}`.replace(/[^a-zA-Z0-9]/g, '_');
174-
lines.push(` ${sId}["${e.source_name}"] --> ${tId}["${e.target_name}"]`);
300+
const sId = nodeId(`${e.source_file}::${e.source_name}`);
301+
const tId = nodeId(`${e.target_file}::${e.target_name}`);
302+
lines.push(` ${sId} -->|${e.edge_kind}| ${tId}`);
303+
}
304+
305+
// Role styling — query roles for all referenced nodes
306+
const allKeys = [...nodeIdMap.keys()];
307+
const roleStyles = [];
308+
for (const key of allKeys) {
309+
const colonIdx = key.indexOf('::');
310+
if (colonIdx === -1) continue;
311+
const file = key.slice(0, colonIdx);
312+
const name = key.slice(colonIdx + 2);
313+
const row = db
314+
.prepare('SELECT role FROM nodes WHERE file = ? AND name = ? AND role IS NOT NULL LIMIT 1')
315+
.get(file, name);
316+
if (row?.role && ROLE_STYLES[row.role]) {
317+
roleStyles.push(` style ${nodeIdMap.get(key)} ${ROLE_STYLES[row.role]}`);
318+
}
175319
}
320+
lines.push(...roleStyles);
176321
}
177322

178323
return lines.join('\n');

‎tests/graph/export.test.js‎

Lines changed: 125 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,17 +43,65 @@ describe('exportDOT', () => {
4343
});
4444

4545
describe('exportMermaid', () => {
46-
it('generates valid Mermaid syntax', () => {
46+
it('generates valid Mermaid syntax with flowchart LR default', () => {
4747
const db = createTestDb();
4848
const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
4949
const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
5050
insertEdge(db, a, b, 'imports');
5151

5252
const mermaid = exportMermaid(db);
53-
expect(mermaid).toContain('graph LR');
53+
expect(mermaid).toContain('flowchart LR');
5454
expect(mermaid).toContain('-->');
5555
db.close();
5656
});
57+
58+
it('uses custom direction option', () => {
59+
const db = createTestDb();
60+
const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
61+
const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
62+
insertEdge(db, a, b, 'imports');
63+
64+
const mermaid = exportMermaid(db, { direction: 'TB' });
65+
expect(mermaid).toContain('flowchart TB');
66+
db.close();
67+
});
68+
69+
it('groups files into directory subgraphs', () => {
70+
const db = createTestDb();
71+
const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
72+
const b = insertNode(db, 'lib/b.js', 'file', 'lib/b.js', 0);
73+
insertEdge(db, a, b, 'imports');
74+
75+
const mermaid = exportMermaid(db);
76+
expect(mermaid).toContain('subgraph');
77+
expect(mermaid).toContain('"src"');
78+
expect(mermaid).toContain('"lib"');
79+
expect(mermaid).toContain('end');
80+
db.close();
81+
});
82+
83+
it('adds edge labels from edge kind', () => {
84+
const db = createTestDb();
85+
const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
86+
const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
87+
insertEdge(db, a, b, 'imports');
88+
89+
const mermaid = exportMermaid(db);
90+
expect(mermaid).toContain('-->|imports|');
91+
db.close();
92+
});
93+
94+
it('collapses imports-type to imports label', () => {
95+
const db = createTestDb();
96+
const a = insertNode(db, 'src/a.js', 'file', 'src/a.js', 0);
97+
const b = insertNode(db, 'src/b.js', 'file', 'src/b.js', 0);
98+
insertEdge(db, a, b, 'imports-type');
99+
100+
const mermaid = exportMermaid(db);
101+
expect(mermaid).toContain('-->|imports|');
102+
expect(mermaid).not.toContain('imports-type');
103+
db.close();
104+
});
57105
});
58106

59107
describe('exportDOT — function-level', () => {
@@ -107,12 +155,86 @@ describe('exportMermaid — function-level', () => {
107155
insertEdge(db, fnA, fnB, 'calls');
108156

109157
const mermaid = exportMermaid(db, { fileLevel: false });
110-
expect(mermaid).toContain('graph LR');
158+
expect(mermaid).toContain('flowchart LR');
111159
expect(mermaid).toContain('doWork');
112160
expect(mermaid).toContain('helper');
113161
expect(mermaid).toContain('-->');
114162
db.close();
115163
});
164+
165+
it('uses stadium shape for functions', () => {
166+
const db = createTestDb();
167+
const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
168+
const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
169+
insertEdge(db, fnA, fnB, 'calls');
170+
171+
const mermaid = exportMermaid(db, { fileLevel: false });
172+
expect(mermaid).toContain('(["doWork"])');
173+
expect(mermaid).toContain('(["helper"])');
174+
db.close();
175+
});
176+
177+
it('uses hexagon shape for classes', () => {
178+
const db = createTestDb();
179+
const cls = insertNode(db, 'MyClass', 'class', 'src/a.js', 5);
180+
const fn = insertNode(db, 'helper', 'function', 'src/b.js', 10);
181+
insertEdge(db, cls, fn, 'calls');
182+
183+
const mermaid = exportMermaid(db, { fileLevel: false });
184+
expect(mermaid).toContain('{{"MyClass"}}');
185+
db.close();
186+
});
187+
188+
it('uses subroutine shape for modules', () => {
189+
const db = createTestDb();
190+
const mod = insertNode(db, 'MyModule', 'module', 'src/a.js', 5);
191+
const fn = insertNode(db, 'helper', 'function', 'src/b.js', 10);
192+
insertEdge(db, mod, fn, 'calls');
193+
194+
const mermaid = exportMermaid(db, { fileLevel: false });
195+
expect(mermaid).toContain('[["MyModule"]]');
196+
db.close();
197+
});
198+
199+
it('adds edge labels for calls', () => {
200+
const db = createTestDb();
201+
const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
202+
const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
203+
insertEdge(db, fnA, fnB, 'calls');
204+
205+
const mermaid = exportMermaid(db, { fileLevel: false });
206+
expect(mermaid).toContain('-->|calls|');
207+
db.close();
208+
});
209+
210+
it('groups functions by file into subgraphs', () => {
211+
const db = createTestDb();
212+
const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
213+
const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
214+
insertEdge(db, fnA, fnB, 'calls');
215+
216+
const mermaid = exportMermaid(db, { fileLevel: false });
217+
expect(mermaid).toContain('subgraph');
218+
expect(mermaid).toContain('"src/a.js"');
219+
expect(mermaid).toContain('"src/b.js"');
220+
expect(mermaid).toContain('end');
221+
db.close();
222+
});
223+
224+
it('applies role styling', () => {
225+
const db = createTestDb();
226+
const fnA = insertNode(db, 'doWork', 'function', 'src/a.js', 5);
227+
const fnB = insertNode(db, 'helper', 'function', 'src/b.js', 10);
228+
// Add role to the nodes
229+
db.prepare('UPDATE nodes SET role = ? WHERE id = ?').run('entry', fnA);
230+
db.prepare('UPDATE nodes SET role = ? WHERE id = ?').run('utility', fnB);
231+
insertEdge(db, fnA, fnB, 'calls');
232+
233+
const mermaid = exportMermaid(db, { fileLevel: false });
234+
expect(mermaid).toContain('fill:#e8f5e9,stroke:#4caf50');
235+
expect(mermaid).toContain('fill:#f5f5f5,stroke:#9e9e9e');
236+
db.close();
237+
});
116238
});
117239

118240
describe('exportJSON', () => {

‎tests/integration/cli.test.js‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -130,9 +130,9 @@ describe('CLI smoke tests', () => {
130130
});
131131

132132
// ─── Export (Mermaid) ────────────────────────────────────────────────
133-
test('export -f mermaid outputs graph LR', () => {
133+
test('export -f mermaid outputs flowchart LR', () => {
134134
const out = run('export', '--db', dbPath, '-f', 'mermaid');
135-
expect(out).toContain('graph LR');
135+
expect(out).toContain('flowchart LR');
136136
});
137137

138138
// ─── Export (JSON) ───────────────────────────────────────────────────

0 commit comments

Comments
 (0)