Support edge labels for one-to-one and through relations 🩷

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"

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<string>();
   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(', ')}]`

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(
+    /\bthrough\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;

src/types.ts

@@ -9,6 +9,7 @@ export interface ParsedField {
   relationKind?: RelationKind;
   onDelete?: string;
   relatedName?: string;
+  throughModel?: string;
   lineNumber: number;
 }
 

test/edge-labels.test.js

@@ -0,0 +1,45 @@
+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.after(() => {
+  Module._load = originalLoad;
+});
+
+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', passthrough='Ignored', 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\]"/);
+});