From 49103ea654745f77a11b034ffc13315d8c7afbd6 Mon Sep 17 00:00:00 2001
From: kairi003 <kairi.satellite@gmail.com>
Date: Sun, 26 Oct 2025 03:01:00 +0900
Subject: [PATCH] fix: update ClassDB and ClassTypes for improved type safety
 and consistency

---
 .../mermaid/src/diagrams/class/classDb.ts     | 30 +++++++++++--------
 .../mermaid/src/diagrams/class/classTypes.ts  |  2 +-
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/packages/mermaid/src/diagrams/class/classDb.ts b/packages/mermaid/src/diagrams/class/classDb.ts
index 152c3e657..7c7d5878c 100644
--- a/packages/mermaid/src/diagrams/class/classDb.ts
+++ b/packages/mermaid/src/diagrams/class/classDb.ts
@@ -17,6 +17,7 @@ import type {
   ClassRelation,
   ClassNode,
   ClassNote,
+  ClassNoteMap,
   ClassMap,
   NamespaceMap,
   NamespaceNode,
@@ -33,9 +34,9 @@ const sanitizeText = (txt: string) => common.sanitizeText(txt, getConfig());
 
 export class ClassDB implements DiagramDB {
   private relations: ClassRelation[] = [];
-  private classes = new Map<string, ClassNode>();
+  private classes: ClassMap = new Map<string, ClassNode>();
   private readonly styleClasses = new Map<string, StyleClass>();
-  private notes = new Map<string, ClassNote>();
+  private notes: ClassNoteMap = new Map<string, ClassNote>();
   private interfaces: Interface[] = [];
   // private static classCounter = 0;
   private namespaces = new Map<string, NamespaceNode>();
@@ -125,7 +126,7 @@ export class ClassDB implements DiagramDB {
       annotations: [],
       styles: [],
       domId: MERMAID_DOM_ID_PREFIX + name + '-' + classCounter,
-    } as ClassNode);
+    });
 
     classCounter++;
   }
@@ -184,7 +185,7 @@ export class ClassDB implements DiagramDB {
     return this.notes.get(key)!;
   }
 
-  public getNotes() {
+  public getNotes(): ClassNoteMap {
     return this.notes;
   }
 
@@ -297,7 +298,7 @@ export class ClassDB implements DiagramDB {
     return note.id;
   }
 
-  public cleanupLabel(label: string) {
+  public cleanupLabel(label: string): string {
     if (label.startsWith(':')) {
       label = label.substring(1);
     }
@@ -363,7 +364,7 @@ export class ClassDB implements DiagramDB {
     });
   }
 
-  public getTooltip(id: string, namespace?: string) {
+  public getTooltip(id: string, namespace?: string): string | undefined {
     if (namespace && this.namespaces.has(namespace)) {
       return this.namespaces.get(namespace)!.classes.get(id)!.tooltip;
     }
@@ -543,10 +544,11 @@ export class ClassDB implements DiagramDB {
 
     this.namespaces.set(id, {
       id: id,
-      classes: new Map(),
-      children: {},
+      classes: new Map<string, ClassNode>(),
+      notes: new Map<string, ClassNote>(),
+      children: new Map<string, NamespaceNode>(),
       domId: MERMAID_DOM_ID_PREFIX + id + '-' + this.namespaceCounter,
-    } as NamespaceNode);
+    });
 
     this.namespaceCounter++;
   }
@@ -648,9 +650,13 @@ export class ClassDB implements DiagramDB {
     }
 
     for (const classNode of this.classes.values()) {
-      const node = classNode as unknown as Node;
-      node.parentId = classNode.parent;
-      node.look = config.look;
+      const node: Node = {
+        ...classNode,
+        type: undefined,
+        isGroup: false,
+        parentId: classNode.parent,
+        look: classNode.look,
+      };
       nodes.push(node);
     }
 
diff --git a/packages/mermaid/src/diagrams/class/classTypes.ts b/packages/mermaid/src/diagrams/class/classTypes.ts
index 1f406e80f..e3351bb18 100644
--- a/packages/mermaid/src/diagrams/class/classTypes.ts
+++ b/packages/mermaid/src/diagrams/class/classTypes.ts
@@ -5,7 +5,7 @@ export interface ClassNode {
   id: string;
   type: string;
   label: string;
-  shape: string;
+  shape: 'classBox';
   text: string;
   cssClasses: string;
   methods: ClassMember[];