From 92a96ee5a6a5a713a2c8d7d99f33742061c9056c Mon Sep 17 00:00:00 2001
From: John Szumski <jszumski@squareup.com>
Date: Tue, 11 Feb 2025 12:16:38 -0500
Subject: [PATCH] Add the ability to add comment commands on the same line as
 an import statement

---
 .../SyntaxAnalysis/ImportSyntaxVisitor.swift  |  2 +-
 .../ImportFixture.swift                       |  4 ++
 .../Syntax/ImportVisitTest.swift              | 43 +++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 Tests/Fixtures/Sources/DeclarationVisitorFixtures/ImportFixture.swift
 create mode 100644 Tests/PeripheryTests/Syntax/ImportVisitTest.swift

diff --git a/Sources/SyntaxAnalysis/ImportSyntaxVisitor.swift b/Sources/SyntaxAnalysis/ImportSyntaxVisitor.swift
index e8b4cc05c..25985eb97 100644
--- a/Sources/SyntaxAnalysis/ImportSyntaxVisitor.swift
+++ b/Sources/SyntaxAnalysis/ImportSyntaxVisitor.swift
@@ -27,7 +27,7 @@ public final class ImportSyntaxVisitor: PeripherySyntaxVisitor {
             isTestable: attributes.contains("testable"),
             isExported: attributes.contains("_exported") || node.modifiers.contains { $0.name.text == "public" },
             location: location,
-            commentCommands: CommentCommand.parseCommands(in: node.leadingTrivia)
+            commentCommands: CommentCommand.parseCommands(in: node.leadingTrivia.merging(node.trailingTrivia))
         )
         importStatements.append(statement)
     }
diff --git a/Tests/Fixtures/Sources/DeclarationVisitorFixtures/ImportFixture.swift b/Tests/Fixtures/Sources/DeclarationVisitorFixtures/ImportFixture.swift
new file mode 100644
index 000000000..212e36228
--- /dev/null
+++ b/Tests/Fixtures/Sources/DeclarationVisitorFixtures/ImportFixture.swift
@@ -0,0 +1,4 @@
+import Foundation
+// periphery:ignore
+import CoreGraphics
+import Swift // periphery:ignore
diff --git a/Tests/PeripheryTests/Syntax/ImportVisitTest.swift b/Tests/PeripheryTests/Syntax/ImportVisitTest.swift
new file mode 100644
index 000000000..64a3acf8b
--- /dev/null
+++ b/Tests/PeripheryTests/Syntax/ImportVisitTest.swift
@@ -0,0 +1,43 @@
+import Foundation
+@testable import SourceGraph
+@testable import SyntaxAnalysis
+@testable import TestShared
+import XCTest
+
+final class ImportVisitTest: XCTestCase {
+    private var results: [ImportStatement]!
+
+    override func setUpWithError() throws {
+        try super.setUpWithError()
+        let multiplexingVisitor = try MultiplexingSyntaxVisitor(file: fixturePath)
+        let visitor = multiplexingVisitor.add(ImportSyntaxVisitor.self)
+        multiplexingVisitor.visit()
+        results = visitor.importStatements
+    }
+
+    override func tearDown() {
+        results = nil
+        super.tearDown()
+    }
+
+    func testCommentCommands() {
+        let expectedIgnored = ["CoreGraphics", "Swift"]
+        let actualIgnored = results.filter { $0.commentCommands.contains(.ignore) }.map(\.module)
+        XCTAssertEqual(actualIgnored, expectedIgnored, "Ignored modules did not match the expected set")
+
+        let actualUnignored = results.filter { !$0.commentCommands.contains(.ignore) }.map(\.module)
+        let expectedUnignored = ["Foundation"]
+        XCTAssertEqual(actualUnignored, expectedUnignored, "Unignored modules did not match the expected set")
+    }
+
+    // MARK: - Private
+
+    private var fixturePath: SourceFile {
+        let path = FixturesProjectPath.appending("Sources/DeclarationVisitorFixtures/ImportFixture.swift")
+        return SourceFile(path: path, modules: ["DeclarationVisitorFixtures"])
+    }
+
+    private func fixtureLocation(line: Int, column: Int = 9) -> Location {
+        Location(file: fixturePath, line: line, column: column)
+    }
+}