Support imports gated by compilation directives (#205)

Some source modules conditionally gate import declarations via compilation directives to support platform-specific dependencies. This change allows import declarations in the generated mock file to preserve compilation directives from the original source file.
typealiased · Jul 9, 2021 · aab1f47 · aab1f47
1 parent 0646c87
commit aab1f47
Show file tree

Hide file tree

Showing 8 changed files with 66 additions and 21 deletions.
diff --git a/Sources/MockingbirdGenerator/Generator/Operations/FileGenerator.swift b/Sources/MockingbirdGenerator/Generator/Operations/FileGenerator.swift
@@ -15,17 +15,17 @@ import os.log
 class FileGenerator {
   let mockableTypes: [MockableType]
   let mockedTypeNames: Set<String>?
-  let imports: Set<String>
+  let parsedFiles: [ParsedFile]
   let config: GenerateFileConfig
 
   init(mockableTypes: [MockableType],
        mockedTypeNames: Set<String>?,
-       imports: Set<String>,
+       parsedFiles: [ParsedFile],
        config: GenerateFileConfig) {
     self.mockableTypes = config.onlyMockProtocols ?
       mockableTypes.filter({ $0.kind == .protocol }) : mockableTypes
     self.mockedTypeNames = mockedTypeNames
-    self.imports = imports
+    self.parsedFiles = parsedFiles
     self.config = config
   }
 
@@ -48,12 +48,33 @@ class FileGenerator {
       headerSections.append("#if \(condition)")
     }
 
-    let moduleImports = (
-      imports.union(["import Foundation", "@testable import Mockingbird"]).union(
-        config.shouldImportModule ? ["@testable import \(config.moduleName)"] : []
-      )
-    ).sorted()
-    headerSections.append(moduleImports.joined(separator: "\n"))
+    let implicitImports = [
+      ImportDeclaration("Foundation"),
+      ImportDeclaration("Mockingbird", testable: true),
+      config.shouldImportModule ? ImportDeclaration(config.moduleName, testable: true) : nil,
+    ].compactMap({ $0?.fullDeclaration })
+
+    let explicitImports = parsedFiles
+      .filter({ $0.shouldMock })
+      .flatMap({ file in
+        file.importDeclarations.map({ importDeclaration -> String in
+          let compilationDirectives = file.compilationDirectives
+            .filter({ $0.range.contains(importDeclaration.offset) })
+          guard !compilationDirectives.isEmpty else {
+            return importDeclaration.fullDeclaration
+          }
+          let start = compilationDirectives.map({ $0.declaration }).joined(separator: "\n")
+          let end = compilationDirectives.map({ _ in "#endif" }).joined(separator: "\n")
+          return [
+            start,
+            importDeclaration.fullDeclaration,
+            end,
+          ].joined(separator: "\n")
+        })
+      })
+
+    let allImports = Set(implicitImports + explicitImports).sorted()
+    headerSections.append(allImports.joined(separator: "\n"))
 
     return PartialFileContent(contents: """
     //

diff --git a/Sources/MockingbirdGenerator/Generator/Operations/GenerateFileOperation.swift b/Sources/MockingbirdGenerator/Generator/Operations/GenerateFileOperation.swift
@@ -62,7 +62,7 @@ public class GenerateFileOperation: BasicOperation {
       let generator = FileGenerator(
         mockableTypes: processTypesResult.mockableTypes,
         mockedTypeNames: findMockedTypesResult?.allMockedTypeNames,
-        imports: processTypesResult.imports,
+        parsedFiles: processTypesResult.parsedFiles,
         config: config
       )
       contents = generator.generate()

diff --git a/Sources/MockingbirdGenerator/Parser/Models/ParsedFile.swift b/Sources/MockingbirdGenerator/Parser/Models/ParsedFile.swift
@@ -55,7 +55,7 @@ struct ParsedFile {
   }
 }
 
-struct CompilationDirective: Comparable {
+struct CompilationDirective: Comparable, Hashable {
   let range: Range<Int64> // Byte offset bounds of the compilation directive declaration.
   let declaration: String
   let condition: String?
@@ -109,16 +109,30 @@ struct CompilationDirective: Comparable {
   static func < (lhs: CompilationDirective, rhs: CompilationDirective) -> Bool {
     return lhs.range.lowerBound < rhs.range.lowerBound
   }
+
+  func hash(into hasher: inout Hasher) {
+    hasher.combine(declaration)
+    hasher.combine(condition)
+  }
 }
 
 struct ImportDeclaration: Hashable {
   let moduleName: String
   let fullPath: String
   let fullDeclaration: String
+  let offset: Int64
 
-  static var implicitSwiftImport: ImportDeclaration {
-    return ImportDeclaration(moduleName: "Swift",
-                             fullPath: "Swift",
-                             fullDeclaration: "import Swift")
+  init(moduleName: String, fullPath: String, fullDeclaration: String, offset: Int64) {
+    self.moduleName = moduleName
+    self.fullPath = fullPath
+    self.fullDeclaration = fullDeclaration
+    self.offset = offset
+  }
+
+  init(_ moduleName: String, testable: Bool = false) {
+    self.moduleName = moduleName
+    self.fullPath = moduleName
+    self.fullDeclaration = (testable ? "@testable " : "") + "import " + moduleName
+    self.offset = 0
   }
 }
diff --git a/Sources/MockingbirdGenerator/Parser/Operations/ParseFilesOperation.swift b/Sources/MockingbirdGenerator/Parser/Operations/ParseFilesOperation.swift
@@ -18,7 +18,6 @@ public class ParseFilesOperation: BasicOperation {
 
   public class Result {
     fileprivate(set) var parsedFiles = [ParsedFile]()
-    fileprivate(set) var imports = Set<String>()
     fileprivate(set) var moduleDependencies = [String: Set<String>]()
   }
 
@@ -63,7 +62,6 @@ public class ParseFilesOperation: BasicOperation {
       })
     }
 
-    result.imports = Set(result.parsedFiles.filter({ $0.shouldMock }).flatMap({ $0.imports }))
     result.moduleDependencies = extractSourcesResult.moduleDependencies
   }
 }
diff --git a/Sources/MockingbirdGenerator/Parser/Operations/ParseSingleFileOperation.swift b/Sources/MockingbirdGenerator/Parser/Operations/ParseSingleFileOperation.swift
@@ -111,7 +111,7 @@ class ParseSwiftSyntaxOperation: BasicOperation {
     retainForever(parser)
 
     // All Swift files implicitly import the Swift standard library.
-    result.importDeclarations = parser.importedPaths.union([.implicitSwiftImport])
+    result.importDeclarations = parser.importedPaths.union([ImportDeclaration("Swift")])
     result.compilationDirectives = parser.directives.sorted()
   }
 }

diff --git a/Sources/MockingbirdGenerator/Parser/Operations/ProcessTypesOperation.swift b/Sources/MockingbirdGenerator/Parser/Operations/ProcessTypesOperation.swift
@@ -17,7 +17,7 @@ public class ProcessTypesOperation: BasicOperation {
 
   public class Result {
     fileprivate(set) var mockableTypes = [MockableType]()
-    fileprivate(set) var imports = Set<String>()
+    fileprivate(set) var parsedFiles = [ParsedFile]()
   }
 
   public let result = Result()
@@ -76,7 +76,7 @@ public class ProcessTypesOperation: BasicOperation {
       result.mockableTypes = flattenInheritanceOperations
         .compactMap({ $0.result.mockableType })
         .filter({ !$0.isContainedType })
-      result.imports = parseFilesResult.imports
+      result.parsedFiles = parseFilesResult.parsedFiles
       log("Created \(result.mockableTypes.count) mockable type\(result.mockableTypes.count != 1 ? "s" : "")")
     }
   }

diff --git a/Sources/MockingbirdGenerator/Parser/Operations/SourceFileAuxiliaryParser.swift b/Sources/MockingbirdGenerator/Parser/Operations/SourceFileAuxiliaryParser.swift
@@ -30,9 +30,13 @@ class SourceFileAuxiliaryParser: SyntaxVisitor {
     guard let moduleName = node.path.first?.name.text else { return .skipChildren }
     let fullPath = node.path.withoutTrivia().description
     let fullDeclaration = node.withoutTrivia().description
+    let sourceRange = node.sourceRange(converter: converter,
+                                       afterLeadingTrivia: true,
+                                       afterTrailingTrivia: true)
     importedPaths.insert(ImportDeclaration(moduleName: moduleName,
                                            fullPath: fullPath,
-                                           fullDeclaration: fullDeclaration))
+                                           fullDeclaration: fullDeclaration,
+                                           offset: Int64(sourceRange.start.offset)))
     return .skipChildren
   }
 

diff --git a/Sources/MockingbirdTestsHost/CompilationDirectives.swift b/Sources/MockingbirdTestsHost/CompilationDirectives.swift
@@ -7,6 +7,14 @@
 
 import Foundation
 
+#if DEBUG
+import Foundation
+#endif
+
+#if NEVER_TRUE
+import InvalidModule
+#endif
+
 #if !(DEBUG)
 #warning("Testing #warning compilation directive keyword handling")
 #endif