Highlight short ptop lines all the way to the end

Doing multiple crops on the same ANSI decorated line puts reset
characters in it, which made us inverting it later fail.

px/px_terminal.py

@@ -294,7 +294,7 @@ def to_screen_lines(procs,  # type: List[px_process.PxProcess]
             max_cpu_time = proc.cpu_time_seconds
             max_cpu_time_s = proc.cpu_time_s
 
-    for proc in procs:
+    for line_number, proc in enumerate(procs):
         cpu_percent_s = proc.cpu_percent_s
         if proc.cpu_percent_s == "0%":
             cpu_percent_s = faint(cpu_percent_s.rjust(cpu_width))
@@ -317,19 +317,13 @@ def to_screen_lines(procs,  # type: List[px_process.PxProcess]
             memory_percent_s, proc.cmdline)
 
         cropped = line
-        if columns is not None:
+        if row_to_highlight == line_number:
+            cropped = get_string_of_length(cropped, columns)
+            cropped = inverse_video(cropped)
+        elif columns is not None:
             cropped = crop_ansi_string_at_length(line, columns)
         lines.append(cropped)
 
-    if row_to_highlight is not None:
-        # The "+ 1" here is to skip the heading line
-        highlighted = lines[row_to_highlight + 1]
-        highlighted = get_string_of_length(highlighted, columns)
-        highlighted = inverse_video(highlighted)
-
-        # The "+ 1" here is to skip the heading line
-        lines[row_to_highlight + 1] = highlighted
-
     lines[0] = heading_line
 
     return lines
@@ -396,14 +390,14 @@ def get_string_of_length(string, length):
     if length is None:
         return string
 
-    incoming_length = visual_length(string)
-    if incoming_length == length:
+    initial_length = visual_length(string)
+    if initial_length == length:
         return string
 
-    if incoming_length < length:
-        return string + u' ' * (length - incoming_length)
+    if initial_length < length:
+        return string + u' ' * (length - initial_length)
 
-    if incoming_length > length:
+    if initial_length > length:
         return crop_ansi_string_at_length(string, length)
 
     assert False  # How did we end up here?

tests/px_terminal_test.py

@@ -65,11 +65,19 @@ def test_to_screen_lines_unicode():
 
 
 def test_get_string_of_length():
+    CSI = u"\x1b["
+
     assert px_terminal.get_string_of_length("12345", 3) == "123"
     assert px_terminal.get_string_of_length("12345", 5) == "12345"
     assert px_terminal.get_string_of_length("12345", 7) == "12345  "
     assert px_terminal.get_string_of_length("12345", None) == "12345"
 
+    # Test with escaped strings
+    mid_bold = "1" + px_terminal.bold("234") + "5"
+    assert px_terminal.get_string_of_length(mid_bold, 6) == mid_bold + " "
+    assert px_terminal.get_string_of_length(mid_bold, 3) == \
+        "1" + CSI + "1m23" + CSI + "0m"
+
 
 def test_crop_ansi_string_at_length():
     CSI = u"\x1b["