Implement automatic terminal detection in DiagnosticsFormatter

Author: simba909

Sources/SwiftDiagnostics/DiagnosticsFormatter.swift

@@ -14,6 +14,21 @@ import SwiftSyntax
 
 public struct DiagnosticsFormatter {
 
+  public static let defaultContextSize: Int = 2
+
+  /// Determines how colors are used in printed diagnostics.
+  public enum ColorMode {
+    /// Colors are used if stderr is detected to be connected to a TTY; otherwise, colors will not
+    /// be used (for example, if stderr is redirected to a file). See: `TerminalHelper.isConnectedToTerminal`
+    case auto
+
+    /// Colors will not be used.
+    case off
+
+    /// Colors will always be used.
+    case on
+  }
+
   /// A wrapper struct for a source line and its diagnostics
   private struct AnnotatedSourceLine {
     var diagnostics: [Diagnostic]
@@ -26,21 +41,46 @@ public struct DiagnosticsFormatter {
   /// Whether to colorize formatted diagnostics.
   public let colorize: Bool
 
-  public init(contextSize: Int = 2, colorize: Bool = false) {
+  public init(contextSize: Int = defaultContextSize, colorize: Bool = false) {
     self.contextSize = contextSize
     self.colorize = colorize
   }
 
+  public init(contextSize: Int = defaultContextSize, colorMode: ColorMode) {
+    var colorize: Bool
+
+    switch colorMode {
+    case .auto:
+      colorize = TerminalHelper.isConnectedToTerminal
+    case .off:
+      colorize = false
+    case .on:
+      colorize = true
+    }
+
+    self.init(contextSize: contextSize, colorize: colorize)
+  }
+
   public static func annotatedSource<SyntaxType: SyntaxProtocol>(
     tree: SyntaxType,
     diags: [Diagnostic],
-    contextSize: Int = 2,
+    contextSize: Int = defaultContextSize,
     colorize: Bool = false
   ) -> String {
     let formatter = DiagnosticsFormatter(contextSize: contextSize, colorize: colorize)
     return formatter.annotatedSource(tree: tree, diags: diags)
   }
 
+  public static func annotatedSource<SyntaxType: SyntaxProtocol>(
+    tree: SyntaxType,
+    diags: [Diagnostic],
+    contextSize: Int = defaultContextSize,
+    colorMode: ColorMode
+  ) -> String {
+    let formatter = DiagnosticsFormatter(contextSize: contextSize, colorMode: colorMode)
+    return formatter.annotatedSource(tree: tree, diags: diags)
+  }
+
   /// Print given diagnostics for a given syntax tree on the command line
   public func annotatedSource<SyntaxType: SyntaxProtocol>(tree: SyntaxType, diags: [Diagnostic]) -> String {
     let slc = SourceLocationConverter(file: "", tree: tree)

Sources/SwiftDiagnostics/TerminalUtils.swift

@@ -0,0 +1,51 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#if canImport(Glibc)
+import Glibc
+#elseif os(Windows)
+import CRT
+#else
+import Darwin.C
+#endif
+
+#if os(Android)
+typealias FILEPointer = OpaquePointer
+#else
+typealias FILEPointer = UnsafeMutablePointer<FILE>
+#endif
+
+enum TerminalHelper {
+  static var isConnectedToTerminal: Bool {
+    return isTTY(stderr)
+  }
+
+  /// Checks if passed file pointer is a tty.
+  static func isTTY(_ filePointer: FILEPointer) -> Bool {
+    return terminalType(filePointer) == .tty
+  }
+
+  /// The type of terminal.
+  enum TerminalType {
+    /// The terminal is a TTY.
+    case tty
+
+    /// The terminal is a file stream.
+    case file
+  }
+
+  /// Computes the terminal type of the stream.
+  static func terminalType(_ filePointer: FILEPointer) -> TerminalType {
+    let isTTY = isatty(fileno(filePointer)) != 0
+    return isTTY ? .tty : .file
+  }
+}

Sources/swift-parser-cli/swift-parser-cli.swift

@@ -198,17 +198,22 @@ class PrintDiags: ParsableCommand {
   @Flag(name: .long, help: "Perform sequence folding with the standard operators")
   var foldSequences: Bool = false
 
-  @Flag(name: .long, help: "Colorize output with ANSI color codes")
+  @Flag(name: .long, help: "Force output coloring with ANSI color codes")
   var colorize: Bool = false
 
   func run() throws {
     let source = try getContentsOfSourceFile(at: sourceFile)
 
     source.withUnsafeBufferPointer { sourceBuffer in
       let tree = Parser.parse(source: sourceBuffer)
-
       var diags = ParseDiagnosticsGenerator.diagnostics(for: tree)
-      print(DiagnosticsFormatter.annotatedSource(tree: tree, diags: diags, colorize: colorize))
+      let annotatedSource = DiagnosticsFormatter.annotatedSource(
+        tree: tree,
+        diags: diags,
+        colorMode: colorize ? .on : .auto
+      )
+
+      print(annotatedSource)
 
       if foldSequences {
         diags += foldAllSequences(tree).1