Add support for optional return types

adincebic · Aug 4, 2023 · 3a582ee · 3a582ee
1 parent 7770ffa
commit 3a582ee
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git a/Sources/existentialannotator/Annotator.swift b/Sources/existentialannotator/Annotator.swift
@@ -126,17 +126,22 @@ final class Annotator: SyntaxRewriter {
             return node
         }
         let currentTypeWithoutAnyKeyword = node.returnType
-        let typeWithAnyKeyword = withExistentialAny(currentTypeWithoutAnyKeyword)
-            .withTrailingTrivia([.spaces(1)])
+
+        var typeWithAnyKeyword: TypeSyntax {
+            if let lastToken = currentTypeWithoutAnyKeyword.lastToken, lastToken.isOptionalNotation {
+                let typeString = "(\(withExistentialAny(currentTypeWithoutAnyKeyword.rawStringRepresentation)))\(lastToken.description)"
+                return TypeSyntax(stringLiteral: typeString)
+            }
+            return withExistentialAny(currentTypeWithoutAnyKeyword)
+                .withTrailingTrivia([.spaces(1)])
+        }
 
         let modifiedNode = node.withReturnType(typeWithAnyKeyword)
         return modifiedNode
     }
 
     private func isOptional(tokens: TokenSequence) -> (isOptional: Bool, optionalNotation: String) {
-        let token = tokens.first(where: {
-            $0.tokenKind == .postfixQuestionMark || $0.tokenKind == .exclamationMark
-        })
+        let token = tokens.first(where: { $0.isOptionalNotation })
         let notation = token?.description ?? ""
         let isOptional = !notation.isEmpty
         return (isOptional: isOptional, optionalNotation: notation)
@@ -166,3 +171,9 @@ extension VariableDeclSyntax {
         return description.contains("[\(type)]")
     }
 }
+
+private extension TokenSyntax {
+    var isOptionalNotation: Bool {
+        tokenKind == .postfixQuestionMark || tokenKind == .exclamationMark
+    }
+}
diff --git a/Tests/existentialannotatorTests/ExistentialAnnotatorTests.swift b/Tests/existentialannotatorTests/ExistentialAnnotatorTests.swift
@@ -511,7 +511,28 @@ final class ExistentialAnnotatorTests: XCTestCase {
         XCTAssertEqual(annotated.description, expected)
     }
 
-    func testThatExistentialIsAnnotatedWhenUsedInCollection() throws {
+    func testThatExistentialIsAnnotatedWhenUsedAsOptionalReturnType() throws {
+        let exampleFile = #"""
+            private func getValue(forKey key: String) -> NSCoding? {
+                return "some key"
+            }
+        """#
+
+        let sut = Annotator(protocols: ["NSCoding"])
+        let parsedSource = try SyntaxParser.parse(source: exampleFile)
+
+        let annotated = sut.visit(parsedSource)
+
+        let expected = #"""
+            private func getValue(forKey key: String) -> (any NSCoding)? {
+                return "some key"
+            }
+        """#
+
+        XCTAssertEqual(annotated.description, expected)
+    }
+
+    func testThatExistentialIsAnnotatedWhenUsedInArray() throws {
         let exampleFile = #"""
         struct MyType {
             let responses: [Decodable]