Insert fileHeader comment after hashbang

Author: nicklockwood

Sources/Rules.swift

@@ -3984,8 +3984,17 @@ public struct _FormatRules {
             }
             header = string
         }
+        var start = 0
         var lastHeaderTokenIndex = -1
         if var startIndex = formatter.index(of: .nonSpaceOrLinebreak, after: -1) {
+            if formatter.tokens[startIndex] == .startOfScope("#!") {
+                guard let endIndex = formatter.index(of: .linebreak, after: startIndex) else {
+                    return
+                }
+                startIndex = formatter.index(of: .nonSpaceOrLinebreak, after: endIndex) ?? endIndex
+                start = startIndex
+                lastHeaderTokenIndex = startIndex - 1
+            }
             switch formatter.tokens[startIndex] {
             case .startOfScope("//"):
                 if case let .commentBody(body)? = formatter.next(.nonSpace, after: startIndex) {
@@ -4048,7 +4057,7 @@ public struct _FormatRules {
             }
         }
         if header.isEmpty {
-            formatter.removeTokens(in: 0 ..< lastHeaderTokenIndex + 1)
+            formatter.removeTokens(in: start ..< lastHeaderTokenIndex + 1)
             return
         }
         var headerTokens = tokenize(header)
@@ -4070,7 +4079,7 @@ public struct _FormatRules {
         }) {
             lastHeaderTokenIndex = index
         }
-        formatter.replaceTokens(in: 0 ..< lastHeaderTokenIndex + 1, with: headerTokens)
+        formatter.replaceTokens(in: start ..< lastHeaderTokenIndex + 1, with: headerTokens)
     }
 
     /// Strip redundant `.init` from type instantiations

Tests/RulesTests.swift

@@ -68,10 +68,12 @@ class RulesTests: XCTestCase {
         if input != output {
             XCTAssertEqual(try format(output, rules: rules, options: options),
                            output, file: file, line: line)
-            for rule in rules {
-                let disabled = "// swiftformat:disable \(rule.name)\n\(input)"
-                XCTAssertEqual(try format(disabled, rules: [rule], options: options),
-                               disabled, "Failed to disable \(rule.name) rule", file: file, line: line)
+            if !input.hasPrefix("#!") {
+                for rule in rules {
+                    let disabled = "// swiftformat:disable \(rule.name)\n\(input)"
+                    XCTAssertEqual(try format(disabled, rules: [rule], options: options),
+                                   disabled, "Failed to disable \(rule.name) rule", file: file, line: line)
+                }
             }
         }
         if input != output2, output != output2 {
@@ -529,6 +531,89 @@ class RulesTests: XCTestCase {
         testFormatting(for: input, output, rule: FormatRules.fileHeader, options: options)
     }
 
+    func testFileHeaderRemovedAfterHashbang() {
+        let input = """
+        #!/usr/bin/swift
+
+        // Header line1
+        // Header line2
+
+        let foo = 5
+        """
+        let output = """
+        #!/usr/bin/swift
+
+        let foo = 5
+        """
+        let options = FormatOptions(fileHeader: "")
+        testFormatting(for: input, output, rule: FormatRules.fileHeader, options: options)
+    }
+
+    func testFileHeaderPlacedAfterHashbang() {
+        let input = """
+        #!/usr/bin/swift
+
+        let foo = 5
+        """
+        let output = """
+        #!/usr/bin/swift
+
+        // Header line1
+        // Header line2
+
+        let foo = 5
+        """
+        let options = FormatOptions(fileHeader: "// Header line1\n// Header line2")
+        testFormatting(for: input, output, rule: FormatRules.fileHeader, options: options)
+    }
+
+    func testBlankLineAfterHashbangNotRemovedByFileHeader() {
+        let input = """
+        #!/usr/bin/swift
+
+        let foo = 5
+        """
+        let options = FormatOptions(fileHeader: "")
+        testFormatting(for: input, rule: FormatRules.fileHeader, options: options)
+    }
+
+    func testLineAfterHashbangNotAffectedByFileHeaderRemoval() {
+        let input = """
+        #!/usr/bin/swift
+        let foo = 5
+        """
+        let options = FormatOptions(fileHeader: "")
+        testFormatting(for: input, rule: FormatRules.fileHeader, options: options)
+    }
+
+    func testDisableFileHeaderCommentRespectedAfterHashbang() {
+        let input = """
+        #!/usr/bin/swift
+        // swiftformat:disable fileHeader
+
+        // Header line1
+        // Header line2
+
+        let foo = 5
+        """
+        let options = FormatOptions(fileHeader: "")
+        testFormatting(for: input, rule: FormatRules.fileHeader, options: options)
+    }
+
+    func testDisableFileHeaderCommentRespectedAfterHashbang2() {
+        let input = """
+        #!/usr/bin/swift
+
+        // swiftformat:disable fileHeader
+        // Header line1
+        // Header line2
+
+        let foo = 5
+        """
+        let options = FormatOptions(fileHeader: "")
+        testFormatting(for: input, rule: FormatRules.fileHeader, options: options)
+    }
+
     // MARK: - strongOutlets
 
     func testRemoveWeakFromOutlet() {

Tests/XCTestManifests.swift

@@ -542,6 +542,7 @@ extension RulesTests {
         ("testBelowCustomStructOrganizationThreshold", testBelowCustomStructOrganizationThreshold),
         ("testBinaryGroupingCustom", testBinaryGroupingCustom),
         ("testBlankCodeCommentBlockLinesNotIndented", testBlankCodeCommentBlockLinesNotIndented),
+        ("testBlankLineAfterHashbangNotRemovedByFileHeader", testBlankLineAfterHashbangNotRemovedByFileHeader),
         ("testBlankLineAfterProtocolBeforeProperty", testBlankLineAfterProtocolBeforeProperty),
         ("testBlankLineBeforeAtAvailabilityOnLineBeforeClass", testBlankLineBeforeAtAvailabilityOnLineBeforeClass),
         ("testBlankLineBeforeAtObjcOnLineBeforeProtocol", testBlankLineBeforeAtObjcOnLineBeforeProtocol),
@@ -696,6 +697,8 @@ extension RulesTests {
         ("testDictionaryLiteralNotEqualYodaCondition", testDictionaryLiteralNotEqualYodaCondition),
         ("testDictionaryLiteralsRuinEverything", testDictionaryLiteralsRuinEverything),
         ("testDictionaryTypeConvertedToSugar", testDictionaryTypeConvertedToSugar),
+        ("testDisableFileHeaderCommentRespectedAfterHashbang", testDisableFileHeaderCommentRespectedAfterHashbang),
+        ("testDisableFileHeaderCommentRespectedAfterHashbang2", testDisableFileHeaderCommentRespectedAfterHashbang2),
         ("testDisableNextRemoveSelf", testDisableNextRemoveSelf),
         ("testDisableRemoveSelf", testDisableRemoveSelf),
         ("testDispatchAsyncAfterClosureArgumentMadeTrailing", testDispatchAsyncAfterClosureArgumentMadeTrailing),
@@ -808,6 +811,8 @@ extension RulesTests {
         ("testFileHeaderCreationDateReplacement", testFileHeaderCreationDateReplacement),
         ("testFileHeaderCreationYearReplacement", testFileHeaderCreationYearReplacement),
         ("testFileHeaderFileReplacement", testFileHeaderFileReplacement),
+        ("testFileHeaderPlacedAfterHashbang", testFileHeaderPlacedAfterHashbang),
+        ("testFileHeaderRemovedAfterHashbang", testFileHeaderRemovedAfterHashbang),
         ("testFileHeaderRuleThrowsIfCreationDateUnavailable", testFileHeaderRuleThrowsIfCreationDateUnavailable),
         ("testFileHeaderRuleThrowsIfFileNameUnavailable", testFileHeaderRuleThrowsIfFileNameUnavailable),
         ("testFileHeaderYearReplacement", testFileHeaderYearReplacement),
@@ -1033,6 +1038,7 @@ extension RulesTests {
         ("testLetRedundantGenericTypeRemovalExplicitType", testLetRedundantGenericTypeRemovalExplicitType),
         ("testLetRedundantTypeRemoval", testLetRedundantTypeRemoval),
         ("testLetRedundantTypeRemovalExplicitType", testLetRedundantTypeRemovalExplicitType),
+        ("testLineAfterHashbangNotAffectedByFileHeaderRemoval", testLineAfterHashbangNotAffectedByFileHeaderRemoval),
         ("testLinebreakAtEndOfFile", testLinebreakAtEndOfFile),
         ("testLinebreakInsertedAtEndOfWrappedFunction", testLinebreakInsertedAtEndOfWrappedFunction),
         ("testLinebreaksRemovedInsideBraces", testLinebreaksRemovedInsideBraces),