Support forced unwrapped type as function return type

Matejkob · Jun 25, 2024 · f4aee27 · f4aee27
1 parent 3209fc1
commit f4aee27
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 6 deletions.
diff --git a/Examples/Sources/ViewModel.swift b/Examples/Sources/ViewModel.swift
@@ -15,6 +15,7 @@ protocol ServiceProtocol {
   func insert(name: (any Codable)?, surname: (any Codable)?)
   func append(name: (any Codable) -> (any Codable)?)
   func get() async throws -> any Codable
+  func read() -> String!
 }
 
 final class ViewModel {

diff --git a/Sources/SpyableMacro/Factories/ClosureFactory.swift b/Sources/SpyableMacro/Factories/ClosureFactory.swift
@@ -32,6 +32,37 @@ struct ClosureFactory {
     variablePrefix: String,
     functionSignature: FunctionSignatureSyntax
   ) throws -> VariableDeclSyntax {
+    let returnClause: ReturnClauseSyntax
+    if let functionReturnClause = functionSignature.returnClause {
+      /*
+       func f() -> String!
+       */
+      if let implicitlyUnwrappedType = functionReturnClause.type.as(ImplicitlyUnwrappedOptionalTypeSyntax.self) {
+        var functionReturnClause = functionReturnClause
+        /*
+         `() -> String!` is not a valid code
+         so we have to convert it to `() -> String?
+         */
+        functionReturnClause.type = TypeSyntax(OptionalTypeSyntax(wrappedType: implicitlyUnwrappedType.wrappedType))
+        returnClause = functionReturnClause
+      /*
+       func f() -> Any
+       func f() -> Any?
+       */
+      } else {
+        returnClause = functionReturnClause
+      }
+    /*
+     func f()
+     */
+    } else {
+      returnClause = ReturnClauseSyntax(
+        type: IdentifierTypeSyntax(
+          name: .identifier("Void")
+        )
+      )
+    }
+
     let elements = TupleTypeElementListSyntax {
       TupleTypeElementSyntax(
         type: FunctionTypeSyntax(
@@ -44,12 +75,7 @@ struct ClosureFactory {
             asyncSpecifier: functionSignature.effectSpecifiers?.asyncSpecifier,
             throwsSpecifier: functionSignature.effectSpecifiers?.throwsSpecifier
           ),
-          returnClause: functionSignature.returnClause
-            ?? ReturnClauseSyntax(
-              type: IdentifierTypeSyntax(
-                name: .identifier("Void")
-              )
-            )
+          returnClause: returnClause
         )
       )
     }

diff --git a/Tests/SpyableMacroTests/Factories/UT_ClosureFactory.swift b/Tests/SpyableMacroTests/Factories/UT_ClosureFactory.swift
@@ -55,6 +55,22 @@ final class UT_ClosureFactory: XCTestCase {
     )
   }
 
+  func testVariableDeclarationOptionalTypeReturnValue() throws {
+    try assertProtocolFunction(
+      withFunctionDeclaration: "func _ignore_() -> Data?",
+      prefixForVariable: "_prefix_",
+      expectingVariableDeclaration: "var _prefix_Closure: (() -> Data? )?"
+    )
+  }
+
+  func testVariableDeclarationForcedUnwrappedOptionalTypeReturnValue() throws {
+    try assertProtocolFunction(
+      withFunctionDeclaration: "func _ignore_() -> Data!",
+      prefixForVariable: "_prefix_",
+      expectingVariableDeclaration: "var _prefix_Closure: (() -> Data?)?"
+    )
+  }
+
   func testVariableDeclarationEverything() throws {
     try assertProtocolFunction(
       withFunctionDeclaration: """

diff --git a/Tests/SpyableMacroTests/Factories/UT_ReturnValueFactory.swift b/Tests/SpyableMacroTests/Factories/UT_ReturnValueFactory.swift
@@ -23,6 +23,14 @@ final class UT_ReturnValueFactory: XCTestCase {
     )
   }
 
+  func testVariableDeclarationForcedUnwrappedType() throws {
+    try assert(
+      functionReturnType: "String!",
+      prefixForVariable: "_prefix_",
+      expectingVariableDeclaration: "var _prefix_ReturnValue: String!"
+    )
+  }
+
   func testVariableDeclarationExistentialType() throws {
     try assert(
       functionReturnType: "any Codable",

diff --git a/Tests/SpyableMacroTests/Factories/UT_SpyFactory.swift b/Tests/SpyableMacroTests/Factories/UT_SpyFactory.swift
@@ -347,6 +347,34 @@ final class UT_SpyFactory: XCTestCase {
     )
   }
 
+  func testDeclarationReturnsForcedUnwrappedType() throws {
+    try assertProtocol(
+      withDeclaration: """
+        protocol ServiceProtocol {
+            func foo() -> String!
+        }
+        """,
+      expectingClassDeclaration: """
+        class ServiceProtocolSpy: ServiceProtocol {
+            var fooCallsCount = 0
+            var fooCalled: Bool {
+                return fooCallsCount > 0
+            }
+            var fooReturnValue: String!
+            var fooClosure: (() -> String?)?
+            func foo() -> String! {
+                fooCallsCount += 1
+                if fooClosure != nil {
+                    return fooClosure!()
+                } else {
+                    return fooReturnValue
+                }
+            }
+        }
+        """
+    )
+  }
+
   func testDeclarationVariable() throws {
     try assertProtocol(
       withDeclaration: """