Fix issues with symbol import declarations

nicklockwood · Aug 23, 2020 · ad50c5d · ad50c5d
1 parent 1453b89
commit ad50c5d
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 2 deletions.
diff --git a/Sources/ParsingHelpers.swift b/Sources/ParsingHelpers.swift
@@ -1011,7 +1011,9 @@ extension Formatter {
                 where: { $0.isDeclarationTypeKeyword || $0 == .startOfScope("#if") }
             ) {
                 // Most declarations will include exactly one token that `isDeclarationTypeKeyword` in
-                // their outermost scope, but `class func` methods will have two (and the first one will be incorrect!)
+                //  - `class func` methods will have two (and the first one will be incorrect!)
+                //  - Symbol imports (like `import class Module.Type`) will have two,
+                //    but the first one will be correct, so we don't make any corrections right here.
                 if parser.tokens[firstDeclarationTypeKeywordIndex].string != "class" {
                     declarationTypeKeywordIndex = firstDeclarationTypeKeywordIndex
                     declarationKeyword = parser.tokens[firstDeclarationTypeKeywordIndex].string
@@ -1050,6 +1052,15 @@ extension Formatter {
                     searchIndex = endOfConditionalCompilationScope
                 }
 
+                // Symbol imports (like `import class Module.Type`) will have an extra `isDeclarationTypeKeyword`
+                // immediately following their `declarationKeyword`, so we need to skip them.
+                if declarationKeyword == "import",
+                    let symbolTypeKeywordIndex = parser.index(of: .nonSpaceOrComment, after: declarationTypeKeywordIndex),
+                    parser.tokens[symbolTypeKeywordIndex].isDeclarationTypeKeyword
+                {
+                    searchIndex = symbolTypeKeywordIndex + 1
+                }
+
                 while searchIndex < parser.tokens.count, nextDeclarationKeywordIndex == nil {
                     // If we encounter a `startOfScope`, we have to skip to the end of the scope.
                     // This accounts for things like function bodies, etc.
@@ -1391,7 +1402,8 @@ extension _FormatRules {
 extension Token {
     /// Whether or not this token "defines" the specific type of declaration
     ///  - A valid declaration will usually include exactly one of these keywords in its outermost scope.
-    ///  - A notable exception is `class func`, which will include two of these keywords.
+    ///  - Notable exceptions are `class func` and symbol imports (like `import class Module.Type`)
+    ///    which will include two of these keywords.
     var isDeclarationTypeKeyword: Bool {
         guard case let .keyword(keyword) = self else {
             return false

diff --git a/Tests/ParsingHelpersTests.swift b/Tests/ParsingHelpersTests.swift
@@ -1054,4 +1054,37 @@ class ParsingHelpersTests: XCTestCase {
         XCTAssertEqual(declarations[1].body?[2].keyword, "struct")
         XCTAssertEqual(declarations[2].keyword, "let")
     }
+
+    func testParseSymbolImportCorrectly() {
+        let input = """
+        import protocol SomeModule.SomeProtocol
+        import class SomeModule.SomeClass
+        import enum SomeModule.SomeEnum
+        import struct SomeModule.SomeStruct
+        import typealias SomeModule.SomeTypealias
+        import let SomeModule.SomeGlobalConstant
+        import var SomeModule.SomeGlobalVariable
+        import func SomeModule.SomeFunc
+
+        struct Foo {
+            init() {}
+            public func instanceMethod() {}
+        }
+        """
+
+        let originalTokens = tokenize(input)
+        let declarations = Formatter(originalTokens).parseDeclarations()
+
+        XCTAssertEqual(declarations[0].keyword, "import")
+        XCTAssertEqual(declarations[1].keyword, "import")
+        XCTAssertEqual(declarations[2].keyword, "import")
+        XCTAssertEqual(declarations[3].keyword, "import")
+        XCTAssertEqual(declarations[4].keyword, "import")
+        XCTAssertEqual(declarations[5].keyword, "import")
+        XCTAssertEqual(declarations[6].keyword, "import")
+        XCTAssertEqual(declarations[7].keyword, "import")
+        XCTAssertEqual(declarations[8].keyword, "struct")
+        XCTAssertEqual(declarations[8].body?[0].keyword, "init")
+        XCTAssertEqual(declarations[8].body?[1].keyword, "func")
+    }
 }
diff --git a/Tests/RulesTests.swift b/Tests/RulesTests.swift
@@ -13916,4 +13916,50 @@ class RulesTests: XCTestCase {
                        options: FormatOptions(ifdefIndent: .noIndent),
                        exclude: ["blankLinesAtStartOfScope", "blankLinesAtEndOfScope"])
     }
+
+    func testOrganizesTypeBelowSymbolImport() {
+        let input = """
+        import protocol SomeModule.SomeProtocol
+        import class SomeModule.SomeClass
+        import enum SomeModule.SomeEnum
+        import struct SomeModule.SomeStruct
+        import typealias SomeModule.SomeTypealias
+        import let SomeModule.SomeGlobalConstant
+        import var SomeModule.SomeGlobalVariable
+        import func SomeModule.SomeFunc
+
+        struct Foo {
+            init() {}
+            public func instanceMethod() {}
+        }
+        """
+
+        let output = """
+        import protocol SomeModule.SomeProtocol
+        import class SomeModule.SomeClass
+        import enum SomeModule.SomeEnum
+        import struct SomeModule.SomeStruct
+        import typealias SomeModule.SomeTypealias
+        import let SomeModule.SomeGlobalConstant
+        import var SomeModule.SomeGlobalVariable
+        import func SomeModule.SomeFunc
+
+        struct Foo {
+
+            // MARK: Lifecycle
+
+            init() {}
+
+            // MARK: Public
+
+            public func instanceMethod() {}
+
+        }
+        """
+
+        testFormatting(
+            for: input, output, rule: FormatRules.organizeDeclarations,
+            exclude: ["blankLinesAtStartOfScope", "blankLinesAtEndOfScope", "sortedImports"]
+        )
+    }
 }
diff --git a/Tests/XCTestManifests.swift b/Tests/XCTestManifests.swift
@@ -372,6 +372,7 @@ extension ParsingHelpersTests {
         ("testParseDeclarationsWithSituationalKeywords", testParseDeclarationsWithSituationalKeywords),
         ("testParseMarkCommentsCorrectly", testParseMarkCommentsCorrectly),
         ("testParseSimpleCompilationBlockCorrectly", testParseSimpleCompilationBlockCorrectly),
+        ("testParseSymbolImportCorrectly", testParseSymbolImportCorrectly),
         ("testParseTrailingCommentsCorrectly", testParseTrailingCommentsCorrectly),
         ("testProcessCaseDeclaredVariablesInIf", testProcessCaseDeclaredVariablesInIf),
         ("testProcessCaseDeclaredVariablesInIf2", testProcessCaseDeclaredVariablesInIf2),
@@ -1328,6 +1329,7 @@ extension RulesTests {
         ("testOrganizeEnumCasesFirst", testOrganizeEnumCasesFirst),
         ("testOrganizePrivateSet", testOrganizePrivateSet),
         ("testOrganizesNestedTypesWithinConditionalCompilationBlock", testOrganizesNestedTypesWithinConditionalCompilationBlock),
+        ("testOrganizesTypeBelowSymbolImport", testOrganizesTypeBelowSymbolImport),
         ("testOrganizesTypesBelowConditionalCompilationBlock", testOrganizesTypesBelowConditionalCompilationBlock),
         ("testOrganizesTypesWithinConditionalCompilationBlock", testOrganizesTypesWithinConditionalCompilationBlock),
         ("testOuterParensRemovedInIf", testOuterParensRemovedInIf),