From 8d09211853d49e49d9771c136e628801b1c532e4 Mon Sep 17 00:00:00 2001
From: Arjan Duijzer <arjan@nxtstep.nl>
Date: Mon, 27 Feb 2017 17:58:07 +0100
Subject: [PATCH] Add option to exclude classes and/or protocols from mock
 generation

---
 .../GenerateMocksCommand.swift                | 31 +++++++++++++++----
 README.md                                     |  4 +++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/Generator/Source/cuckoo_generator/GenerateMocksCommand.swift b/Generator/Source/cuckoo_generator/GenerateMocksCommand.swift
index c6b0fb48..e5ca489f 100644
--- a/Generator/Source/cuckoo_generator/GenerateMocksCommand.swift
+++ b/Generator/Source/cuckoo_generator/GenerateMocksCommand.swift
@@ -12,9 +12,9 @@ import SourceKittenFramework
 import FileKit
 import CuckooGeneratorFramework
 
-private func curry<P1, P2, P3, P4, P5, P6, P7, P8, R>(_ f: @escaping (P1, P2, P3, P4, P5, P6, P7, P8) -> R)
-    -> (P1) -> (P2) -> (P3) -> (P4) -> (P5) -> (P6) -> (P7) -> (P8) -> R {
-        return { p1 in { p2 in { p3 in { p4 in { p5 in { p6 in { p7 in { p8 in f(p1, p2, p3, p4, p5, p6, p7, p8) } } } } } } } }
+private func curry<P1, P2, P3, P4, P5, P6, P7, P8, P9, R>(_ f: @escaping (P1, P2, P3, P4, P5, P6, P7, P8, P9) -> R)
+    -> (P1) -> (P2) -> (P3) -> (P4) -> (P5) -> (P6) -> (P7) -> (P8) -> (P9) -> R {
+        return { p1 in { p2 in { p3 in { p4 in { p5 in { p6 in { p7 in { p8 in { p9 in f(p1, p2, p3, p4, p5, p6, p7, p8, p9) } } } } } } } } }
 }
 
 public struct GenerateMocksCommand: CommandProtocol {
@@ -27,7 +27,9 @@ public struct GenerateMocksCommand: CommandProtocol {
         let inputFiles = inputPathValues.map { File(path: $0) }
         let tokens = inputFiles.flatMap { $0 }.map { Tokenizer(sourceFile: $0).tokenize() }
         let tokensWithInheritance = options.noInheritance ? tokens : mergeInheritance(tokens)
-        let parsedFiles = options.noClassMocking ? removeClasses(tokensWithInheritance) : tokensWithInheritance
+        let tokensWithoutClasses = options.noClassMocking ? removeClasses(tokensWithInheritance) : tokensWithInheritance
+        // filter excluded classes/protocols
+        let parsedFiles = removeClassesAndProtocols(from: tokensWithoutClasses, in: options.exclude)
 
         let headers = parsedFiles.map { options.noHeader ? "" : FileHeaderHandler.getHeader(of: $0, includeTimestamp: !options.noTimestamp) }
         let imports = parsedFiles.map { FileHeaderHandler.getImports(of: $0, testableFrameworks: options.testableFrameworks) }
@@ -68,6 +70,19 @@ public struct GenerateMocksCommand: CommandProtocol {
                 return FileRepresentation(sourceFile: $0.sourceFile, declarations: declarations)
             }.filter { !$0.declarations.isEmpty }
     }
+
+    private func removeClassesAndProtocols(from files: [FileRepresentation], in excluded: [String]) -> [FileRepresentation] {
+        return files.map {
+            FileRepresentation(sourceFile: $0.sourceFile, declarations: $0.declarations.filter { token in
+                guard let token = token as? ContainerToken else {
+                    return true
+                }
+                return !excluded.contains(token.name)
+            })
+        }.filter {
+            !$0.declarations.isEmpty
+        }
+    }
     
     public struct Options: OptionsProtocol {
         let files: [String]
@@ -76,12 +91,14 @@ public struct GenerateMocksCommand: CommandProtocol {
         let noTimestamp: Bool
         let noInheritance: Bool
         let testableFrameworks: [String]
+        let exclude: [String]
         let filePrefix: String
         let noClassMocking: Bool
 
-        public init(output: String, testableFrameworks: String, noHeader: Bool, noTimestamp: Bool, noInheritance: Bool, filePrefix: String, noClassMocking: Bool, files: [String]) {
+        public init(output: String, testableFrameworks: String, exclude: String, noHeader: Bool, noTimestamp: Bool, noInheritance: Bool, filePrefix: String, noClassMocking: Bool, files: [String]) {
             self.output = output
             self.testableFrameworks = testableFrameworks.components(separatedBy: ",").filter { !$0.isEmpty }
+            self.exclude = exclude.components(separatedBy: ",").filter { !$0.isEmpty }.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
             self.noHeader = noHeader
             self.noTimestamp = noTimestamp
             self.noInheritance = noInheritance
@@ -94,6 +111,8 @@ public struct GenerateMocksCommand: CommandProtocol {
             let output: Result<String, CommandantError<ClientError>> = m <| Option(key: "output", defaultValue: "GeneratedMocks.swift", usage: "Where to put the generated mocks.\nIf a path to a directory is supplied, each input file will have a respective output file with mocks.\nIf a path to a Swift file is supplied, all mocks will be in a single file.\nDefault value is `GeneratedMocks.swift`.")
 
             let testable: Result<String, CommandantError<ClientError>> = m <| Option(key: "testable", defaultValue: "", usage: "A comma separated list of frameworks that should be imported as @testable in the mock files.")
+			
+			let exclude: Result<String, CommandantError<ClientError>> = m <| Option(key: "exclude", defaultValue: "", usage: "A comma separated list of classes and protocols that should be skipped during mock generation.")
 
             let noHeader: Result<Bool, CommandantError<ClientError>> = m <| Option(key: "no-header", defaultValue: false, usage: "Do not generate file headers.")
 
@@ -110,13 +129,13 @@ public struct GenerateMocksCommand: CommandProtocol {
             return curry(Options.init)
                 <*> output
                 <*> testable
+                <*> exclude
                 <*> noHeader
                 <*> noTimestamp
                 <*> noInheritance
                 <*> filePrefix
                 <*> noClassMocking
                 <*> input
-
         }
     }
 }
diff --git a/README.md b/README.md
index 9126afa7..29466d13 100644
--- a/README.md
+++ b/README.md
@@ -391,6 +391,10 @@ Default value is `GeneratedMocks.swift`.
 
 A comma separated list of frameworks that should be imported as @testable in the mock files.
 
+##### `--exclude` (string)
+
+A comma separated list of classes and protocols that should be skipped during mock generation.  
+
 ##### `--no-header`
 
 Do not generate file headers.