Add option to wrap function effects (throws, async) (#1301)

Rules.md

@@ -2177,6 +2177,7 @@ Option | Description
 `--wrapreturntype` | Wrap return type: "if-multiline", "preserve" (default)
 `--wrapconditions` | Wrap conditions: "before-first", "after-first", "preserve"
 `--wraptypealiases` | Wrap typealiases: "before-first", "after-first", "preserve"
+`--wrapeffects` | Wrap effects: "if-multiline", "never", "preserve"
 
 <details>
 <summary>Examples</summary>

Sources/FormattingHelpers.swift

@@ -290,29 +290,64 @@ extension Formatter {
             }
         }
 
-        func wrapReturnIfNecessary(
+        func wrapReturnAndEffectsIfNecessary(
             startOfScope: Int,
             endOfFunctionScope: Int
         ) {
-            switch options.wrapReturnType {
-            case .preserve:
-                break
-            case .ifMultiline:
-                guard token(at: startOfScope) == .startOfScope("("),
-                      let openBracket = index(of: .startOfScope, after: endOfFunctionScope),
-                      token(at: openBracket) == .startOfScope("{"),
-                      let returnArrowIndex = index(of: .operator("->", .infix), after: endOfFunctionScope),
-                      returnArrowIndex < openBracket
-                else { return }
+            guard token(at: startOfScope) == .startOfScope("("),
+                  let openBracket = index(of: .startOfScope, after: endOfFunctionScope),
+                  token(at: openBracket) == .startOfScope("{")
+            else { return }
 
-                // If the return arrow is on the same line as the closing paren, wrap it
-                if startOfLine(at: endOfFunctionScope) == startOfLine(at: returnArrowIndex) {
-                    insertSpace(indentForLine(at: returnArrowIndex), at: returnArrowIndex)
-                    insertLinebreak(at: returnArrowIndex)
+            func wrap(before index: Int) {
+                insertSpace(indentForLine(at: index), at: index)
+                insertLinebreak(at: index)
 
-                    // Remove any trailing whitespace that is now orphaned on the previous line
-                    if tokens[returnArrowIndex - 1].is(.space) {
-                        removeToken(at: returnArrowIndex - 1)
+                // Remove any trailing whitespace that is now orphaned on the previous line
+                if tokens[index - 1].is(.space) {
+                    removeToken(at: index - 1)
+                }
+            }
+
+            if let effectIndex = index(after: endOfFunctionScope, where: { $0.string == "throws" || $0.string == "async" }),
+               effectIndex < openBracket
+            {
+                switch options.wrapEffects {
+                case .preserve:
+                    break
+                case .ifMultiline:
+                    // If the effect is on the same line as the closing paren, wrap it
+                    if startOfLine(at: endOfFunctionScope) == startOfLine(at: effectIndex) {
+                        wrap(before: effectIndex)
+
+                        // When wrapping the effect, we should also un-wrap any return type
+                        if
+                            let returnArrowIndex = index(of: .operator("->", .infix), after: endOfFunctionScope),
+                            returnArrowIndex < openBracket,
+                            let tokenBeforeArrowIndex = index(of: .nonSpaceOrCommentOrLinebreak, before: returnArrowIndex),
+                            startOfLine(at: tokenBeforeArrowIndex) != startOfLine(at: returnArrowIndex)
+                        {
+                            replaceTokens(in: endOfLine(at: tokenBeforeArrowIndex) ..< returnArrowIndex, with: [.space(" ")])
+                        }
+                    }
+                case .never:
+                    if startOfLine(at: endOfFunctionScope) != startOfLine(at: effectIndex) {
+                        replaceTokens(in: endOfLine(at: endOfFunctionScope) ..< effectIndex, with: [.space(" ")])
+                    }
+                }
+            }
+
+            if
+                let returnArrowIndex = index(of: .operator("->", .infix), after: endOfFunctionScope),
+                returnArrowIndex < openBracket
+            {
+                switch options.wrapReturnType {
+                case .preserve:
+                    break
+                case .ifMultiline:
+                    // If the return arrow is on the same line as the closing paren, wrap it
+                    if startOfLine(at: endOfFunctionScope) == startOfLine(at: returnArrowIndex) {
+                        wrap(before: returnArrowIndex)
                     }
                 }
             }
@@ -384,7 +419,7 @@ extension Formatter {
                 }
             }
 
-            wrapReturnIfNecessary(
+            wrapReturnAndEffectsIfNecessary(
                 startOfScope: i,
                 endOfFunctionScope: endOfScope
             )
@@ -439,7 +474,7 @@ extension Formatter {
                 insertLinebreak(at: breakIndex)
             }
 
-            wrapReturnIfNecessary(
+            wrapReturnAndEffectsIfNecessary(
                 startOfScope: i,
                 endOfFunctionScope: endOfScope
             )

Sources/OptionDescriptor.swift

@@ -470,6 +470,12 @@ struct _Descriptors {
         help: "Wrap return type: \"if-multiline\", \"preserve\" (default)",
         keyPath: \.wrapReturnType
     )
+    let wrapEffects = OptionDescriptor(
+        argumentName: "wrapeffects",
+        displayName: "Wrap Function Effects (throws, async)",
+        help: "Wrap effects: \"if-multiline\", \"never\", \"preserve\"",
+        keyPath: \.wrapEffects
+    )
     let wrapConditions = OptionDescriptor(
         argumentName: "wrapconditions",
         displayName: "Wrap Conditions",

Sources/Options.swift

@@ -110,6 +110,16 @@ public enum WrapReturnType: String, CaseIterable {
     case preserve
 }
 
+/// Wrapping behavior for effects (`async`, `throws`)
+public enum WrapEffects: String, CaseIterable {
+    case preserve
+    /// `async` and `throws` are wrapped to the line after the closing paren
+    /// if the function spans multiple lines
+    case ifMultiline = "if-multiline"
+    /// `async` and `throws` are never wrapped, and are always included on the same line as the closing paren
+    case never
+}
+
 /// Annotation which should be kept when removing a redundant type
 public enum RedundantType: String, CaseIterable {
     /// Preserves the type as a part of the property definition:
@@ -413,6 +423,7 @@ public struct FormatOptions: CustomStringConvertible {
     public var removeStartOrEndBlankLinesFromTypes: Bool
     public var genericTypes: String
     public var useSomeAny: Bool
+    public var wrapEffects: WrapEffects
 
     // Deprecated
     public var indentComments: Bool
@@ -508,6 +519,7 @@ public struct FormatOptions: CustomStringConvertible {
                 removeStartOrEndBlankLinesFromTypes: Bool = true,
                 genericTypes: String = "",
                 useSomeAny: Bool = true,
+                wrapEffects: WrapEffects = .preserve,
                 // Doesn't really belong here, but hard to put elsewhere
                 fragment: Bool = false,
                 ignoreConflictMarkers: Bool = false,
@@ -594,6 +606,7 @@ public struct FormatOptions: CustomStringConvertible {
         self.removeStartOrEndBlankLinesFromTypes = removeStartOrEndBlankLinesFromTypes
         self.genericTypes = genericTypes
         self.useSomeAny = useSomeAny
+        self.wrapEffects = wrapEffects
         // Doesn't really belong here, but hard to put elsewhere
         self.fragment = fragment
         self.ignoreConflictMarkers = ignoreConflictMarkers

Sources/Rules.swift

@@ -4172,7 +4172,7 @@ public struct _FormatRules {
         help: "Wrap lines that exceed the specified maximum width.",
         options: ["maxwidth", "nowrapoperators", "assetliterals", "wrapternary"],
         sharedOptions: ["wraparguments", "wrapparameters", "wrapcollections", "closingparen", "indent",
-                        "trimwhitespace", "linebreaks", "tabwidth", "maxwidth", "smarttabs", "wrapreturntype", "wrapconditions", "wraptypealiases", "wrapternary"]
+                        "trimwhitespace", "linebreaks", "tabwidth", "maxwidth", "smarttabs", "wrapreturntype", "wrapconditions", "wraptypealiases", "wrapternary", "wrapeffects"]
     ) { formatter in
         let maxWidth = formatter.options.maxWidth
         guard maxWidth > 0 else { return }
@@ -4229,7 +4229,7 @@ public struct _FormatRules {
         help: "Align wrapped function arguments or collection elements.",
         orderAfter: ["wrap"],
         options: ["wraparguments", "wrapparameters", "wrapcollections", "closingparen",
-                  "wrapreturntype", "wrapconditions", "wraptypealiases"],
+                  "wrapreturntype", "wrapconditions", "wraptypealiases", "wrapeffects"],
         sharedOptions: ["indent", "trimwhitespace", "linebreaks",
                         "tabwidth", "maxwidth", "smarttabs", "assetliterals", "wrapternary"]
     ) { formatter in

Tests/MetadataTests.swift

@@ -177,8 +177,8 @@ class MetadataTests: XCTestCase {
                     referencedOptions += [
                         Descriptors.wrapArguments, Descriptors.wrapParameters, Descriptors.wrapCollections,
                         Descriptors.closingParenOnSameLine, Descriptors.linebreak, Descriptors.truncateBlankLines,
-                        Descriptors.indent, Descriptors.tabWidth, Descriptors.smartTabs,
-                        Descriptors.maxWidth, Descriptors.assetLiteralWidth, Descriptors.wrapReturnType,
+                        Descriptors.indent, Descriptors.tabWidth, Descriptors.smartTabs, Descriptors.maxWidth,
+                        Descriptors.assetLiteralWidth, Descriptors.wrapReturnType, Descriptors.wrapEffects,
                         Descriptors.wrapConditions, Descriptors.wrapTypealiases, Descriptors.wrapTernaryOperators,
                     ]
                 case .identifier("indexWhereLineShouldWrapInLine"), .identifier("indexWhereLineShouldWrap"):

Tests/RulesTests+Wrapping.swift

@@ -2719,6 +2719,95 @@ class WrappingTests: RulesTests {
         testFormatting(for: input, output, rule: FormatRules.wrapArguments, options: options)
     }
 
+    func testWrapReturnAndEffectOnMultilineFunctionDeclaration() {
+        let input = """
+        func multilineFunction(
+            foo _: String,
+            bar _: String) async throws -> String {}
+        """
+
+        let output = """
+        func multilineFunction(
+            foo _: String,
+            bar _: String)
+            async throws -> String {}
+        """
+
+        let options = FormatOptions(
+            wrapArguments: .beforeFirst,
+            closingParenOnSameLine: true,
+            wrapReturnType: .ifMultiline,
+            wrapEffects: .ifMultiline
+        )
+
+        testFormatting(for: input, output, rule: FormatRules.wrapArguments, options: options)
+    }
+
+    func testDoesntWrapReturnAndEffectOnSingleLineFunctionDeclaration() {
+        let input = """
+        func singleLineFunction() async throws -> String {}
+        """
+
+        let options = FormatOptions(
+            wrapArguments: .beforeFirst,
+            closingParenOnSameLine: true,
+            wrapReturnType: .ifMultiline,
+            wrapEffects: .ifMultiline
+        )
+
+        testFormatting(for: input, rule: FormatRules.wrapArguments, options: options)
+    }
+
+    func testWrapEffectOnMultilineFunctionDeclaration() {
+        let input = """
+        func multilineFunction(
+            foo _: String,
+            bar _: String) async throws
+            -> String {}
+        """
+
+        let output = """
+        func multilineFunction(
+            foo _: String,
+            bar _: String)
+            async throws -> String {}
+        """
+
+        let options = FormatOptions(
+            wrapArguments: .beforeFirst,
+            closingParenOnSameLine: true,
+            wrapReturnType: .ifMultiline,
+            wrapEffects: .ifMultiline
+        )
+
+        testFormatting(for: input, output, rule: FormatRules.wrapArguments, options: options)
+    }
+
+    func testUnwrapEffectOnMultilineFunctionDeclaration() {
+        let input = """
+        func multilineFunction(
+            foo _: String,
+            bar _: String)
+            async throws -> String {}
+        """
+
+        let output = """
+        func multilineFunction(
+            foo _: String,
+            bar _: String) async throws
+            -> String {}
+        """
+
+        let options = FormatOptions(
+            wrapArguments: .beforeFirst,
+            closingParenOnSameLine: true,
+            wrapReturnType: .ifMultiline,
+            wrapEffects: .never
+        )
+
+        testFormatting(for: input, output, rule: FormatRules.wrapArguments, options: options)
+    }
+
     func testWrapReturnOnMultilineFunctionDeclarationWithAfterFirst() {
         let input = """
         func multilineFunction(foo _: String,
@@ -2767,6 +2856,31 @@ class WrappingTests: RulesTests {
         )
     }
 
+    func testWrapReturnAndEffectOnMultilineThrowingFunctionDeclarationWithAfterFirst() {
+        let input = """
+        func multilineFunction(foo _: String,
+                               bar _: String) throws -> String {}
+        """
+
+        let output = """
+        func multilineFunction(foo _: String,
+                               bar _: String)
+                               throws -> String {}
+        """
+
+        let options = FormatOptions(
+            wrapArguments: .afterFirst,
+            closingParenOnSameLine: true,
+            wrapReturnType: .ifMultiline,
+            wrapEffects: .ifMultiline
+        )
+
+        testFormatting(
+            for: input, output, rule: FormatRules.wrapArguments, options: options,
+            exclude: ["indent"]
+        )
+    }
+
     func testDoesntWrapReturnOnMultilineThrowingFunction() {
         let input = """
         func multilineFunction(foo _: String,
@@ -3278,6 +3392,35 @@ class WrappingTests: RulesTests {
         ], options: options)
     }
 
+    func testWrapsMultilineStatementConsistently2_withEffects() {
+        let input = """
+        func aFunc(
+            one _: Int,
+            two _: Int) async throws -> String {
+            "one"
+        }
+        """
+
+        let output = """
+        func aFunc(
+            one _: Int,
+            two _: Int
+        ) async throws -> String {
+            "one"
+        }
+        """
+
+        let options = FormatOptions(
+            wrapArguments: .beforeFirst,
+            closingParenOnSameLine: false,
+            wrapEffects: .never
+        )
+        testFormatting(for: input, [output], rules: [
+            FormatRules.wrapMultilineStatementBraces,
+            FormatRules.wrapArguments,
+        ], options: options)
+    }
+
     func testWrapsMultilineStatementConsistently3() {
         let input = """
         func aFunc(