From c1201e2ab55e98305a9db6cc1ba2de7b8e57d2d0 Mon Sep 17 00:00:00 2001 From: Dale Myers Date: Sat, 14 Mar 2026 10:41:26 +0000 Subject: [PATCH] Add import extraction support to Structure --- Source/SourceKittenFramework/ImportInfo.swift | 168 +++++++++++++ Source/SourceKittenFramework/Structure.swift | 28 ++- Source/sourcekitten/Structure.swift | 6 +- .../StructureTests.swift | 229 ++++++++++++++++++ 4 files changed, 427 insertions(+), 4 deletions(-) create mode 100644 Source/SourceKittenFramework/ImportInfo.swift diff --git a/Source/SourceKittenFramework/ImportInfo.swift b/Source/SourceKittenFramework/ImportInfo.swift new file mode 100644 index 000000000..fb49b30e5 --- /dev/null +++ b/Source/SourceKittenFramework/ImportInfo.swift @@ -0,0 +1,168 @@ +/// An attribute on an import statement (e.g. `@testable`, `@_exported`). +public struct ImportAttribute: Equatable { + /// The SourceKit attribute key (e.g. "source.decl.attribute.testable"). + public let kind: String + /// Byte offset of the attribute in the source file. + public let offset: ByteCount + /// Byte length of the attribute. + public let length: ByteCount +} + +/// Represents an import statement extracted from a Swift source file's syntax map. +public struct ImportInfo: Equatable { + /// Full import path (e.g. "Foundation", "Foundation.NSObject", "Darwin.exit"). + public let name: String + /// Byte offset of the import statement (including any preceding attributes like @testable). + public let offset: ByteCount + /// Byte length of the full import statement. + public let length: ByteCount + /// Attributes on the import (e.g. `@testable`, `@_exported`). + public let attributes: [ImportAttribute] + /// The declaration kind keyword for kind-qualified imports (e.g. "class", "struct", "func"). + /// Nil for plain imports like `import Foundation`. + public let importKind: String? + /// The module name (first path component) for dotted imports. + /// Nil for plain module imports like `import Foundation`. + public let moduleName: String? + + /// Dictionary representation using SourceKit key conventions. + public var dictionaryRepresentation: [String: SourceKitRepresentable] { + var dict: [String: SourceKitRepresentable] = [ + "key.name": name, + "key.offset": Int64(offset.value), + "key.length": Int64(length.value) + ] + if !attributes.isEmpty { + dict["key.attributes"] = attributes.map { attribute -> SourceKitRepresentable in + [ + "key.attribute": attribute.kind, + "key.offset": Int64(attribute.offset.value), + "key.length": Int64(attribute.length.value) + ] as [String: SourceKitRepresentable] + } as [SourceKitRepresentable] + } + if let importKind = importKind { + dict["key.import_kind"] = importKind + } + if let moduleName = moduleName { + dict["key.module_name"] = moduleName + } + return dict + } + + /// Extract import information from a syntax map and source file. + /// + /// Walks the syntax map tokens looking for `import` keywords, then collects + /// subsequent identifier tokens on the same line to determine the module name + /// and full import path. Handles `@testable` and `@_exported` attributes + /// preceding the import keyword, as well as kind-qualified imports like + /// `import class Foundation.NSObject`. + /// + /// - Parameters: + /// - syntaxMap: The syntax map from a SourceKit `editor.open` response. + /// - file: The source file, used to extract text from byte ranges. + /// - Returns: Array of `ImportInfo` for each import statement found. + public static func extractImports(from syntaxMap: SyntaxMap, in file: File) -> [ImportInfo] { + let tokens = syntaxMap.tokens + var imports = [ImportInfo]() + var index = 0 + + while index < tokens.count { + let token = tokens[index] + + // Look for keyword tokens that are "import" + guard token.type == SyntaxKind.keyword.rawValue, + file.stringView.substringWithByteRange(token.range) == "import" else { + index += 1 + continue + } + + // Collect preceding attribute tokens (e.g. @testable, @_exported) + var collectedAttributes = [ImportAttribute]() + var attributeStart: ByteCount? + var lookback = index - 1 + while lookback >= 0 { + let prev = tokens[lookback] + guard prev.type == SyntaxKind.attributeBuiltin.rawValue || + prev.type == SyntaxKind.attributeID.rawValue, + let attrText = file.stringView.substringWithByteRange(prev.range), + attrText.hasPrefix("@") else { + break + } + let sourceKitKey = "source.decl.attribute.\(attrText.dropFirst())" + guard let kind = SwiftDeclarationAttributeKind(rawValue: sourceKitKey) else { + break + } + attributeStart = prev.offset + collectedAttributes.append(ImportAttribute(kind: kind.rawValue, offset: prev.offset, length: prev.length)) + lookback -= 1 + } + // Reverse so attributes appear in source order + collectedAttributes.reverse() + + let importKeywordEnd = token.offset + token.length + let statementStart = attributeStart ?? token.offset + + // Collect subsequent tokens on the same line to build the full import path + var identifiers = [String]() + var importKind: String? + var lastTokenEnd = importKeywordEnd + var nextIndex = index + 1 + + while nextIndex < tokens.count { + let nextToken = tokens[nextIndex] + + // Check that the gap between the previous token end and this token doesn't contain a newline + let gapStart = lastTokenEnd + let gapLength = nextToken.offset - gapStart + if gapLength > ByteCount(0) { + let gapRange = ByteRange(location: gapStart, length: gapLength) + if let gapText = file.stringView.substringWithByteRange(gapRange), + gapText.contains("\n") { + break + } + } + + if nextToken.type == SyntaxKind.identifier.rawValue || + nextToken.type == SyntaxKind.typeidentifier.rawValue { + if let text = file.stringView.substringWithByteRange(nextToken.range) { + identifiers.append(text) + } + lastTokenEnd = nextToken.offset + nextToken.length + } else if nextToken.type == SyntaxKind.keyword.rawValue { + // The kind keyword (class/struct/func/enum/protocol/typealias/var/let) + // e.g. `import class Foundation.NSObject` + if let text = file.stringView.substringWithByteRange(nextToken.range) { + importKind = text + } + lastTokenEnd = nextToken.offset + nextToken.length + } else if nextToken.type == SyntaxKind.operator.rawValue { + // Dot operator in submodule imports like `import Foundation.NSObject` + lastTokenEnd = nextToken.offset + nextToken.length + } else { + break + } + + nextIndex += 1 + } + + if !identifiers.isEmpty { + let fullPath = identifiers.joined(separator: ".") + let moduleName: String? = identifiers.count > 1 ? identifiers[0] : nil + let totalLength = lastTokenEnd - statementStart + imports.append(ImportInfo( + name: fullPath, + offset: statementStart, + length: totalLength, + attributes: collectedAttributes, + importKind: importKind, + moduleName: moduleName + )) + } + + index = nextIndex + } + + return imports + } +} diff --git a/Source/SourceKittenFramework/Structure.swift b/Source/SourceKittenFramework/Structure.swift index a9701e827..60ecf54e4 100644 --- a/Source/SourceKittenFramework/Structure.swift +++ b/Source/SourceKittenFramework/Structure.swift @@ -14,14 +14,38 @@ public struct Structure { dictionary = sourceKitResponse } + /** + Create a Structure from a SourceKit `editor.open` response, extracting import + information from the syntax map before discarding it. + + - parameter sourceKitResponse: SourceKit `editor.open` response. + - parameter file: The source file, used to resolve import names from byte ranges. + */ + public init(sourceKitResponse: [String: SourceKitRepresentable], file: File, extractImports: Bool = false) { + var sourceKitResponse = sourceKitResponse + if extractImports, let syntaxMapData = SwiftDocKey.getSyntaxMap(sourceKitResponse) { + let syntaxMap = SyntaxMap(data: syntaxMapData) + let imports = ImportInfo.extractImports(from: syntaxMap, in: file) + if !imports.isEmpty { + sourceKitResponse["key.imports"] = imports.map { $0.dictionaryRepresentation } as [SourceKitRepresentable] + } else { + sourceKitResponse["key.imports"] = [SourceKitRepresentable]() + } + } + _ = sourceKitResponse.removeValue(forKey: SwiftDocKey.syntaxMap.rawValue) + dictionary = sourceKitResponse + } + /** Initialize a Structure by passing in a File. - parameter file: File to parse for structural information. + - parameter extractImports: Whether to extract import information from the syntax map. Defaults to `false`. - throws: Request.Error */ - public init(file: File) throws { - self.init(sourceKitResponse: try Request.editorOpen(file: file).send()) + public init(file: File, extractImports: Bool = false) throws { + let response = try Request.editorOpen(file: file).send() + self.init(sourceKitResponse: response, file: file, extractImports: extractImports) } } diff --git a/Source/sourcekitten/Structure.swift b/Source/sourcekitten/Structure.swift index ecb8a1d9a..74ebf594f 100644 --- a/Source/sourcekitten/Structure.swift +++ b/Source/sourcekitten/Structure.swift @@ -9,17 +9,19 @@ extension SourceKitten { var file: String = "" @Option(help: "Swift code text to parse") var text: String = "" + @Flag(help: "Include import statements in the output") + var imports: Bool = false mutating func run() throws { if !file.isEmpty { if let file = File(path: file) { - print(try SourceKittenFramework.Structure(file: file)) + print(try SourceKittenFramework.Structure(file: file, extractImports: imports)) return } throw SourceKittenError.readFailed(path: file) } if !text.isEmpty { - print(try SourceKittenFramework.Structure(file: File(contents: text))) + print(try SourceKittenFramework.Structure(file: File(contents: text), extractImports: imports)) return } throw SourceKittenError.invalidArgument( diff --git a/Tests/SourceKittenFrameworkTests/StructureTests.swift b/Tests/SourceKittenFrameworkTests/StructureTests.swift index a9c682e8f..c02331ecf 100644 --- a/Tests/SourceKittenFrameworkTests/StructureTests.swift +++ b/Tests/SourceKittenFrameworkTests/StructureTests.swift @@ -177,6 +177,235 @@ class StructureTests: XCTestCase { XCTAssertEqual(toNSDictionary(structure.dictionary), expected, "should generate expected structure") } + func testStructureOmitsImportsByDefault() throws { + let structure = try Structure(file: File(contents: "import Foundation\nclass Foo {}")) + XCTAssertNil(structure.dictionary["key.imports"], "key.imports should be absent when extractImports is not set") + } + + func testStructureIncludesImports() throws { + let structure = try Structure(file: File(contents: "import Foundation\nclass Foo {}"), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation") + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + XCTAssertEqual(first?["key.length"] as? Int64, 17) + } + + func testStructureMultipleImports() throws { + let source = "import Foundation\nimport UIKit\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 2) + let first = imports?[0] as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation") + let second = imports?[1] as? [String: SourceKitRepresentable] + XCTAssertEqual(second?["key.name"] as? String, "UIKit") + } + + func testStructureNoImports() throws { + let structure = try Structure(file: File(contents: "class Foo {}"), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports, "key.imports should be present when extractImports is true") + XCTAssertEqual(imports?.count, 0, "key.imports should be empty when there are no imports") + } + + func testStructureTestableImport() throws { + let source = "@testable import XCTest\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "XCTest") + // Offset should be 0 (start of @testable), not the import keyword + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // Length should cover "@testable import XCTest" = 23 + XCTAssertEqual(first?["key.length"] as? Int64, 23) + let attrs = first?["key.attributes"] as? [SourceKitRepresentable] + XCTAssertEqual(attrs?.count, 1) + let firstAttr = attrs?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(firstAttr?["key.attribute"] as? String, "source.decl.attribute.testable") + // "@testable" starts at 0, length 9 + XCTAssertEqual(firstAttr?["key.offset"] as? Int64, 0) + XCTAssertEqual(firstAttr?["key.length"] as? Int64, 9) + } + + func testStructureExportedTestableImport() throws { + let source = "@_exported @testable import XCTest\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "XCTest") + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "@_exported @testable import XCTest" = 34 + XCTAssertEqual(first?["key.length"] as? Int64, 34) + let attrs = first?["key.attributes"] as? [SourceKitRepresentable] + XCTAssertEqual(attrs?.count, 2) + let firstAttr = attrs?[0] as? [String: SourceKitRepresentable] + let secondAttr = attrs?[1] as? [String: SourceKitRepresentable] + XCTAssertEqual(firstAttr?["key.attribute"] as? String, "source.decl.attribute._exported") + // "@_exported" starts at 0, length 10 + XCTAssertEqual(firstAttr?["key.offset"] as? Int64, 0) + XCTAssertEqual(firstAttr?["key.length"] as? Int64, 10) + XCTAssertEqual(secondAttr?["key.attribute"] as? String, "source.decl.attribute.testable") + // "@testable" starts at 11, length 9 + XCTAssertEqual(secondAttr?["key.offset"] as? Int64, 11) + XCTAssertEqual(secondAttr?["key.length"] as? Int64, 9) + } + + func testStructurePlainImportHasNoAttributes() throws { + let structure = try Structure(file: File(contents: "import Foundation\nclass Foo {}"), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertNil(first?["key.attributes"], "plain imports should not have key.attributes") + } + + func testStructureSubmoduleImport() throws { + let source = "import Foundation.NSObject\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation.NSObject") + XCTAssertEqual(first?["key.module_name"] as? String, "Foundation") + XCTAssertNil(first?["key.import_kind"], "bare submodule import should not have import_kind") + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "import Foundation.NSObject" = 26 + XCTAssertEqual(first?["key.length"] as? Int64, 26) + } + + func testStructureSubmoduleImportNSDate() throws { + let source = "import Foundation.NSDate\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation.NSDate") + XCTAssertEqual(first?["key.module_name"] as? String, "Foundation") + XCTAssertNil(first?["key.import_kind"]) + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "import Foundation.NSDate" = 24 + XCTAssertEqual(first?["key.length"] as? Int64, 24) + } + + func testStructureKindQualifiedImportClass() throws { + let source = "import class Foundation.NSObject\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation.NSObject") + XCTAssertEqual(first?["key.module_name"] as? String, "Foundation") + XCTAssertEqual(first?["key.import_kind"] as? String, "class") + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "import class Foundation.NSObject" = 32 + XCTAssertEqual(first?["key.length"] as? Int64, 32) + } + + func testStructureKindQualifiedImportFunc() throws { + let source = "import func Darwin.exit\nfunc foo() {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Darwin.exit") + XCTAssertEqual(first?["key.module_name"] as? String, "Darwin") + XCTAssertEqual(first?["key.import_kind"] as? String, "func") + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "import func Darwin.exit" = 23 + XCTAssertEqual(first?["key.length"] as? Int64, 23) + } + + func testStructureKindQualifiedImportStruct() throws { + let source = "import struct Foundation.URL\nstruct Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation.URL") + XCTAssertEqual(first?["key.module_name"] as? String, "Foundation") + XCTAssertEqual(first?["key.import_kind"] as? String, "struct") + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "import struct Foundation.URL" = 28 + XCTAssertEqual(first?["key.length"] as? Int64, 28) + } + + func testStructureKindQualifiedImportEnum() throws { + let source = "import enum Foundation.JSONSerialization.ReadingOptions\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation.JSONSerialization.ReadingOptions") + XCTAssertEqual(first?["key.module_name"] as? String, "Foundation") + XCTAssertEqual(first?["key.import_kind"] as? String, "enum") + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "import enum Foundation.JSONSerialization.ReadingOptions" = 55 + XCTAssertEqual(first?["key.length"] as? Int64, 55) + } + + func testStructureMixedImportStyles() throws { + let source = "import Foundation\nimport func Darwin.exit\nimport struct Foundation.URL\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 3) + let names = imports?.compactMap { ($0 as? [String: SourceKitRepresentable])?["key.name"] as? String } + XCTAssertEqual(names, ["Foundation", "Darwin.exit", "Foundation.URL"]) + let kinds = imports?.map { ($0 as? [String: SourceKitRepresentable])?["key.import_kind"] as? String } + XCTAssertEqual(kinds?.count, 3) + XCTAssertNil(kinds?[0]) + XCTAssertEqual(kinds?[1], "func") + XCTAssertEqual(kinds?[2], "struct") + } + + func testStructureTestableSubmoduleImport() throws { + let source = "@testable import Foundation.NSObject\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 1) + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation.NSObject") + XCTAssertEqual(first?["key.module_name"] as? String, "Foundation") + XCTAssertNil(first?["key.import_kind"]) + XCTAssertEqual(first?["key.offset"] as? Int64, 0) + // "@testable import Foundation.NSObject" = 36 + XCTAssertEqual(first?["key.length"] as? Int64, 36) + let attrs = first?["key.attributes"] as? [SourceKitRepresentable] + XCTAssertEqual(attrs?.count, 1) + } + + func testStructurePlainImportHasNoModulenameOrKind() throws { + let structure = try Structure(file: File(contents: "import Foundation\nclass Foo {}"), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + let first = imports?.first as? [String: SourceKitRepresentable] + XCTAssertEqual(first?["key.name"] as? String, "Foundation") + XCTAssertNil(first?["key.module_name"], "plain module import should not have key.modulename") + XCTAssertNil(first?["key.import_kind"], "plain module import should not have key.import_kind") + } + + func testStructureConsecutiveImportsNoBlankLine() throws { + let source = "import Foundation\nimport UIKit\nimport CoreData\nclass Foo {}" + let structure = try Structure(file: File(contents: source), extractImports: true) + let imports = structure.dictionary["key.imports"] as? [SourceKitRepresentable] + XCTAssertNotNil(imports) + XCTAssertEqual(imports?.count, 3) + let names = imports?.compactMap { ($0 as? [String: SourceKitRepresentable])?["key.name"] as? String } + XCTAssertEqual(names, ["Foundation", "UIKit", "CoreData"]) + } + func testStructurePrintValidJSON() throws { let structure = try Structure(file: File(contents: "struct A { func b() {} }")) let expectedStructure: NSDictionary = [