Add fixed source to assertMacroExpansion

Author: kimdv

Sources/SwiftSyntaxMacrosTestSupport/Assertions.swift

@@ -254,13 +254,18 @@ func assertDiagnostic(
 ///   - macros: The macros that should be expanded, provided as a dictionary
 ///     mapping macro names (e.g., `"stringify"`) to implementation types
 ///     (e.g., `StringifyMacro.self`).
+///   - applyFixIts: If specified, filters the Fix-Its that are applied to generate `fixedSource` to only those whose message occurs in this array. If `nil`, all Fix-Its from the diagnostics are applied.
+///   - fixedSource: If specified, asserts that the source code after applying Fix-Its matches this string.
 ///   - testModuleName: The name of the test module to use.
 ///   - testFileName: The name of the test file name to use.
+///   - indentationWidth: The indentation width used in the expansion.
 public func assertMacroExpansion(
   _ originalSource: String,
   expandedSource expectedExpandedSource: String,
   diagnostics: [DiagnosticSpec] = [],
   macros: [String: Macro.Type],
+  applyFixIts: [String]? = nil,
+  fixedSource expectedFixedSource: String? = nil,
   testModuleName: String = "TestModule",
   testFileName: String = "test.swift",
   indentationWidth: Trivia = .spaces(4),
@@ -316,4 +321,16 @@ public func assertMacroExpansion(
       assertDiagnostic(actualDiag, in: context, expected: expectedDiag)
     }
   }
+
+  // Applying Fix-Its
+  if let expectedFixedSource = expectedFixedSource {
+    let fixedTree = FixItApplier.applyFixes(in: context.diagnostics, withMessages: applyFixIts, to: expandedSourceFile)
+    let fixedTreeDescription = fixedTree.description
+    assertStringsEqualWithDiff(
+      fixedTreeDescription.trimmingTrailingWhitespace(),
+      expectedFixedSource.trimmingTrailingWhitespace(),
+      file: file,
+      line: line
+    )
+  }
 }

Sources/_SwiftSyntaxTestSupport/FixItApplier.swift

@@ -0,0 +1,70 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+import SwiftDiagnostics
+
+public class FixItApplier: SyntaxRewriter {
+  var changes: [FixIt.Change]
+
+  init(diagnostics: [Diagnostic], withMessages messages: [String]?) {
+    let messages = messages ?? diagnostics.compactMap { $0.fixIts.first?.message.message }
+
+    self.changes =
+      diagnostics
+      .flatMap { $0.fixIts }
+      .filter {
+        return messages.contains($0.message.message)
+      }
+      .flatMap { $0.changes }
+
+    super.init(viewMode: .all)
+  }
+
+  public override func visitAny(_ node: Syntax) -> Syntax? {
+    for change in changes {
+      switch change {
+      case .replace(oldNode: let oldNode, newNode: let newNode):
+        if oldNode.id == node.id {
+          return newNode
+        } else {
+          break
+        }
+      default:
+        break
+      }
+    }
+    return nil
+  }
+
+  public override func visit(_ node: TokenSyntax) -> TokenSyntax {
+    var modifiedNode = node
+    for change in changes {
+      switch change {
+      case .replaceLeadingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
+        modifiedNode = node.with(\.leadingTrivia, newTrivia)
+      case .replaceTrailingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
+        modifiedNode = node.with(\.trailingTrivia, newTrivia)
+      default:
+        break
+      }
+    }
+    return modifiedNode
+  }
+
+  /// If `messages` is `nil`, applies all Fix-Its in `diagnostics` to `tree` and returns the fixed syntax tree.
+  /// If `messages` is not `nil`, applies only Fix-Its whose message is in `messages`.
+  public static func applyFixes<T: SyntaxProtocol>(in diagnostics: [Diagnostic], withMessages messages: [String]?, to tree: T) -> Syntax {
+    let applier = FixItApplier(diagnostics: diagnostics, withMessages: messages)
+    return applier.rewrite(tree)
+  }
+}

Tests/SwiftParserTest/Assertions.swift

@@ -276,58 +276,6 @@ struct DiagnosticSpec {
   }
 }
 
-class FixItApplier: SyntaxRewriter {
-  var changes: [FixIt.Change]
-
-  init(diagnostics: [Diagnostic], withMessages messages: [String]?) {
-    let messages = messages ?? diagnostics.compactMap { $0.fixIts.first?.message.message }
-
-    self.changes =
-      diagnostics
-      .flatMap { $0.fixIts }
-      .filter {
-        return messages.contains($0.message.message)
-      }
-      .flatMap { $0.changes }
-
-    super.init(viewMode: .all)
-  }
-
-  public override func visitAny(_ node: Syntax) -> Syntax? {
-    for change in changes {
-      switch change {
-      case .replace(oldNode: let oldNode, newNode: let newNode) where oldNode.id == node.id:
-        return newNode
-      default:
-        break
-      }
-    }
-    return nil
-  }
-
-  override func visit(_ node: TokenSyntax) -> TokenSyntax {
-    var modifiedNode = node
-    for change in changes {
-      switch change {
-      case .replaceLeadingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
-        modifiedNode = node.with(\.leadingTrivia, newTrivia)
-      case .replaceTrailingTrivia(token: let changedNode, newTrivia: let newTrivia) where changedNode.id == node.id:
-        modifiedNode = node.with(\.trailingTrivia, newTrivia)
-      default:
-        break
-      }
-    }
-    return modifiedNode
-  }
-
-  /// If `messages` is `nil`, applies all Fix-Its in `diagnostics` to `tree` and returns the fixed syntax tree.
-  /// If `messages` is not `nil`, applies only Fix-Its whose message is in `messages`.
-  public static func applyFixes<T: SyntaxProtocol>(in diagnostics: [Diagnostic], withMessages messages: [String]?, to tree: T) -> Syntax {
-    let applier = FixItApplier(diagnostics: diagnostics, withMessages: messages)
-    return applier.rewrite(tree)
-  }
-}
-
 /// Assert that `location` is the same as that of `locationMarker` in `tree`.
 func assertLocation<T: SyntaxProtocol>(
   _ location: SourceLocation,

Tests/SwiftSyntaxMacroExpansionTest/MacroSystemTests.swift

@@ -284,9 +284,37 @@ public struct AddCompletionHandler: PeerMacro {
 
     // This only makes sense for async functions.
     if funcDecl.signature.effectSpecifiers?.asyncSpecifier == nil {
-      throw MacroExpansionErrorMessage(
-        "@addCompletionHandler requires an async function"
+      let newEffects: FunctionEffectSpecifiersSyntax
+      if let existingEffects = funcDecl.signature.effectSpecifiers {
+        newEffects = existingEffects.with(\.asyncSpecifier, .keyword(.async))
+      } else {
+        newEffects = FunctionEffectSpecifiersSyntax(asyncSpecifier: .keyword(.async))
+      }
+
+      let newSignature = funcDecl.signature.with(\.effectSpecifiers, newEffects)
+
+      let diag = Diagnostic(
+        node: Syntax(funcDecl.funcKeyword),
+        message: MacroExpansionErrorMessage(
+          "can only add a completion-handler variant to an 'async' function"
+        ),
+        fixIts: [
+          FixIt(
+            message: MacroExpansionFixItMessage(
+              "add 'async'"
+            ),
+            changes: [
+              FixIt.Change.replace(
+                oldNode: Syntax(funcDecl.signature),
+                newNode: Syntax(newSignature)
+              )
+            ]
+          )
+        ]
       )
+
+      context.diagnose(diag)
+      return []
     }
 
     // Form the completion handler parameter.
@@ -1092,6 +1120,33 @@ final class MacroSystemTests: XCTestCase {
     )
   }
 
+  func testAddCompletionHandlerWithNoAsync() {
+    assertMacroExpansion(
+      """
+      @addCompletionHandler
+      func f(a: Int, for b: String, _ value: Double) -> String { }
+      """,
+      expandedSource: """
+        func f(a: Int, for b: String, _ value: Double) -> String { }
+        """,
+      diagnostics: [
+        DiagnosticSpec(
+          message: "can only add a completion-handler variant to an 'async' function",
+          line: 2,
+          column: 1,
+          fixIts: [
+            FixItSpec(message: "add 'async'")
+          ]
+        )
+      ],
+      macros: testMacros,
+      fixedSource: """
+        func f(a: Int, for b: String, _ value: Double) async -> String { }
+        """,
+      indentationWidth: indentationWidth
+    )
+  }
+
   func testAddBackingStorage() {
     assertMacroExpansion(
       """
@@ -1563,5 +1618,4 @@ final class MacroSystemTests: XCTestCase {
       indentationWidth: indentationWidth
     )
   }
-
 }