diff --git a/.changeset/pretty-falcons-say.md b/.changeset/pretty-falcons-say.md new file mode 100644 index 000000000..f1a0a95c3 --- /dev/null +++ b/.changeset/pretty-falcons-say.md @@ -0,0 +1,5 @@ +--- +'mermaid': patch +--- + +chore: Updated TreeMapDB to use class based approach diff --git a/packages/mermaid/src/diagrams/treemap/db.ts b/packages/mermaid/src/diagrams/treemap/db.ts index 6a68857f7..23326bc85 100644 --- a/packages/mermaid/src/diagrams/treemap/db.ts +++ b/packages/mermaid/src/diagrams/treemap/db.ts @@ -1,10 +1,10 @@ -import { getConfig as commonGetConfig } from '../../config.js'; -import DEFAULT_CONFIG from '../../defaultConfig.js'; +import type { DiagramDB } from '../../diagram-api/types.js'; import type { DiagramStyleClassDef } from '../../diagram-api/types.js'; -import { isLabelStyle } from '../../rendering-util/rendering-elements/shapes/handDrawnShapeStyles.js'; - +import type { TreemapDiagramConfig, TreemapNode } from './types.js'; +import DEFAULT_CONFIG from '../../defaultConfig.js'; +import { getConfig as commonGetConfig } from '../../config.js'; import { cleanAndMerge } from '../../utils.js'; -import { ImperativeState } from '../../utils/imperativeState.js'; +import { isLabelStyle } from '../../rendering-util/rendering-elements/shapes/handDrawnShapeStyles.js'; import { clear as commonClear, getAccDescription, @@ -14,99 +14,82 @@ import { setAccTitle, setDiagramTitle, } from '../common/commonDb.js'; -import type { TreemapDB, TreemapData, TreemapDiagramConfig, TreemapNode } from './types.js'; +export class TreeMapDB implements DiagramDB { + private nodes: TreemapNode[] = []; + private levels: Map = new Map(); + private outerNodes: TreemapNode[] = []; + private classes: Map = new Map(); + private root?: TreemapNode; -const defaultTreemapData: TreemapData = { - nodes: [], - levels: new Map(), - outerNodes: [], - classes: new Map(), -}; - -const state = new ImperativeState(() => structuredClone(defaultTreemapData)); - -const getConfig = (): Required => { - // Use type assertion with unknown as intermediate step - const defaultConfig = DEFAULT_CONFIG as unknown as { treemap: Required }; - const userConfig = commonGetConfig() as unknown as { treemap?: Partial }; - - return cleanAndMerge({ - ...defaultConfig.treemap, - ...(userConfig.treemap ?? {}), - }) as Required; -}; - -const getNodes = (): TreemapNode[] => state.records.nodes; - -const addNode = (node: TreemapNode, level: number) => { - const data = state.records; - data.nodes.push(node); - data.levels.set(node, level); - - if (level === 0) { - data.outerNodes.push(node); + public getNodes() { + return this.nodes; } - // Set the root node if this is a level 0 node and we don't have a root yet - if (level === 0 && !data.root) { - data.root = node; + public getConfig() { + const defaultConfig = DEFAULT_CONFIG as unknown as { treemap: Required }; + const userConfig = commonGetConfig() as unknown as { treemap?: Partial }; + return cleanAndMerge({ + ...defaultConfig.treemap, + ...(userConfig.treemap ?? {}), + }) as Required; } -}; -const getRoot = (): TreemapNode | undefined => ({ name: '', children: state.records.outerNodes }); + public addNode(node: TreemapNode, level: number) { + this.nodes.push(node); + this.levels.set(node, level); + if (level === 0) { + this.outerNodes.push(node); + this.root ??= node; + } + } -const addClass = (id: string, _style: string) => { - const classes = state.records.classes; - const styleClass = classes.get(id) ?? { id, styles: [], textStyles: [] }; - classes.set(id, styleClass); + public getRoot() { + return { name: '', children: this.outerNodes }; + } - const styles = _style.replace(/\\,/g, '§§§').replace(/,/g, ';').replace(/§§§/g, ',').split(';'); - - if (styles) { - styles.forEach((s) => { - if (isLabelStyle(s)) { - if (styleClass?.textStyles) { - styleClass.textStyles.push(s); - } else { - styleClass.textStyles = [s]; + public addClass(id: string, _style: string) { + const styleClass = this.classes.get(id) ?? { id, styles: [], textStyles: [] }; + const styles = _style.replace(/\\,/g, '§§§').replace(/,/g, ';').replace(/§§§/g, ',').split(';'); + if (styles) { + styles.forEach((s) => { + if (isLabelStyle(s)) { + if (styleClass?.textStyles) { + styleClass.textStyles.push(s); + } else { + styleClass.textStyles = [s]; + } } - } - if (styleClass?.styles) { - styleClass.styles.push(s); - } else { - styleClass.styles = [s]; - } - }); + if (styleClass?.styles) { + styleClass.styles.push(s); + } else { + styleClass.styles = [s]; + } + }); + } + this.classes.set(id, styleClass); } - classes.set(id, styleClass); -}; -const getClasses = (): Map => { - return state.records.classes; -}; + public getClasses() { + return this.classes; + } -const getStylesForClass = (classSelector: string): string[] => { - return state.records.classes.get(classSelector)?.styles ?? []; -}; + public getStylesForClass(classSelector: string): string[] { + return this.classes.get(classSelector)?.styles ?? []; + } -const clear = () => { - commonClear(); - state.reset(); -}; + public clear() { + commonClear(); + this.nodes = []; + this.levels = new Map(); + this.outerNodes = []; + this.classes = new Map(); + this.root = undefined; + } -export const db: TreemapDB = { - getNodes, - addNode, - getRoot, - getConfig, - clear, - setAccTitle, - getAccTitle, - setDiagramTitle, - getDiagramTitle, - getAccDescription, - setAccDescription, - addClass, - getClasses, - getStylesForClass, -}; + public setAccTitle = setAccTitle; + public getAccTitle = getAccTitle; + public setDiagramTitle = setDiagramTitle; + public getDiagramTitle = getDiagramTitle; + public getAccDescription = getAccDescription; + public setAccDescription = setAccDescription; +} diff --git a/packages/mermaid/src/diagrams/treemap/diagram.ts b/packages/mermaid/src/diagrams/treemap/diagram.ts index dd599174e..2f8ff92f3 100644 --- a/packages/mermaid/src/diagrams/treemap/diagram.ts +++ b/packages/mermaid/src/diagrams/treemap/diagram.ts @@ -1,12 +1,14 @@ import type { DiagramDefinition } from '../../diagram-api/types.js'; -import { db } from './db.js'; +import { TreeMapDB } from './db.js'; import { parser } from './parser.js'; import { renderer } from './renderer.js'; import styles from './styles.js'; export const diagram: DiagramDefinition = { parser, - db, + get db() { + return new TreeMapDB(); + }, renderer, styles, }; diff --git a/packages/mermaid/src/diagrams/treemap/parser.ts b/packages/mermaid/src/diagrams/treemap/parser.ts index 82efb5911..4d71ff470 100644 --- a/packages/mermaid/src/diagrams/treemap/parser.ts +++ b/packages/mermaid/src/diagrams/treemap/parser.ts @@ -2,15 +2,15 @@ import { parse } from '@mermaid-js/parser'; 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, TreemapAst } from './types.js'; +import type { TreemapNode, TreemapAst, TreemapDB } from './types.js'; import { buildHierarchy } from './utils.js'; +import { TreeMapDB } from './db.js'; /** * Populates the database with data from the Treemap AST * @param ast - The Treemap AST */ -const populate = (ast: TreemapAst) => { +const populate = (ast: TreemapAst, db: TreemapDB) => { // We need to bypass the type checking for populateCommonDb // eslint-disable-next-line @typescript-eslint/no-explicit-any populateCommonDb(ast as any, db); @@ -84,6 +84,8 @@ const getItemName = (item: { name?: string | number }): string => { }; export const parser: ParserDefinition = { + // @ts-expect-error - TreeMapDB is not assignable to DiagramDB + parser: { yy: undefined }, parse: async (text: string): Promise => { try { // Use a generic parse that accepts any diagram type @@ -91,7 +93,13 @@ export const parser: ParserDefinition = { const parseFunc = parse as (diagramType: string, text: string) => Promise; const ast = await parseFunc('treemap', text); log.debug('Treemap AST:', ast); - populate(ast); + const db = parser.parser?.yy; + if (!(db instanceof TreeMapDB)) { + throw new Error( + 'parser.parser?.yy was not a TreemapDB. This is due to a bug within Mermaid, please report this issue at https://github.com/mermaid-js/mermaid/issues.' + ); + } + populate(ast, db); } catch (error) { log.error('Error parsing treemap:', error); throw error;