Added support for @main attribute based scripts (#194)

* Added support for @main attribute based scripts, included async-main-count-lines Example

* removed unnecessary manual #warning

* Improved Package.swift manifest file creation

* Minor improvements

* Update Sources/Command/edit().swift

Applying changes suggested by @Amzd for Swift <= 5.5

Co-authored-by: Casper Zandbergen <[email protected]>

* Fix compilation in older Swifts

---------

Co-authored-by: Casper Zandbergen <[email protected]>
Co-authored-by: Max Howell <[email protected]>

Author: 3 people

Examples/async-main-count-lines

@@ -0,0 +1,16 @@
+#!/user/bin/swift sh
+import Foundation
+import ArgumentParser // apple/swift-argument-parser ~> 1.4.0
+
+/// example borrowed from [here](https://swiftpackageindex.com/apple/swift-argument-parser/1.4.0/documentation/argumentparser/asyncparsablecommand)
+@main
+struct CountLines: AsyncParsableCommand {
+    @Argument(transform: URL.init(fileURLWithPath:))
+    var inputFile: URL
+    mutating func run() async throws {
+        let fileHandle = try FileHandle(forReadingFrom: inputFile)
+        let lineCount = try await fileHandle.bytes.lines.reduce(into: 0)
+            { count, _ in count += 1 }
+        print(lineCount)
+    }
+}

Sources/Command/edit().swift

@@ -7,8 +7,15 @@ import Path
 public func edit(path: Path) throws -> Never {
 #if os(macOS)
     let input:Script.Input = .path(path)
-    let deps = try StreamReader(path: path).compactMap { try ImportSpecification(line: $0, from: input) }
-    let script = Script(for: .path(path), dependencies: deps)
+    let reader = try StreamReader(path: path)
+    var style: ExecutableTargetMainStyle = .topLevelCode
+    let deps: [ImportSpecification] = try reader.compactMap { line in
+        if line.contains("@main") && !(line.contains("//") || line.contains("/*")) {
+            style = .mainAttribute
+        }
+        return try ImportSpecification(line: line, from: input)
+    }
+    let script = Script(for: .path(path), style: style, dependencies: deps)
     try script.write()
     try exec(arg0: "/usr/bin/swift", args: ["sh-edit", path.string, script.buildDirectory.string])
 #else

Sources/Command/editor().swift

@@ -6,8 +6,15 @@ import Path
 
 public func editor(path: Path) throws -> Never {
     let input:Script.Input = .path(path)
-    let deps = try StreamReader(path: path).compactMap { try ImportSpecification(line: $0, from: input) }
-    let script = Script(for: .path(path), dependencies: deps)
+    let reader = try StreamReader(path: path)
+    var style: ExecutableTargetMainStyle = .topLevelCode
+    let deps: [ImportSpecification] = try reader.compactMap { line in
+        if line.contains("@main") && !(line.contains("//") || line.contains("/*")) {
+            style = .mainAttribute
+        }
+        return try ImportSpecification(line: line, from: input)
+    }
+    let script = Script(for: .path(path), style: style, dependencies: deps)
     try script.write()
 
     guard let editor = ProcessInfo.processInfo.environment["EDITOR"] else {

Sources/Command/run().swift

@@ -27,6 +27,7 @@ private func run<T>(reader: StreamReader, input: Input, arguments: T) throws ->
 
     var deps = [ImportSpecification]()
     var lines = [String]()
+    var mainStyle: ExecutableTargetMainStyle = .topLevelCode
 
     var transformedInput: Script.Input {
         switch input {
@@ -36,12 +37,14 @@ private func run<T>(reader: StreamReader, input: Input, arguments: T) throws ->
             return .path(path)
         }
     }
-
     for (index, line) in reader.enumerated() {
         if index == 0, line.hasPrefix("#!") {
             lines.append("// shebang removed")  // keep line numbers in sync
             continue
         }
+        if line.contains("@main") && !(line.contains("//") || line.contains("/*")) {
+            mainStyle = .mainAttribute
+        }
         if let result = try ImportSpecification(line: line, from: transformedInput) {
             deps.append(result)
         }
@@ -53,7 +56,7 @@ private func run<T>(reader: StreamReader, input: Input, arguments: T) throws ->
         }
     }
 
-    let script = Script(for: transformedInput, dependencies: deps, arguments: Array(arguments))
+    let script = Script(for: transformedInput, style: mainStyle, dependencies: deps, arguments: Array(arguments))
     try script.run()
 }
 

Sources/Script/ExecutableTargetMainStyle.swift

@@ -0,0 +1,8 @@
+
+/// wheter the script has @main
+public enum ExecutableTargetMainStyle {
+    /// script has @main, script source file cannot be named main.swift, as this would be compilation error
+    case mainAttribute
+    /// script doesn't have @main, and it's ok to have main.swift file as script source
+    case topLevelCode
+}

Sources/Script/Script.swift

@@ -8,6 +8,7 @@ public class Script {
     let input: Input
     let deps: [ImportSpecification]
     let args: [String]
+    let mainStyle: ExecutableTargetMainStyle
 
     private let inputPathHash: String?
 
@@ -30,18 +31,22 @@ public class Script {
     }
 
     public var mainSwift: Path {
-        return buildDirectory/"main.swift"
+        switch mainStyle {
+            case .mainAttribute: return buildDirectory/"Root.swift"
+            case .topLevelCode: return buildDirectory/"main.swift"
+        }
     }
 
     public enum Input {
         case path(Path)
         case string(name: String, content: String)
     }
 
-    public init(for: Input, dependencies: [ImportSpecification], arguments: [String] = []) {
+    public init(for: Input, style: ExecutableTargetMainStyle, dependencies: [ImportSpecification], arguments: [String] = []) {
         input = `for`
         deps = dependencies
         args = arguments
+        mainStyle = style
 
         // cache hash if appropriate since accessed often and involves work
         if case let Input.path(path) = input {
@@ -68,20 +73,30 @@ public class Script {
     public func write() throws {
         //NOTE we only support Swift >= 4.2 basically
         //TODO dependency module names might not correspond the products that packages export, must parse `swift package dump-package` output
-
         if depsCache != deps {
             // this check because SwiftPM has to reparse the manifest if we rewrite it
             // this is noticably slow, so avoid it if possible
-
             var macOS: String {
                 let version = ProcessInfo.processInfo.operatingSystemVersion
                 return ".macOS(\"\(version.majorVersion).\(version.minorVersion)\")"
             }
-
             try buildDirectory.mkdir(.p)
+            let targetDefinition: String
+            let swiftToolsVersion: String
+            let sourceFile: String
+            switch mainStyle {
+                case .mainAttribute:
+                    targetDefinition = "executableTarget"
+                    swiftToolsVersion = "5.5"
+                    sourceFile = "Root.swift"
+                case .topLevelCode:
+                    targetDefinition = "target"
+                    swiftToolsVersion = "5.1"
+                    sourceFile = "main.swift"
+            }
             // we are using tools version 5.1 while we still can as >= 5.3 makes specifying deps significantly more complex
             try """
-                // swift-tools-version:5.1
+                // swift-tools-version:\(swiftToolsVersion)
                 import PackageDescription
 
                 let pkg = Package(name: "\(name)")
@@ -93,12 +108,12 @@ public class Script {
                     \(deps.packageLines)
                 ]
                 pkg.targets = [
-                    .target(
+                    .\(targetDefinition)(
                         name: "\(name)",
                         dependencies: [\(deps.mainTargetDependencies)],
                         path: ".",
                         exclude: ["deps.json"],
-                        sources: ["main.swift"]
+                        sources: ["\(sourceFile)"]
                     )
                 ]
 
@@ -109,19 +124,26 @@ public class Script {
                 #endif
 
                 """.write(to: manifestPath)
-
             try JSONEncoder().encode(deps).write(to: depsCachePath)
         }
-
         switch input {
         case .path(let userPath):
             func mklink() throws { try userPath.symlink(as: mainSwift) }
-
-            if let linkdst = try? mainSwift.readlink(), linkdst != userPath {
-                try mainSwift.delete()
-                try mklink()
-            } else if !mainSwift.exists {
-                try mklink()
+            switch mainStyle {
+                case .mainAttribute:
+                    try mainSwift.delete()
+                    let reader = try StreamReader(path: userPath)
+                    let source = reader.compactMap { line in
+                        line.contains("#!") ? .none : line
+                    }.joined(separator: "\n")
+                    try source.write(to: mainSwift)
+                case .topLevelCode:
+                    if let linkdst = try? mainSwift.readlink(), linkdst != userPath {
+                        try mainSwift.delete()
+                        try mklink()
+                    } else if !mainSwift.exists {
+                        try mklink()
+                    }
             }
         case .string(_, let contents):
             if let currentContents = try? String(contentsOf: mainSwift), currentContents == contents { break }
@@ -178,12 +200,12 @@ public class Script {
     public func run() throws -> Never {
         if scriptChanged {
             try write()
-
             // first arg has to be same as executable path
             let task = Process()
             task.launchPath = Path.swift.string
             task.arguments = ["build", "-Xswiftc", "-suppress-warnings"]
             task.currentDirectoryPath = buildDirectory.string
+
           #if !os(Linux)
             task.standardOutput = task.standardError
           #else