Add organizeDeclarations support for extensions

Rules.md

@@ -675,9 +675,11 @@ Option | Description
 `--categorymark` | Template for category mark comments. Defaults to "MARK: %c"
 `--beforemarks` | Declarations placed before first mark (e.g. `typealias,struct`)
 `--lifecycle` | Names of additional Lifecycle methods (e.g. `viewDidLoad`)
+`--organizetypes` | Declarations to organize (defaults to `struct,class,enum`)
 `--structthreshold` | Minimum line count to organize struct body. Defaults to 0
 `--classthreshold` | Minimum line count to organize class body. Defaults to 0
 `--enumthreshold` | Minimum line count to organize enum body. Defaults to 0
+`--extensionlength` | Minimum line count to organize extension body. Defaults to 0
 
 <details>
 <summary>Examples</summary>

Sources/FormattingHelpers.swift

@@ -797,7 +797,7 @@ private extension Formatter {
         _ typeDeclaration: (kind: String, open: [Token], body: [Formatter.Declaration], close: [Token])
     ) -> (kind: String, open: [Token], body: [Formatter.Declaration], close: [Token]) {
         // Only organize the body of classes, structs, and enums (not protocols and extensions)
-        guard ["class", "struct", "enum"].contains(typeDeclaration.kind) else {
+        guard options.organizeTypes.contains(typeDeclaration.kind) else {
             return typeDeclaration
         }
 
@@ -810,6 +810,8 @@ private extension Formatter {
             organizationThreshold = options.organizeStructThreshold
         case "enum":
             organizationThreshold = options.organizeEnumThreshold
+        case "extension":
+            organizationThreshold = options.organizeExtensionThreshold
         default:
             organizationThreshold = 0
         }

Sources/OptionDescriptor.swift

@@ -664,6 +664,12 @@ struct _Descriptors {
         help: "Names of additional Lifecycle methods (e.g. `viewDidLoad`)",
         keyPath: \.lifecycleMethods
     )
+    let organizeTypes = OptionDescriptor(
+        argumentName: "organizetypes",
+        displayName: "Declaration Types to Organize",
+        help: "Declarations to organize (defaults to `struct,class,enum`)",
+        keyPath: \.organizeTypes
+    )
     let organizeStructThreshold = OptionDescriptor(
         argumentName: "structthreshold",
         displayName: "Organize Struct Threshold",
@@ -682,6 +688,12 @@ struct _Descriptors {
         help: "Minimum line count to organize enum body. Defaults to 0",
         keyPath: \.organizeEnumThreshold
     )
+    let organizeExtensionThreshold = OptionDescriptor(
+        argumentName: "extensionlength",
+        displayName: "Organize Extension Threshold",
+        help: "Minimum line count to organize extension body. Defaults to 0",
+        keyPath: \.organizeExtensionThreshold
+    )
     let funcAttributes = OptionDescriptor(
         argumentName: "funcattributes",
         displayName: "Function Attributes",

Sources/Options.swift

@@ -312,9 +312,11 @@ public struct FormatOptions: CustomStringConvertible {
     public var categoryMarkComment: String
     public var beforeMarks: Set<String>
     public var lifecycleMethods: Set<String>
+    public var organizeTypes: Set<String>
     public var organizeClassThreshold: Int
     public var organizeStructThreshold: Int
     public var organizeEnumThreshold: Int
+    public var organizeExtensionThreshold: Int
     public var yodaSwap: YodaMode
 
     // Deprecated
@@ -378,9 +380,11 @@ public struct FormatOptions: CustomStringConvertible {
                 categoryMarkComment: String = "MARK: %c",
                 beforeMarks: Set<String> = [],
                 lifecycleMethods: Set<String> = [],
+                organizeTypes: Set<String> = ["class", "struct", "enum"],
                 organizeClassThreshold: Int = 0,
                 organizeStructThreshold: Int = 0,
                 organizeEnumThreshold: Int = 0,
+                organizeExtensionThreshold: Int = 0,
                 yodaSwap: YodaMode = .always,
                 // Doesn't really belong here, but hard to put elsewhere
                 fragment: Bool = false,
@@ -438,9 +442,11 @@ public struct FormatOptions: CustomStringConvertible {
         self.categoryMarkComment = categoryMarkComment
         self.beforeMarks = beforeMarks
         self.lifecycleMethods = lifecycleMethods
+        self.organizeTypes = organizeTypes
         self.organizeClassThreshold = organizeClassThreshold
         self.organizeStructThreshold = organizeStructThreshold
         self.organizeEnumThreshold = organizeEnumThreshold
+        self.organizeExtensionThreshold = organizeExtensionThreshold
         self.yodaSwap = yodaSwap
         // Doesn't really belong here, but hard to put elsewhere
         self.fragment = fragment

Sources/Rules.swift

@@ -4912,8 +4912,8 @@ public struct _FormatRules {
         help: "Organizes declarations within class, struct, and enum bodies.",
         runOnceOnly: true,
         disabledByDefault: true,
-        options: ["categorymark", "beforemarks", "lifecycle", "structthreshold",
-                  "classthreshold", "enumthreshold"]
+        options: ["categorymark", "beforemarks", "lifecycle", "organizetypes",
+                  "structthreshold", "classthreshold", "enumthreshold", "extensionlength"]
     ) { formatter in
         // Parse the file into declarations and organize the body of individual types
         let organizedDeclarations = formatter

Tests/MetadataTests.swift

@@ -158,9 +158,11 @@ class MetadataTests: XCTestCase {
                         Descriptors.categoryMarkComment,
                         Descriptors.beforeMarks,
                         Descriptors.lifecycleMethods,
+                        Descriptors.organizeTypes,
                         Descriptors.organizeStructThreshold,
                         Descriptors.organizeClassThreshold,
                         Descriptors.organizeEnumThreshold,
+                        Descriptors.organizeExtensionThreshold,
                     ]
                 default:
                     continue

Tests/RulesTests+Organization.swift

@@ -517,6 +517,59 @@ extension RulesTests {
         )
     }
 
+    func testBelowCustomExtensionOrganizationThreshold() {
+        let input = """
+        extension FooBelowThreshold {
+            func bar() {}
+        }
+        """
+
+        testFormatting(
+            for: input,
+            rule: FormatRules.organizeDeclarations,
+            options: FormatOptions(
+                organizeTypes: ["class", "struct", "enum", "extension"],
+                organizeExtensionThreshold: 2
+            )
+        )
+    }
+
+    func testAboveCustomExtensionOrganizationThreshold() {
+        let input = """
+        extension FooBelowThreshold {
+            public func bar() {}
+            func baaz() {}
+            private func quux() {}
+        }
+        """
+
+        let output = """
+        extension FooBelowThreshold {
+
+            // MARK: Public
+
+            public func bar() {}
+
+            // MARK: Internal
+
+            func baaz() {}
+
+            // MARK: Private
+
+            private func quux() {}
+        }
+        """
+
+        testFormatting(
+            for: input, output,
+            rule: FormatRules.organizeDeclarations,
+            options: FormatOptions(
+                organizeTypes: ["class", "struct", "enum", "extension"],
+                organizeExtensionThreshold: 2
+            ), exclude: ["blankLinesAtStartOfScope"]
+        )
+    }
+
     func testPreservesExistingMarks() {
         let input = """
         class Foo {

Tests/XCTestManifests.swift

@@ -434,6 +434,7 @@ extension RulesTests {
     //   `swift test --generate-linuxmain`
     // to regenerate.
     static let __allTests__RulesTests = [
+        ("testAboveCustomExtensionOrganizationThreshold", testAboveCustomExtensionOrganizationThreshold),
         ("testAboveCustomStructOrganizationThreshold", testAboveCustomStructOrganizationThreshold),
         ("testAddSpaceAfterFuncEquals", testAddSpaceAfterFuncEquals),
         ("testAddSpaceAfterOperatorEquals", testAddSpaceAfterOperatorEquals),
@@ -495,6 +496,7 @@ extension RulesTests {
         ("testBeforeFirstPreservedAndTrailingCommaAddedInSingleLineNestedDictionaryWithOneNestedItem", testBeforeFirstPreservedAndTrailingCommaAddedInSingleLineNestedDictionaryWithOneNestedItem),
         ("testBeforeFirstPreservedIndentFixed", testBeforeFirstPreservedIndentFixed),
         ("testBeforeFirstPreservedNewlineAdded", testBeforeFirstPreservedNewlineAdded),
+        ("testBelowCustomExtensionOrganizationThreshold", testBelowCustomExtensionOrganizationThreshold),
         ("testBelowCustomStructOrganizationThreshold", testBelowCustomStructOrganizationThreshold),
         ("testBinaryGroupingCustom", testBinaryGroupingCustom),
         ("testBlankCodeCommentBlockLinesNotIndented", testBlankCodeCommentBlockLinesNotIndented),