Merge pull request #6 from scottrhoyt/console_escape_stripping

Console Escape Stripping
scottrhoyt · Mar 11, 2016 · eda101d · eda101d
2 parents 0865906 + 979f0ea
commit eda101d
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 3 deletions.
diff --git a/Source/SwiftyTextTable/TextTable.swift b/Source/SwiftyTextTable/TextTable.swift
@@ -8,6 +8,13 @@
 
 import Foundation
 
+// MARK: Console Escape Stripping
+private let strippingPattern = "(?:\u{001B}\\[(?:[0-9]|;)+m)*(.*?)(?:\u{001B}\\[0m)+"
+
+// We can safely force try this regex because the pattern has be tested to work.
+// swiftlint:disable:next force_try
+private let strippingRegex = try! NSRegularExpression(pattern: strippingPattern, options: [])
+
 extension String: CustomStringConvertible {
     public var description: String {
         return self
@@ -23,6 +30,19 @@ private extension String {
         }
         return self
     }
+
+    func stripped() -> String {
+        let matches = strippingRegex
+            .matchesInString(self, options: [], range: NSRange(location: 0, length: self.characters.count))
+            .map {
+                (self as NSString).substringWithRange($0.rangeAtIndex(1))
+        }
+        return matches.isEmpty ? self : matches.joinWithSeparator("")
+    }
+
+    func strippedLength() -> Int {
+        return stripped().characters.count
+    }
 }
 
 private func fence(strings: [String], separator: String) -> String {
@@ -37,8 +57,8 @@ public struct TextTableColumn {
         self.header = header
     }
 
-    var width: Int {
-        return max(header.characters.count, values.reduce(0) { max($0, $1.characters.count) })
+    public var width: Int {
+        return max(header.strippedLength(), values.reduce(0) { max($0, $1.strippedLength()) })
     }
 }
 

diff --git a/Source/SwiftyTextTableTests/SwiftyTextTableTests.swift b/Source/SwiftyTextTableTests/SwiftyTextTableTests.swift
@@ -50,9 +50,33 @@ class SwiftyTextTableTests: XCTestCase {
         XCTAssertEqual(output, expected)
     }
 
+    func testStripping() {
+        let c1 = TextTableColumn(header: "\u{001B}[0mHello\u{001B}[0m")
+        XCTAssertEqual(c1.width, 5)
+
+        let c2 = TextTableColumn(header: "\u{001B}[31m\u{001B}[4;31;93mHello World\u{001B}[0m\u{001B}[0m")
+        XCTAssertEqual(c2.width, 11)
+
+        let c3 = TextTableColumn(header: "\u{001B}[0m\u{001B}[0m")
+        XCTAssertEqual(c3.width, 0)
+
+        let c4 = TextTableColumn(header: "\u{001B}[31mHello World\u{001B}[0m")
+        XCTAssertEqual(c4.width, 11)
+
+        let c5 = TextTableColumn(header: "\u{001B}[4;31;42;93;5mHello World\u{001B}[0m")
+        XCTAssertEqual(c5.width, 11)
+
+        let c6 = TextTableColumn(header: "\u{001B}[4;31;93mHello World\u{001B}[0m")
+        XCTAssertEqual(c6.width, 11)
+
+        let c7 = TextTableColumn(header: "Hello World")
+        XCTAssertEqual(c7.width, 11)
+    }
+
     // MARK: - protocol XCTestCaseProvider for SPM
     lazy var allTests: [(String, () throws -> Void)] = [
         ("testRenderDefault", self.testRenderDefault),
-        ("testRenderCustom", self.testRenderCustom)
+        ("testRenderCustom", self.testRenderCustom),
+        ("testStripping", self.testRenderCustom)
     ]
 }