Merge branch 'ldoktor-iter-tabular-console-chars2'

avocado/utils/astring.py

@@ -25,6 +25,7 @@ string. Even with the dot notation, people may try to do things like
 And not notice until their code starts failing.
 """
 
+import itertools
 import os.path
 import re
 
@@ -148,38 +149,48 @@ def iter_tabular_output(matrix, header=None):
     :param matrix: Matrix representation (list with n rows of m elements).
     :param header: Optional tuple or list with header elements to be displayed.
     """
-    if type(header) is list:
-        header = tuple(header)
-    lengths = []
+    def _get_matrix_with_header():
+        return itertools.chain([header], matrix)
+
+    def _get_matrix_no_header():
+        return matrix
+
+    if header is None:
+        header = []
     if header:
-        for column in header:
-            lengths.append(len(column))
+        get_matrix = _get_matrix_with_header
+    else:
+        get_matrix = _get_matrix_no_header
+
+    lengths = []
+    len_matrix = []
     str_matrix = []
-    for row in matrix:
-        str_matrix.append([])
-        for i, column in enumerate(row):
-            column = string_safe_encode(column)
-            str_matrix[-1].append(column)
-            col_len = len(column.decode("utf-8"))
+    for row in get_matrix():
+        len_matrix.append([])
+        str_matrix.append([string_safe_encode(column) for column in row])
+        for i, column in enumerate(str_matrix[-1]):
+            col_len = len(strip_console_codes(column.decode("utf-8")))
+            len_matrix[-1].append(col_len)
             try:
                 max_len = lengths[i]
                 if col_len > max_len:
                     lengths[i] = col_len
             except IndexError:
                 lengths.append(col_len)
-
-    if not lengths:     # No items...
-        raise StopIteration
-    format_string = " ".join(["%-" + str(leng) + "s"
-                              for leng in lengths[:-1]] +
-                             ["%s"])
-
-    if header:
-        out_line = format_string % header
-        yield out_line
-    for row in str_matrix:
-        out_line = format_string % tuple(row)
-        yield out_line
+        # For different no cols we need to calculate `lengths` of the last item
+        # but later in `yield` we don't want it in `len_matrix`
+        len_matrix[-1] = len_matrix[-1][:-1]
+
+    for row, row_lens in itertools.izip(str_matrix, len_matrix):
+        out = []
+        padding = [" " * (lengths[i] - row_lens[i])
+                   for i in xrange(len(row_lens))]
+        out = ["%s%s" % line for line in itertools.izip(row, padding)]
+        try:
+            out.append(row[-1])
+        except IndexError:
+            continue    # Skip empty rows
+        yield " ".join(out)
 
 
 def tabular_output(matrix, header=None):

selftests/unit/test_astring.py

@@ -17,6 +17,29 @@ class AstringTest(unittest.TestCase):
                           'foo               bar\n'
                           '/bin/bar/sbrubles /home/myuser/sbrubles'))
 
+    def testTabularWithConsoleCodes(self):
+        matrix = [("a", "bb", "ccc", "dddd", "last"),
+                  ("\x1b[94ma",             # {BLUE}a
+                   "\033[0mbb",             # {END}bb
+                   "cc\033[91mc",   # cc{RED}c
+                   # {RED}d{GREEN}d{BLUE}d{GREY}d{END}
+                   "\033[91md\033[92md\033[94md\033[90md\033[0m",
+                   "last")]
+        header = ['0', '1', '2', '3', '4']
+        self.assertEqual(astring.tabular_output(matrix, header),
+                         "0 1  2   3    4\n"
+                         "a bb ccc dddd last\n"
+                         "\x1b[94ma \x1b[0mbb cc\033[91mc "
+                         "\033[91md\033[92md\033[94md\033[90md\033[0m last")
+
+    def testTabularOutputDifferentNOCols(self):
+        matrix = [[], [1], [2, 2], [333, 333, 333], [4, 4, 4, 4444]]
+        self.assertEqual(astring.tabular_output(matrix),
+                         "1\n"
+                         "2   2\n"
+                         "333 333 333\n"
+                         "4   4   4   4444")
+
     def testUnicodeTabular(self):
         """
         Verifies tabular can handle utf-8 chars properly