Support elseif/else directive

Sources/MockoloFramework/Models/ConditionalImportBlock.swift

@@ -0,0 +1,45 @@
+//
+//  Copyright (c) 2018. Uber Technologies
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+/// Represents import content: either a simple import statement or a nested conditional block
+public indirect enum ImportContent {
+    case simple(Import)
+    case conditional(ConditionalImportBlock)
+}
+
+/// Represents a conditional import block (#if/#elseif/#else/#endif)
+public struct ConditionalImportBlock {
+    /// Represents a single clause in a conditional import block
+    public struct Clause {
+        public let type: ClauseType
+        public let condition: String?  // nil for #else
+        public var contents: [ImportContent]
+
+        public init(type: ClauseType, condition: String?, contents: [ImportContent]) {
+            self.type = type
+            self.condition = condition
+            self.contents = contents
+        }
+    }
+
+    public let clauses: [Clause]
+    public let offset: Int64
+
+    public init(clauses: [Clause], offset: Int64) {
+        self.clauses = clauses
+        self.offset = offset
+    }
+}

Sources/MockoloFramework/Models/IfMacroModel.swift

@@ -14,36 +14,69 @@
 //  limitations under the License.
 //
 
+/// Represents the type of a clause in an #if/#elseif/#else block
+public enum ClauseType: Comparable {
+    case `if`
+    case elseif(order: Int)
+    case `else`
+
+    // Comparable implementation: if < elseif(0) < elseif(1) < ... < else
+    public static func < (lhs: ClauseType, rhs: ClauseType) -> Bool {
+        switch (lhs, rhs) {
+        case (.if, .elseif), (.if, .else), (.elseif, .else):
+            true
+        case (.elseif(let l), .elseif(let r)):
+            l < r
+        default:
+            false
+        }
+    }
+}
+
 final class IfMacroModel: Model {
-    let name: String
+    /// Represents a single clause in a conditional compilation block
+    struct Clause {
+        let type: ClauseType
+        let condition: String?  // nil for #else
+        let entities: [(String, Model)]
+    }
+
+    let clauses: [Clause]
     let offset: Int64
-    let entities: [(String, Model)]
 
     var modelType: ModelType {
-        return .macro
+        .macro
+    }
+
+    var name: String {
+        clauses.first?.condition ?? ""
     }
 
     var fullName: String {
-        return entities.map {$0.0}.joined(separator: "_")
+        clauses.flatMap(\.entities).map { $0.0 }.joined(separator: "_")
     }
-
-    init(name: String,
-         offset: Int64,
-         entities: [(String, Model)]) {
-        self.name = name
-        self.entities = entities
+
+    /// Creates an IfMacroModel with multiple clauses
+    init(clauses: [Clause], offset: Int64) {
+        self.clauses = clauses
         self.offset = offset
     }
-
+
+    /// Initializer for simple #if blocks
+    convenience init(name: String,
+                     offset: Int64,
+                     entities: [(String, Model)]) {
+        let clause = Clause(type: .if, condition: name, entities: entities)
+        self.init(clauses: [clause], offset: offset)
+    }
+
     func render(
         context: RenderContext,
         arguments: GenerationArguments
     ) -> String? {
-        return applyMacroTemplate(
-            name: name,
+        applyMacroTemplate(
             context: context,
-            arguments: arguments,
-            entities: entities
+            arguments: arguments
         )
     }
 }

Sources/MockoloFramework/Models/Import.swift

@@ -15,10 +15,10 @@
 //
 
 /// A structure defining an "import" statement parsed by `Generator`, including various modifiers.
-struct Import: CustomStringConvertible {
+public struct Import: CustomStringConvertible {
 
     /// The access level of the import
-    enum ACL: String {
+    public enum ACL: String {
         case `private`
         case `fileprivate`
         case `internal`
@@ -37,18 +37,18 @@ struct Import: CustomStringConvertible {
     }
 
     /// A modifier that precedes the "import" keyword. ACL and "@testable" are mutually exclusive.
-    enum Modifier: RawRepresentable {
+    public enum Modifier: RawRepresentable {
         case acl(ACL)
         case testable
 
-        var rawValue: String {
+        public var rawValue: String {
             switch self {
             case .acl(let acl): acl.rawValue
             case .testable: "@testable"
             }
         }
 
-        init?(rawValue: String) {
+        public init?(rawValue: String) {
             if rawValue == "@testable" {
                 self = .testable
             } else if let acl = ACL(rawValue: rawValue) {
@@ -60,18 +60,18 @@ struct Import: CustomStringConvertible {
     }
 
     /// Name of the module
-    let moduleName: String
+    public let moduleName: String
 
     /// A modifier preceding the "import" keyword (e.g. public, internal, @testable)
-    var modifier: Modifier?
+    public var modifier: Modifier?
 
     /// An opaque string preceding the entire import statement (typically `#if FOO\n` for nested macro support)
-    let prefix: String?
+    public let prefix: String?
 
     /// An opaque string following the entire import statement (typically `\n#endif` for nested macro support)
-    let suffix: String?
+    public let suffix: String?
 
-    var description: String {
+    public var description: String {
         let line: String
         if let modifier {
             line = "\(modifier.rawValue) import \(moduleName)"
@@ -81,7 +81,7 @@ struct Import: CustomStringConvertible {
         return [prefix, line, suffix].compactMap { $0 }.joined()
     }
 
-    init(
+    public init(
         moduleName: String,
         modifier: Modifier? = nil,
         prefix: String? = nil,

Sources/MockoloFramework/Models/ParsedEntity.swift

@@ -150,7 +150,20 @@ struct GenerationArguments {
     )
 }
 
-public typealias ImportMap = [String: [String: [String]]]
+/// Structured import data parsed from a source file
+public struct ParsedImports {
+    /// Top-level imports without conditional compilation
+    public var topLevel: [Import]
+    /// Conditional import blocks (#if/#elseif/#else/#endif)
+    public var conditional: [ConditionalImportBlock]
+
+    public init(topLevel: [Import] = [], conditional: [ConditionalImportBlock] = []) {
+        self.topLevel = topLevel
+        self.conditional = conditional
+    }
+}
+
+public typealias ImportMap = [String: ParsedImports]
 
 /// Metadata for a type being mocked
 public final class Entity {

Sources/MockoloFramework/Operations/ImportsHandler.swift

@@ -22,63 +22,113 @@ func handleImports(pathToImportsMap: ImportMap,
                    testableImports: [String]?,
                    relevantPaths: [String]) -> String {
 
-    var imports = [String: [Import]]()
-    let defaultKey = ""
-    if imports[defaultKey] == nil {
-        imports[defaultKey] = []
-    }
+    var topLevelImports: [Import] = []
+    var conditionalBlocks: [ConditionalImportBlock] = []
 
-    for (path, importMap) in pathToImportsMap {
+    // 1. Collect imports from all relevant files
+    for (path, parsedImports) in pathToImportsMap {
         guard relevantPaths.contains(path) else { continue }
-        for (k, v) in importMap {
-            if imports[k] == nil {
-                imports[k] = []
-            }
 
-            if let ex = excludeImports {
-                let filtered = v.filter{ !ex.contains($0.moduleNameInImport) }
-                imports[k]?.append(contentsOf: filtered.compactMap { Import(line: $0) })
-            } else {
-                imports[k]?.append(contentsOf: v.compactMap { Import(line: $0) })
-            }
-        }
+        // Collect top-level imports
+        topLevelImports.append(contentsOf: parsedImports.topLevel)
+
+        // Collect conditional blocks
+        conditionalBlocks.append(contentsOf: parsedImports.conditional)
+    }
+
+    // 2. Apply excludeImports filter to top-level imports
+    if let excludes = excludeImports {
+        topLevelImports = topLevelImports.filter { !excludes.contains($0.moduleName) }
     }
 
+    // 3. Add custom imports
     if let customImports = customImports {
-        imports[defaultKey]?.append(contentsOf: customImports.compactMap { Import(moduleName: $0) })
+        topLevelImports.append(contentsOf: customImports.map { Import(moduleName: $0) })
+    }
+
+    // 4. Resolve duplicates in top-level imports
+    var resolvedTopLevel = topLevelImports.resolved()
+
+    // 5. Apply testableImports modifier
+    if let testableImportNames = testableImports, !testableImportNames.isEmpty {
+        let (passthroughImports, candidateImports) = resolvedTopLevel.partitioned(by: { testableImportNames.contains($0.moduleName) })
+        let mappedImports = candidateImports.map(\.asTestable)
+        let newImports: [Import] = testableImportNames.compactMap { name in
+            guard !mappedImports.contains(where: { $0.moduleName == name }) else { return nil }
+            return Import(moduleName: name, modifier: .testable)
+        }
+        resolvedTopLevel = (passthroughImports + mappedImports + newImports).resolved()
+    }
+
+    // 6. Sort conditional blocks by offset (file appearance order)
+    let sortedBlocks = conditionalBlocks.sorted(by: { $0.offset < $1.offset })
+
+    // 7. Generate output
+    var lines: [String] = []
+
+    if !resolvedTopLevel.isEmpty {
+        lines.append(resolvedTopLevel.lines())
     }
 
-    var sortedImports = [String: [Import]]()
-    for (k, v) in imports {
-        sortedImports[k] = v.resolved()
+    for block in sortedBlocks {
+        lines.append(renderConditionalBlock(block, excludeImports: excludeImports))
     }
 
-    if let existingSet = sortedImports[defaultKey] {
-        if let testableImportNames = testableImports, !testableImportNames.isEmpty {
-            let (passthroughImports, candidateImports) = existingSet.partitioned(by: { testableImportNames.contains($0.moduleName) })
-            let mappedImports = candidateImports.map(\.asTestable)
-            let newImports: [Import] = testableImportNames.compactMap { name in
-                guard !mappedImports.contains(where: { $0.moduleName == name }) else { return nil }
-                return Import(moduleName: name, modifier: .testable)
+    return lines.joined(separator: "\n")
+}
+
+/// Recursively renders a ConditionalImportBlock
+private func renderConditionalBlock(_ block: ConditionalImportBlock, excludeImports: [String]?) -> String {
+    var result = ""
+
+    for (index, clause) in block.clauses.enumerated() {
+        // Render directive line
+        switch clause.type {
+        case .if:
+            result += "#if \(clause.condition ?? "")\n"
+        case .elseif:
+            result += "#elseif \(clause.condition ?? "")\n"
+        case .else:
+            result += "#else\n"
+        }
+
+        // Render contents of this clause
+        var clauseLines: [String] = []
+        var simpleImports: [Import] = []
+
+        for content in clause.contents {
+            switch content {
+            case .simple(let imp):
+                // Filter excluded imports
+                if let excludes = excludeImports, excludes.contains(imp.moduleName) {
+                    continue
+                }
+                simpleImports.append(imp)
+            case .conditional(let nestedBlock):
+                // First output accumulated simple imports
+                if !simpleImports.isEmpty {
+                    clauseLines.append(simpleImports.resolved().lines())
+                    simpleImports = []
+                }
+                // Recursively render nested block
+                clauseLines.append(renderConditionalBlock(nestedBlock, excludeImports: excludeImports))
             }
-            sortedImports[defaultKey] = (passthroughImports + mappedImports + newImports).resolved()
         }
-    }
 
-    let sortedKeys = sortedImports.keys.sorted()
-    let importsStr = sortedKeys.map { k in
-        let v = sortedImports[k]
-        let lines = v?.lines() ?? ""
-        if k.isEmpty {
-            return lines
-        } else {
-            return """
-            #if \(k)
-            \(lines)
-            #endif
-            """
+        // Output remaining simple imports
+        if !simpleImports.isEmpty {
+            clauseLines.append(simpleImports.resolved().lines())
+        }
+
+        let clauseContent = clauseLines.joined(separator: "\n")
+        if !clauseContent.isEmpty {
+            result += clauseContent
+            if index < block.clauses.count - 1 {
+                result += "\n"
+            }
         }
-    }.joined(separator: "\n")
+    }
 
-    return importsStr
+    result += "\n#endif"
+    return result
 }