fix: relationships when the same contract names are used in different…

… folders

lib/parser.d.ts

@@ -1,3 +1,3 @@
 import { ASTNode } from '@solidity-parser/parser';
 import { UmlClass } from './umlClass';
-export declare function convertNodeToUmlClass(node: ASTNode, codeSource: string): UmlClass[];
+export declare function convertNodeToUmlClass(node: ASTNode, codePath: string): UmlClass[];

lib/umlClass.d.ts

@@ -47,7 +47,8 @@ export interface Association {
 }
 export interface ClassProperties {
     name: string;
-    codeSource: string;
+    codePath: string;
+    importedFileNames?: string[];
     stereotype?: ClassStereotype;
     enums?: {
         [name: string]: string[];
@@ -62,7 +63,8 @@ export declare class UmlClass implements ClassProperties {
     static idCounter: number;
     id: number;
     name: string;
-    codeSource: string;
+    codePath: string;
+    importedPaths?: string[];
     stereotype?: ClassStereotype;
     attributes: Attribute[];
     operators: Operator[];

package.json

@@ -1,6 +1,6 @@
 {
   "name": "sol2uml",
-  "version": "1.1.20",
+  "version": "1.1.21",
   "description": "Unified Modeling Language (UML) class diagram generator for Solidity contracts",
   "main": "./lib/index.js",
   "types": "./lib/index.d.ts",

src/contracts/relativeImports/1/Parent.sol

@@ -0,0 +1,4 @@
+pragma solidity ^0.6.2;
+
+contract Parent {
+}

src/contracts/relativeImports/1/sub/Child.sol

@@ -0,0 +1,7 @@
+pragma solidity ^0.6.2;
+
+import "../Parent.sol";
+
+contract Child is Parent {
+
+}

src/contracts/relativeImports/2/Parent.sol

@@ -0,0 +1,4 @@
+pragma solidity ^0.6.2;
+
+contract Parent {
+}

src/contracts/relativeImports/2/sub/Child.sol

@@ -0,0 +1,7 @@
+pragma solidity ^0.6.2;
+
+import "../Parent.sol";
+
+contract Child is Parent {
+
+}

src/ts/__tests__/fileParser.test.ts

@@ -6,7 +6,7 @@ describe('Parser', () => {
             const files = await getSolidityFilesFromFolderOrFile(
                 './src/contracts'
             )
-            expect(files).toHaveLength(14)
+            expect(files).toHaveLength(18)
         })
 
         test('get Solidity files from folder with no sol files', async () => {
@@ -23,7 +23,7 @@ describe('Parser', () => {
 
         test('get Solidity files including Open Zeppelin', async () => {
             const files = await getSolidityFilesFromFolderOrFile('.')
-            expect(files).toHaveLength(100)
+            expect(files).toHaveLength(104)
         })
 
         describe('Failures', () => {

src/ts/converter.ts

@@ -84,29 +84,28 @@ arrowhead=open
 node [shape=record, style=filled, fillcolor=gray95]`
 
     // Sort UML Classes by folder of source file
-    const umlClassesSortedBySourceFiles = sortUmlClassesBySourceFolder(
-        umlClasses
-    )
+    const umlClassesSortedByCodePath = sortUmlClassesByCodePath(umlClasses)
 
-    let sourceFolder = ''
-    for (const umlClass of umlClassesSortedBySourceFiles) {
-        if (sourceFolder !== umlClass.codeSource) {
+    let currentCodeFolder = ''
+    for (const umlClass of umlClassesSortedByCodePath) {
+        const codeFolder = path.dirname(umlClass.codePath)
+        if (currentCodeFolder !== codeFolder) {
             // Need to close off the last subgraph if not the first
-            if (sourceFolder != '') {
+            if (currentCodeFolder != '') {
                 dotString += '\n}'
             }
 
             dotString += `
 subgraph ${getSubGraphName(clusterFolders)} {
-label="${umlClass.codeSource}"`
+label="${codeFolder}"`
 
-            sourceFolder = umlClass.codeSource
+            currentCodeFolder = codeFolder
         }
         dotString += dotUmlClass(umlClass, classOptions)
     }
 
     // Need to close off the last subgraph if not the first
-    if (sourceFolder != '') {
+    if (currentCodeFolder != '') {
         dotString += '\n}'
     }
 
@@ -128,12 +127,12 @@ function getSubGraphName(clusterFolders: boolean = false) {
     return ` graph_${subGraphCount++}`
 }
 
-function sortUmlClassesBySourceFolder(umlClasses: UmlClass[]): UmlClass[] {
+function sortUmlClassesByCodePath(umlClasses: UmlClass[]): UmlClass[] {
     return umlClasses.sort((a, b) => {
-        if (a.codeSource < b.codeSource) {
+        if (a.codePath < b.codePath) {
             return -1
         }
-        if (a.codeSource > b.codeSource) {
+        if (a.codePath > b.codePath) {
             return 1
         }
         return 0
@@ -142,18 +141,23 @@ function sortUmlClassesBySourceFolder(umlClasses: UmlClass[]): UmlClass[] {
 
 export function addAssociationsToDot(umlClasses: UmlClass[]): string {
     let dotString: string = ''
-    let nameToIdMap: { [className: string]: UmlClass } = {}
-
-    for (const umlClass of umlClasses) {
-        nameToIdMap[umlClass.name] = umlClass
-    }
 
     // for each class
     for (const sourceUmlClass of umlClasses) {
         // for each association in that class
         for (const association of Object.values(sourceUmlClass.associations)) {
-            // find the target class
-            const targetUmlClass = nameToIdMap[association.targetUmlClassName]
+            // find the target class with the same class name and
+            // codePath of the target in the importedPaths of the source class OR
+            // the codePath of the target is the same as the codePath pf the source class
+            const targetUmlClass = umlClasses.find((targetUmlClass) => {
+                return (
+                    targetUmlClass.name === association.targetUmlClassName &&
+                    (sourceUmlClass.importedPaths.includes(
+                        targetUmlClass.codePath
+                    ) ||
+                        sourceUmlClass.codePath === targetUmlClass.codePath)
+                )
+            })
             if (targetUmlClass) {
                 dotString += addAssociationToDot(
                     sourceUmlClass,

src/ts/fileParser.ts

@@ -1,5 +1,5 @@
 import { lstatSync, readFileSync } from 'fs'
-import { basename, dirname, extname, relative } from 'path'
+import { basename, extname, relative } from 'path'
 import klaw from 'klaw'
 import { ASTNode, parse } from '@solidity-parser/parser'
 import { VError } from 'verror'
@@ -25,10 +25,9 @@ export const parseUmlClassesFromFiles = async (
     for (const file of files) {
         const node = await parseSolidityFile(file)
 
-        const sourceFolder = dirname(file)
-        const relativeSourceFolder = relative(process.cwd(), sourceFolder)
+        const relativePath = relative(process.cwd(), file)
 
-        const umlClass = convertNodeToUmlClass(node, relativeSourceFolder)
+        const umlClass = convertNodeToUmlClass(node, relativePath)
         umlClasses = umlClasses.concat(umlClass)
     }
 

src/ts/parser.ts

@@ -13,14 +13,16 @@ import {
     UmlClass,
     Visibility,
 } from './umlClass'
+import { dirname, join } from 'path'
 
 const debug = require('debug')('sol2uml')
 
 export function convertNodeToUmlClass(
     node: ASTNode,
-    codeSource: string
+    codePath: string
 ): UmlClass[] {
     let umlClasses: UmlClass[] = []
+    const importedPaths: string[] = []
 
     if (node.type === 'SourceUnit') {
         node.children.forEach((childNode) => {
@@ -29,21 +31,26 @@ export function convertNodeToUmlClass(
 
                 let umlClass = new UmlClass({
                     name: childNode.name,
-                    codeSource: codeSource,
+                    codePath,
                 })
 
                 umlClass = parseContractDefinition(umlClass, childNode)
 
                 umlClasses.push(umlClass)
             } else if (childNode.type === 'ImportDirective') {
-                // TODO travers to parse imports
-                // importedContracts.push(contract)
+                const codeFolder = dirname(codePath)
+                const importPath = join(codeFolder, childNode.path)
+                importedPaths.push(importPath)
             }
         })
     } else {
         throw new Error(`AST node not of type SourceUnit`)
     }
 
+    umlClasses.forEach((umlClass) => {
+        umlClass.importedPaths = importedPaths
+    })
+
     return umlClasses
 }
 

src/ts/umlClass.ts

@@ -56,7 +56,8 @@ export interface Association {
 
 export interface ClassProperties {
     name: string
-    codeSource: string
+    codePath: string
+    importedFileNames?: string[]
     stereotype?: ClassStereotype
     enums?: { [name: string]: string[] }
     attributes?: Attribute[]
@@ -69,7 +70,8 @@ export class UmlClass implements ClassProperties {
 
     id: number
     name: string
-    codeSource: string
+    codePath: string
+    importedPaths?: string[]
     stereotype?: ClassStereotype
 
     attributes: Attribute[] = []