fix: file level structs and enums not being hidden

naddison36 · Aug 12, 2021 · 2840fc3 · 2840fc3 · ToTheHighestTopMoon · Dec 18, 2021
1 parent e76b808
commit 2840fc3
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 8 deletions.
diff --git a/lib/converter.js b/lib/converter.js
diff --git a/lib/dotGenerator.js b/lib/dotGenerator.js
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "sol2uml",
-  "version": "1.1.29",
+  "version": "1.1.30",
   "description": "Unified Modeling Language (UML) class diagram generator for Solidity contracts",
   "main": "./lib/index.js",
   "types": "./lib/index.d.ts",

diff --git a/src/ts/converter.ts b/src/ts/converter.ts
@@ -187,7 +187,11 @@ function addAssociationToDot(
             (sourceUmlClass.stereotype === ClassStereotype.Library ||
                 targetUmlClass.stereotype === ClassStereotype.Library)) ||
         (classOptions.hideInterfaces &&
-            targetUmlClass.stereotype === ClassStereotype.Interface)
+            targetUmlClass.stereotype === ClassStereotype.Interface) ||
+        (classOptions.hideStructs &&
+            targetUmlClass.stereotype === ClassStereotype.Struct) ||
+        (classOptions.hideEnums &&
+            targetUmlClass.stereotype === ClassStereotype.Enum)
     ) {
         return ''
     }

diff --git a/src/ts/dotGenerator.ts b/src/ts/dotGenerator.ts
@@ -23,12 +23,15 @@ export const dotUmlClass = (
     umlClass: UmlClass,
     options: ClassOptions = {}
 ): string => {
-    // do not include library or interface classes if hidden
+    // do not include library, interface, struct or enum classes if hidden
     if (
         (options.hideLibraries &&
             umlClass.stereotype === ClassStereotype.Library) ||
         (options.hideInterfaces &&
-            umlClass.stereotype === ClassStereotype.Interface)
+            umlClass.stereotype === ClassStereotype.Interface) ||
+        (options.hideStructs &&
+            umlClass.stereotype === ClassStereotype.Struct) ||
+        (options.hideEnums && umlClass.stereotype === ClassStereotype.Enum)
     ) {
         return ''
     }