Work correctly with functions

WhoopInc · Jan 12, 2024 · 9de9db1 · 9de9db1
1 parent c4ba07a
commit 9de9db1
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/Sources/WhoopDIKitMacros/InjectableMacro.swift b/Sources/WhoopDIKitMacros/InjectableMacro.swift
@@ -200,6 +200,11 @@ extension VariableDeclSyntax {
     }
 
     var typeName: TypeSyntax? {
-        self.bindings.first?.typeAnnotation?.type.trimmed
+        guard let annotationType = self.bindings.first?.typeAnnotation?.type.trimmed else { return nil }
+        if (annotationType.is(FunctionTypeSyntax.self)) {
+            return "@escaping \(annotationType)"
+        } else {
+            return annotationType
+        }
     }
 }
diff --git a/Tests/WhoopDIKitTests/InjectableTests.swift b/Tests/WhoopDIKitTests/InjectableTests.swift
@@ -105,4 +105,30 @@ final class InjectableTests: XCTestCase {
         """,
         macros: ["Injectable": InjectableMacro.self])
     }
+
+    func testInjectWithClosures() {
+        assertMacroExpansion(
+            """
+            @Injectable struct ClosureHolder {
+                let closure: () -> String
+            }
+            """,
+            expandedSource: """
+            struct ClosureHolder {
+                let closure: () -> String
+
+                internal static func inject() -> Self {
+                    Self.init(closure: WhoopDI.inject(nil))
+                }
+
+                internal init(closure: @escaping () -> String) {
+                    self.closure = closure
+                }
+            }
+            
+            extension ClosureHolder : Injectable {
+            }
+            """,
+            macros: ["Injectable": InjectableMacro.self])
+    }
 }