From a479842262874272d44386cdf04e5be6d59a5f59 Mon Sep 17 00:00:00 2001 From: Bao Date: Mon, 13 Jul 2026 05:41:48 +0700 Subject: [PATCH 1/2] fix: label one-to-one and through relations --- package.json | 1 + src/graphWebview.ts | 3 ++- src/parser.ts | 11 +++++++++++ src/types.ts | 1 + test/edge-labels.test.js | 41 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 test/edge-labels.test.js diff --git a/package.json b/package.json index 28201b3..278c6bf 100644 --- a/package.json +++ b/package.json @@ -154,6 +154,7 @@ "scripts": { "vscode:prepublish": "npm run build", "build": "tsc -p ./", + "test": "npm run build && node --test", "watch": "tsc -watch -p ./", "package": "vsce package", "publish": "vsce publish" diff --git a/src/graphWebview.ts b/src/graphWebview.ts index 06f97b9..5e68afc 100644 --- a/src/graphWebview.ts +++ b/src/graphWebview.ts @@ -8,7 +8,7 @@ function escapeLabel(s: string): string { return s.replace(/[^A-Za-z0-9_]/g, '_'); } -function buildMermaid(index: WorkspaceIndex): string { +export function buildMermaid(index: WorkspaceIndex): string { const lines: string[] = ['erDiagram']; const modelNames = new Set(); for (const app of index.apps) { @@ -46,6 +46,7 @@ function buildMermaid(index: WorkspaceIndex): string { : '}o--||'; const parts: string[] = [f.name]; if (f.onDelete) parts.push(f.onDelete); + if (f.throughModel) parts.push(`through ${f.throughModel}`); if (f.relatedName) parts.push(`as ${f.relatedName}`); const label = parts.length > 1 ? `${parts[0]} [${parts.slice(1).join(', ')}]` diff --git a/src/parser.ts b/src/parser.ts index 0dd1f7b..a30d03c 100644 --- a/src/parser.ts +++ b/src/parser.ts @@ -109,6 +109,13 @@ function extractRelatedName(argsBlock: string): string | undefined { return m ? m[1] || m[2] : undefined; } +function extractThroughModel(argsBlock: string): string | undefined { + const m = argsBlock.match( + /through\s*=\s*(?:'([^']+)'|"([^"]+)"|([A-Za-z_][A-Za-z0-9_.]*))/ + ); + return m ? m[1] || m[2] || m[3] : undefined; +} + function readBalancedArgs(lines: string[], startIdx: number): { argsBlock: string; endIdx: number; @@ -241,6 +248,10 @@ export function parseModelsFile(filePath: string, content: string): ParsedModel[ if (onDelete) field.onDelete = onDelete; const relatedName = extractRelatedName(argsInner); if (relatedName) field.relatedName = relatedName; + if (fieldType === 'ManyToManyField') { + const throughModel = extractThroughModel(argsInner); + if (throughModel) field.throughModel = throughModel; + } } model.fields.push(field); j = endIdx + 1; diff --git a/src/types.ts b/src/types.ts index 45ce6d7..84a784c 100644 --- a/src/types.ts +++ b/src/types.ts @@ -9,6 +9,7 @@ export interface ParsedField { relationKind?: RelationKind; onDelete?: string; relatedName?: string; + throughModel?: string; lineNumber: number; } diff --git a/test/edge-labels.test.js b/test/edge-labels.test.js new file mode 100644 index 0000000..dc50743 --- /dev/null +++ b/test/edge-labels.test.js @@ -0,0 +1,41 @@ +const assert = require('node:assert/strict'); +const Module = require('node:module'); +const test = require('node:test'); + +const originalLoad = Module._load; +Module._load = function (request, parent, isMain) { + if (request === 'vscode') { + return {}; + } + return originalLoad.call(this, request, parent, isMain); +}; + +const { buildMermaid } = require('../out/graphWebview'); +const { parseModelsFile } = require('../out/parser'); + +test('labels OneToOneField metadata and ManyToManyField through models', () => { + const models = parseModelsFile('/project/app/models.py', ` +class User(models.Model): + profile = models.OneToOneField( + 'Profile', on_delete=models.CASCADE, related_name='user' + ) + +class Profile(models.Model): + pass + +class Group(models.Model): + members = models.ManyToManyField( + 'User', through='Membership', related_name='groups' + ) + +class Membership(models.Model): + pass +`); + const mermaid = buildMermaid({ + apps: [{ name: 'app', path: '/project/app', models }], + scannedAt: 0, + }); + + assert.match(mermaid, /User \|\|--\|\| Profile : "profile \[CASCADE, as user\]"/); + assert.match(mermaid, /Group \}o--o\{ User : "members \[through Membership, as groups\]"/); +}); \ No newline at end of file From f8ed09da4d0a4bc584f67cf1977ca0cbf7f5e5e5 Mon Sep 17 00:00:00 2001 From: Bao Date: Mon, 13 Jul 2026 08:02:00 +0700 Subject: [PATCH 2/2] test: harden through relation parsing --- src/parser.ts | 2 +- test/edge-labels.test.js | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/parser.ts b/src/parser.ts index a30d03c..5032831 100644 --- a/src/parser.ts +++ b/src/parser.ts @@ -111,7 +111,7 @@ function extractRelatedName(argsBlock: string): string | undefined { function extractThroughModel(argsBlock: string): string | undefined { const m = argsBlock.match( - /through\s*=\s*(?:'([^']+)'|"([^"]+)"|([A-Za-z_][A-Za-z0-9_.]*))/ + /\bthrough\s*=\s*(?:'([^']+)'|"([^"]+)"|([A-Za-z_][A-Za-z0-9_.]*))/ ); return m ? m[1] || m[2] || m[3] : undefined; } diff --git a/test/edge-labels.test.js b/test/edge-labels.test.js index dc50743..0d43e32 100644 --- a/test/edge-labels.test.js +++ b/test/edge-labels.test.js @@ -13,6 +13,10 @@ Module._load = function (request, parent, isMain) { const { buildMermaid } = require('../out/graphWebview'); const { parseModelsFile } = require('../out/parser'); +test.after(() => { + Module._load = originalLoad; +}); + test('labels OneToOneField metadata and ManyToManyField through models', () => { const models = parseModelsFile('/project/app/models.py', ` class User(models.Model): @@ -25,7 +29,7 @@ class Profile(models.Model): class Group(models.Model): members = models.ManyToManyField( - 'User', through='Membership', related_name='groups' + 'User', passthrough='Ignored', through='Membership', related_name='groups' ) class Membership(models.Model):