Merge pull request #2428 from jhasse/elide-middle-ansi-escape

Correctly handle ANSI Escape Codes in ElideMiddle, fix #713

src/util.cc

@@ -26,6 +26,7 @@
 #include <assert.h>
 #include <errno.h>
 #include <fcntl.h>
+#include <regex>
 #include <stdarg.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -925,14 +926,43 @@ string ElideMiddle(const string& str, size_t width) {
       case 3: return "...";
   }
   const int kMargin = 3;  // Space for "...".
-  string result = str;
-  if (result.size() > width) {
-    size_t elide_size = (width - kMargin) / 2;
-    result =
-        result.substr(0, elide_size) + "..." +
-        result.substr(result.size() - elide_size - ((width - kMargin) % 2));
+  const static std::regex ansi_escape("\\x1b[^m]*m");
+  std::string result = std::regex_replace(str, ansi_escape, "");
+  if (result.size() <= width) {
+    return str;
+  }
+  int32_t elide_size = (width - kMargin) / 2;
+
+  std::vector<std::pair<int32_t, std::string>> escapes;
+  size_t added_len = 0;  // total number of characters
+
+  std::sregex_iterator it(str.begin(), str.end(), ansi_escape);
+  std::sregex_iterator end;
+  while (it != end) {
+    escapes.emplace_back(it->position() - added_len, it->str());
+    added_len += it->str().size();
+    ++it;
+  }
+
+  std::string new_status =
+      result.substr(0, elide_size) + "..." +
+      result.substr(result.size() - elide_size - ((width - kMargin) % 2));
+
+  added_len = 0;
+  // We need to put all ANSI escape codes back in:
+  for (const auto& escape : escapes) {
+    int32_t pos = escape.first;
+    if (pos > elide_size) {
+      pos -= result.size() - width;
+      if (pos < static_cast<int32_t>(width) - elide_size) {
+        pos = width - elide_size - (width % 2 == 0 ? 1 : 0);
+      }
+    }
+    pos += added_len;
+    new_status.insert(pos, escape.second);
+    added_len += escape.second.size();
   }
-  return result;
+  return new_status;
 }
 
 bool Truncate(const string& path, size_t size, string* err) {

src/util_test.cc

@@ -523,3 +523,41 @@ TEST(ElideMiddle, ElideInTheMiddle) {
   EXPECT_EQ("01234567...23456789", ElideMiddle(input, 19));
   EXPECT_EQ("01234567890123456789", ElideMiddle(input, 20));
 }
+
+TEST(ElideMiddle, ElideAnsiEscapeCodes) {
+  std::string input = "012345\x1B[0;35m67890123456789";
+  EXPECT_EQ("012...\x1B[0;35m6789", ElideMiddle(input, 10));
+  EXPECT_EQ("012345\x1B[0;35m67...23456789", ElideMiddle(input, 19));
+
+  EXPECT_EQ("Nothing \33[m string.", ElideMiddle("Nothing \33[m string.", 18));
+  EXPECT_EQ("0\33[m12...6789", ElideMiddle("0\33[m1234567890123456789", 10));
+
+  input = "abcd\x1b[1;31mefg\x1b[0mhlkmnopqrstuvwxyz";
+  EXPECT_EQ("", ElideMiddle(input, 0));
+  EXPECT_EQ(".", ElideMiddle(input, 1));
+  EXPECT_EQ("..", ElideMiddle(input, 2));
+  EXPECT_EQ("...", ElideMiddle(input, 3));
+  EXPECT_EQ("...\x1B[1;31m\x1B[0mz", ElideMiddle(input, 4));
+  EXPECT_EQ("a...\x1B[1;31m\x1B[0mz", ElideMiddle(input, 5));
+  EXPECT_EQ("a...\x1B[1;31m\x1B[0myz", ElideMiddle(input, 6));
+  EXPECT_EQ("ab...\x1B[1;31m\x1B[0myz", ElideMiddle(input, 7));
+  EXPECT_EQ("ab...\x1B[1;31m\x1B[0mxyz", ElideMiddle(input, 8));
+  EXPECT_EQ("abc...\x1B[1;31m\x1B[0mxyz", ElideMiddle(input, 9));
+  EXPECT_EQ("abc...\x1B[1;31m\x1B[0mwxyz", ElideMiddle(input, 10));
+  EXPECT_EQ("abcd\x1B[1;31m...\x1B[0mwxyz", ElideMiddle(input, 11));
+  EXPECT_EQ("abcd\x1B[1;31m...\x1B[0mvwxyz", ElideMiddle(input, 12));
+
+  EXPECT_EQ("abcd\x1B[1;31mef...\x1B[0muvwxyz", ElideMiddle(input, 15));
+  EXPECT_EQ("abcd\x1B[1;31mef...\x1B[0mtuvwxyz", ElideMiddle(input, 16));
+  EXPECT_EQ("abcd\x1B[1;31mefg\x1B[0m...tuvwxyz", ElideMiddle(input, 17));
+  EXPECT_EQ("abcd\x1B[1;31mefg\x1B[0m...stuvwxyz", ElideMiddle(input, 18));
+  EXPECT_EQ("abcd\x1B[1;31mefg\x1B[0mh...stuvwxyz", ElideMiddle(input, 19));
+
+  input = "abcdef\x1b[31mA\x1b[0mBC";
+  EXPECT_EQ("...\x1B[31m\x1B[0mC", ElideMiddle(input, 4));
+  EXPECT_EQ("a...\x1B[31m\x1B[0mC", ElideMiddle(input, 5));
+  EXPECT_EQ("a...\x1B[31m\x1B[0mBC", ElideMiddle(input, 6));
+  EXPECT_EQ("ab...\x1B[31m\x1B[0mBC", ElideMiddle(input, 7));
+  EXPECT_EQ("ab...\x1B[31mA\x1B[0mBC", ElideMiddle(input, 8));
+  EXPECT_EQ("abcdef\x1b[31mA\x1b[0mBC", ElideMiddle(input, 9));
+}