cleanup

uber · Dec 6, 2019 · da9d6e3 · da9d6e3
1 parent ddc3893
commit da9d6e3
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 47 deletions.
diff --git a/Sources/MockoloFramework/Operations/Generator.swift b/Sources/MockoloFramework/Operations/Generator.swift
@@ -34,7 +34,7 @@ public func generate(sourceDirs: [String]?,
                      annotation: String,
                      header: String?,
                      macro: String?,
-                     parserType: ParserType = .swiftSyntax,
+                     parserType: ParserType,
                      to outputFilePath: String,
                      loggingLevel: Int,
                      concurrencyLimit: Int?,

diff --git a/Sources/MockoloFramework/Operations/OutputWriter.swift b/Sources/MockoloFramework/Operations/OutputWriter.swift
@@ -29,15 +29,14 @@ func write(candidates: [(String, Int64)],
     for path in relevantPaths {
         if let lines = pathToImportsMap[path] {
             importLines.append(contentsOf: lines)
-        } else {
-            for (filepath, filecontent, offset) in pathToContentMap {
-                if path == filepath {
-                    let v = findImportLines(data: filecontent, offset: offset)
-                    importLines.append(contentsOf: v)
-                }
-            }
         }
     }
+    for (filepath, filecontent, offset) in pathToContentMap {
+        let v = findImportLines(data: filecontent, offset: offset)
+        importLines.append(contentsOf: v)
+        break
+    }
+
     let importsSet = Set(importLines.map{$0.trimmingCharacters(in: .whitespaces)})
     let importLineStr = importsSet.sorted().joined(separator: "\n")
 

diff --git a/Sources/MockoloFramework/Operations/ProcessedTypeMapGenerator.swift b/Sources/MockoloFramework/Operations/ProcessedTypeMapGenerator.swift
@@ -45,7 +45,7 @@ func generateProcessedTypeMap(_ paths: [String],
             }
         }
     default:
-        var treeVisitor = EntityVisitor()
+        var treeVisitor = EntityVisitor(entityType: .classType)
         for filePath in paths {
             generateProcessedModels(filePath, treeVisitor: &treeVisitor, lock: nil, process: process)
         }
@@ -106,24 +106,3 @@ private func generateProcessedModels(_ path: String,
         fatalError(error.localizedDescription)
     }
 }
-
-private func generateProcessedModels(_ path: String,
-                                     treeVisitor: inout EntityVisitor,
-                                     lock: NSLock?,
-                                     process: @escaping ([Entity], [String]) -> ()) {
-    do {
-        var results = [Entity]()
-        let node = try SyntaxParser.parse(path)
-        node.walk(&treeVisitor)
-        let ret = treeVisitor.entities
-        results.append(contentsOf: ret)
-
-        let imports = treeVisitor.imports
-        treeVisitor.reset()
-        lock?.lock()
-        process(results, imports)
-        lock?.unlock()
-    } catch {
-        fatalError(error.localizedDescription)
-    }
-}
diff --git a/Sources/MockoloFramework/Operations/ProtocolMapGenerator.swift b/Sources/MockoloFramework/Operations/ProtocolMapGenerator.swift
@@ -76,7 +76,7 @@ private func generateProtcolMap(dirs: [String],
             }
         }
     default:
-        var treeVisitor = EntityVisitor(annotation: annotation)
+        var treeVisitor = EntityVisitor(annotation: annotation, entityType: .protocolType)
         scanPaths(dirs) { filePath in
             generateProtcolMap(filePath,
                                exclusionSuffixes: exclusionSuffixes,
@@ -128,7 +128,7 @@ private func generateProtcolMap(files: [String],
             }
         }
     default:
-        var treeVisitor = EntityVisitor(annotation: annotation)
+        var treeVisitor = EntityVisitor(annotation: annotation, entityType: .protocolType)
         for filePath in files {
             generateProtcolMap(filePath,
                                exclusionSuffixes: exclusionSuffixes,

diff --git a/Sources/MockoloFramework/Utils/SwiftSyntaxExtensions.swift b/Sources/MockoloFramework/Utils/SwiftSyntaxExtensions.swift
@@ -336,13 +336,20 @@ extension AssociatedtypeDeclSyntax {
 
 
 }
+
+enum EntityType {
+    case protocolType, classType
+}
+
 final class EntityVisitor: SyntaxVisitor {
     var entities: [Entity] = []
     var imports: [String] = []
     let annotation: String
+    let entityType: EntityType
 
-    init(annotation: String = "") {
+    init(annotation: String = "", entityType: EntityType) {
         self.annotation = annotation
+        self.entityType = entityType
     }
 
     func reset() {
@@ -351,25 +358,29 @@ final class EntityVisitor: SyntaxVisitor {
     }
 
     func visit(_ node: ProtocolDeclSyntax) -> SyntaxVisitorContinueKind {
-        var isAnnotated = false
-        var overrides: [String: String]? = nil
-        if !annotation.isEmpty {
-            let metadata = node.annotationMetadata(with: annotation)
-            isAnnotated = metadata != nil
-            overrides = metadata?.typealiases
+        if entityType == .protocolType {
+            var isAnnotated = false
+            var overrides: [String: String]? = nil
+            if !annotation.isEmpty {
+                let metadata = node.annotationMetadata(with: annotation)
+                isAnnotated = metadata != nil
+                overrides = metadata?.typealiases
+            }
+
+            let ent = Entity(entityNode: node,
+                             isAnnotated: isAnnotated,
+                             overrides: overrides,
+                             isProcessed: false)
+            entities.append(ent)
         }
-
-        let ent = Entity(entityNode: node,
-                         isAnnotated: isAnnotated,
-                         overrides: overrides,
-                         isProcessed: false)
-        entities.append(ent)
         return .skipChildren
     }
 
     func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind {
-        let ent = Entity(entityNode: node, isAnnotated: false, overrides: nil, isProcessed: true)
-        entities.append(ent)
+        if entityType == .classType {
+            let ent = Entity(entityNode: node, isAnnotated: false, overrides: nil, isProcessed: true)
+            entities.append(ent)
+        }
         return .skipChildren
     }