From 58fbfc3c38b8f79ec1c88e9b462c313b2811d121 Mon Sep 17 00:00:00 2001
From: Justin Greywolf <jgreywolf@lostware.com>
Date: Mon, 6 Jan 2020 16:21:11 -0800
Subject: [PATCH] 1119 Support method return types

Small refactor to split out logic for determining method display text and style. Updated documentation
Used regex to parse method statements in class diagrams to extract discrete elements to set display appropriately. Added tests and updated docs
---
 .../rendering/classDiagram.spec.js            | 15 +++++++++
 docs/classDiagram.md                          |  7 ++--
 src/diagrams/class/classDiagram.spec.js       | 32 ++++++++++++++++++-
 src/diagrams/class/classRenderer.js           | 30 +++++++++--------
 4 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/cypress/integration/rendering/classDiagram.spec.js b/cypress/integration/rendering/classDiagram.spec.js
index 3f022b55d..078094519 100644
--- a/cypress/integration/rendering/classDiagram.spec.js
+++ b/cypress/integration/rendering/classDiagram.spec.js
@@ -274,4 +274,19 @@ describe('Class diagram', () => {
     );
     cy.get('svg');
   });
+
+  it('11: should render a simple class diagram with return type on method', () => {
+    imgSnapshotTest(
+      `
+    classDiagram
+      class Class10~T~ {
+        int[] id
+        test(int[] ids) bool
+        testArray() bool[]
+      }
+      `,
+      {}
+    );
+    cy.get('svg');
+  });
 });
diff --git a/docs/classDiagram.md b/docs/classDiagram.md
index 814fee196..f13d60d83 100644
--- a/docs/classDiagram.md
+++ b/docs/classDiagram.md
@@ -115,7 +115,7 @@ There are two ways to define the members of a class, and regardless of whichever
   class BankAccount
   BankAccount : +String owner
   BankAccount : +BigDecimal balance
-  BankAccount : +deposit(amount) [bool]
+  BankAccount : +deposit(amount) bool
   BankAccount : +withdrawal(amount)
   ```
   ```  mermaid
@@ -132,7 +132,7 @@ There are two ways to define the members of a class, and regardless of whichever
 class BankAccount{
     +String owner
     +BigDecimal balance
-    +deposit(amount) [bool]
+    +deposit(amount) bool
     +withdrawl(amount)
 }
 ```
@@ -147,6 +147,9 @@ class BankAccount{
 ```
 
 
+#### Return Type
+Optionally you can end the method/function definition with the data type that will be returned
+
 #### Visibility
 To specify the visibility of a class member (i.e. any attribute or method), these notations may be placed before the member's name, but it is optional:
 
diff --git a/src/diagrams/class/classDiagram.spec.js b/src/diagrams/class/classDiagram.spec.js
index 89bfdc01b..8510b086d 100644
--- a/src/diagrams/class/classDiagram.spec.js
+++ b/src/diagrams/class/classDiagram.spec.js
@@ -105,7 +105,7 @@ describe('class diagram, ', function () {
       parser.parse(str);
     });
 
-    it('should handle parsing of method statements  grouped by brackets', function () {
+    it('should handle parsing of method statements grouped by brackets', function () {
       const str =
         'classDiagram\n' +
         'class Dummy_Class {\n' +
@@ -121,6 +121,36 @@ describe('class diagram, ', function () {
       parser.parse(str);
     });
 
+    it('should handle return types on methods', function () {
+      const str =
+        'classDiagram\n' +
+        'Object <|-- ArrayList\n' +
+        'Object : equals()\n' +
+        'Object : -Object[] objects\n' +
+        'Object : +getObjects() Object[]\n' +
+        'ArrayList : Dummy elementData\n' +
+        'ArrayList : getDummy() Dummy';
+
+      parser.parse(str);
+    });
+
+    it('should handle return types on methods grouped by brackets', function () {
+      const str =
+        'classDiagram\n' +
+        'class Dummy_Class {\n' +
+        'string data\n' +
+        'getDummy() Dummy\n' +
+        '}\n' +
+        '\n' +
+        'class Flight {\n' +
+        '   int flightNumber\n' +
+        '   datetime departureTime\n' +
+        '   getDepartureTime() datetime\n' +
+        '}';
+
+      parser.parse(str);
+    });
+
     it('should handle parsing of separators', function () {
       const str =
         'classDiagram\n' +
diff --git a/src/diagrams/class/classRenderer.js b/src/diagrams/class/classRenderer.js
index 181df2ece..f90e3d620 100644
--- a/src/diagrams/class/classRenderer.js
+++ b/src/diagrams/class/classRenderer.js
@@ -314,11 +314,24 @@ const drawClass = function(elem, classDef) {
 
   const id = classDef.id;
   const buildDisplayTextForMethod = function(txt) {
+    let regEx = /(\+|-|~|#)?(\w+)\((\w+\[?\]?)?\s?(\w+)?\)([*|$])?\s?(\w+\[?\]?)?/;
     let cssStyle = '';
-    let methodEnd = txt.indexOf(')') + 1;
-    let methodName = txt.substring(0, methodEnd);
+    let displayText = txt;
+    let methodName = txt;
+    let classifier = '';
 
-    let classifier = txt.substring(methodEnd, methodEnd + 1);
+    let parsedText = txt.match(regEx);
+
+    if (parsedText) {
+      let visibility = parsedText[1] ? parsedText[1] : '';
+      methodName = parsedText[2] ? parsedText[2] : '';
+      let parameterType = parsedText[3] ? parsedText[3] : '';
+      let parameterName = parsedText[4] ? parsedText[4] : '';
+      classifier = parsedText[5] ? parsedText[5] : '';
+      let returnType = parsedText[6] ? ' : ' + parsedText[6] : '';
+      displayText =
+        visibility + methodName + '(' + parameterType + ' ' + parameterName + ')' + returnType;
+    }
 
     switch (classifier) {
       case '*':
@@ -331,19 +344,10 @@ const drawClass = function(elem, classDef) {
 
     let method = {
       methodname: methodName,
-      displayText: methodName,
+      displayText: displayText,
       cssStyle: cssStyle
     };
 
-    let returnTypeStart = txt.indexOf('[') + 1;
-    let returnTypeEnd = txt.indexOf(']');
-
-    if (returnTypeStart > 1 && returnTypeEnd > returnTypeStart) {
-      let returnType = txt.substring(returnTypeStart, returnTypeEnd);
-
-      method.displayText = methodName + ' : ' + returnType;
-    }
-
     return method;
   };