From 3782b69744e70a72a5868535b28b81f6fd8152fe Mon Sep 17 00:00:00 2001
From: yari-dewalt <yari.dewalt@gmail.com>
Date: Fri, 13 Sep 2024 12:55:10 -0700
Subject: [PATCH] Add support for classDef statement

---
 .../class/classDiagram-styles.spec.js         | 42 +++++++++++++++++++
 .../diagrams/class/parser/classDiagram.jison  | 11 +++++
 2 files changed, 53 insertions(+)

diff --git a/packages/mermaid/src/diagrams/class/classDiagram-styles.spec.js b/packages/mermaid/src/diagrams/class/classDiagram-styles.spec.js
index 957188401..04529302f 100644
--- a/packages/mermaid/src/diagrams/class/classDiagram-styles.spec.js
+++ b/packages/mermaid/src/diagrams/class/classDiagram-styles.spec.js
@@ -69,5 +69,47 @@ describe('class diagram, ', function () {
       expect(styleElements[1]).toBe('stroke:#333');
       expect(styleElements[2]).toBe('stroke-width:4px');
     });
+    it('should be possible to define and assign a class inside the diagram', function () {
+      const str =
+        'classDiagram\n' + 'class Class01\n cssClass "Class01" pink\n classDef pink fill:#f9f';
+
+      parser.parse(str);
+
+      expect(parser.yy.getClass('Class01').cssClasses[0]).toBe('pink');
+    });
+    it('should be possible to define and assign a class using shorthand inside the diagram', function () {
+      const str = 'classDiagram\n' + 'class Class01:::pink\n classDef pink fill:#f9f';
+
+      parser.parse(str);
+
+      expect(parser.yy.getClass('Class01').cssClasses[0]).toBe('pink');
+    });
+    it('should properly assign styles from a class defined inside the diagram', function () {
+      const str =
+        'classDiagram\n' +
+        'class Class01:::pink\n classDef pink fill:#f9f,stroke:#333,stroke-width:6px';
+
+      parser.parse(str);
+
+      expect(parser.yy.getClass('Class01').styles).toStrictEqual([
+        'fill:#f9f',
+        'stroke:#333',
+        'stroke-width:6px',
+      ]);
+    });
+    it('should properly assign multiple classes and styles from classes defined inside the diagram', function () {
+      const str =
+        'classDiagram\n' +
+        'class Class01:::pink\n cssClass "Class01" bold\n classDef pink fill:#f9f\n classDef bold stroke:#333,stroke-width:6px';
+
+      parser.parse(str);
+
+      expect(parser.yy.getClass('Class01').styles).toStrictEqual([
+        'fill:#f9f',
+        'stroke:#333',
+        'stroke-width:6px',
+      ]);
+      expect(parser.yy.getClass('Class01').cssClasses).toStrictEqual(['pink', 'bold']);
+    });
   });
 });
diff --git a/packages/mermaid/src/diagrams/class/parser/classDiagram.jison b/packages/mermaid/src/diagrams/class/parser/classDiagram.jison
index ed4e72027..282bd9a41 100644
--- a/packages/mermaid/src/diagrams/class/parser/classDiagram.jison
+++ b/packages/mermaid/src/diagrams/class/parser/classDiagram.jison
@@ -61,6 +61,7 @@ Function arguments are optional: 'call <callback_name>()' simply executes 'callb
 <string>[^"]*                   return "STR";
 <*>["]                          this.begin("string");
 "style"                         return 'STYLE';
+"classDef"                      return 'CLASSDEF';
 
 <INITIAL,namespace>"namespace"  { this.begin('namespace'); return 'NAMESPACE'; }
 <namespace>\s*(\r?\n)+          { this.popState(); return 'NEWLINE'; }
@@ -263,6 +264,7 @@ statement
     | styleStatement
     | cssClassStatement
     | noteStatement
+    | classDefStatement
     | direction
     | acc_title acc_title_value  { $$=$2.trim();yy.setAccTitle($$); }
     | acc_descr acc_descr_value  { $$=$2.trim();yy.setAccDescription($$); }
@@ -324,6 +326,15 @@ noteStatement
     | NOTE noteText                { yy.addNote($2); }
     ;
 
+classDefStatement
+    : CLASSDEF classList stylesOpt {$$ = $CLASSDEF;yy.defineClass($classList,$stylesOpt);}
+    ;
+
+classList
+    : ALPHA { $$ = [$ALPHA]; }
+    | classList COMMA ALPHA = { $$ = $classList.concat([$ALPHA]); }
+    ;
+
 direction
     : direction_tb
     { yy.setDirection('TB');}