From 3e07a5acff05d93d0dbd0001a88bd3540d4b87d3 Mon Sep 17 00:00:00 2001
From: Knut Sveidqvist <knsv@sveido.com>
Date: Mon, 12 May 2025 17:46:19 +0200
Subject: [PATCH] Added support for valueFormat directive

---
 packages/mermaid/src/diagrams/treemap/db.ts   |  4 +--
 .../mermaid/src/diagrams/treemap/parser.ts    | 25 +++++++++++++------
 .../mermaid/src/diagrams/treemap/types.ts     | 24 ++++++++++++++++++
 3 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/packages/mermaid/src/diagrams/treemap/db.ts b/packages/mermaid/src/diagrams/treemap/db.ts
index ce76ab2dd..550ad6476 100644
--- a/packages/mermaid/src/diagrams/treemap/db.ts
+++ b/packages/mermaid/src/diagrams/treemap/db.ts
@@ -78,8 +78,8 @@ const getClasses = (): Map<string, DiagramStyleClassDef> => {
   return classes;
 };
 
-const getStylesForClass = (classSelector: string) => {
-  return classes.get(classSelector)?.styles;
+const getStylesForClass = (classSelector: string): string[] => {
+  return classes.get(classSelector)?.styles ?? [];
 };
 
 const clear = () => {
diff --git a/packages/mermaid/src/diagrams/treemap/parser.ts b/packages/mermaid/src/diagrams/treemap/parser.ts
index b7bab2e8b..34c218d34 100644
--- a/packages/mermaid/src/diagrams/treemap/parser.ts
+++ b/packages/mermaid/src/diagrams/treemap/parser.ts
@@ -3,27 +3,34 @@ import type { ParserDefinition } from '../../diagram-api/types.js';
 import { log } from '../../logger.js';
 import { populateCommonDb } from '../common/populateCommonDb.js';
 import { db } from './db.js';
-import type { TreemapNode } from './types.js';
+import type { TreemapNode, TreemapAst } from './types.js';
 import { buildHierarchy } from './utils.js';
 
 /**
  * Populates the database with data from the Treemap AST
  * @param ast - The Treemap AST
  */
-const populate = (ast: any) => {
+const populate = (ast: TreemapAst) => {
   populateCommonDb(ast, db);
 
-  const items = [];
+  const items: {
+    level: number;
+    name: string;
+    type: string;
+    value?: number;
+    classSelector?: string;
+    cssCompiledStyles?: string;
+  }[] = [];
 
   // Extract classes and styles from the treemap
-  for (const row of ast.TreemapRows || []) {
+  for (const row of ast.TreemapRows ?? []) {
     if (row.$type === 'ClassDefStatement') {
-      db.addClass(row.className, row.styleText);
+      db.addClass(row.className ?? '', row.styleText ?? '');
     }
   }
 
   // Extract data from each row in the treemap
-  for (const row of ast.TreemapRows || []) {
+  for (const row of ast.TreemapRows ?? []) {
     const item = row.item;
 
     if (!item) {
@@ -33,13 +40,17 @@ const populate = (ast: any) => {
     const level = row.indent ? parseInt(row.indent) : 0;
     const name = getItemName(item);
 
+    // Get styles as a string if they exist
+    const styles = item.classSelector ? db.getStylesForClass(item.classSelector) : [];
+    const cssCompiledStyles = styles.length > 0 ? styles.join(';') : undefined;
+
     const itemData = {
       level,
       name,
       type: item.$type,
       value: item.value,
       classSelector: item.classSelector,
-      cssCompiledStyles: item.classSelector ? db.getStylesForClass(item.classSelector) : undefined,
+      cssCompiledStyles,
     };
 
     items.push(itemData);
diff --git a/packages/mermaid/src/diagrams/treemap/types.ts b/packages/mermaid/src/diagrams/treemap/types.ts
index 161fd0ec1..d15915b80 100644
--- a/packages/mermaid/src/diagrams/treemap/types.ts
+++ b/packages/mermaid/src/diagrams/treemap/types.ts
@@ -40,6 +40,30 @@ export interface TreemapData {
   root?: TreemapNode;
 }
 
+export interface TreemapItem {
+  $type: string;
+  name: string;
+  value?: number;
+  classSelector?: string;
+}
+
+export interface TreemapRow {
+  $type: string;
+  indent?: string;
+  item?: TreemapItem;
+  className?: string;
+  styleText?: string;
+}
+
+export interface TreemapAst {
+  TreemapRows?: TreemapRow[];
+  title?: string;
+  description?: string;
+  accDescription?: string;
+  accTitle?: string;
+  diagramTitle?: string;
+}
+
 // Define the TreemapDiagramConfig interface
 export interface TreemapDiagramConfig extends BaseDiagramConfig {
   padding?: number;