diff --git a/px/px_top.py b/px/px_top.py index 955edae..c111a10 100644 --- a/px/px_top.py +++ b/px/px_top.py @@ -126,6 +126,14 @@ def get_toplist(baseline): return toplist +def writebytes(bytestring): + if sys.version_info.major == 2: + sys.stdout.write(bytestring) + else: + # http://stackoverflow.com/a/908440/473672 + sys.stdout.buffer.write(bytestring) + + def clear_screen(): """ Clear the screen and move cursor to top left corner: @@ -139,28 +147,17 @@ def clear_screen(): """ CSI = b"\x1b[" - if sys.version_info.major == 2: - sys.stdout.write(CSI + b"1J") - sys.stdout.write(CSI + b"H") - else: - # http://stackoverflow.com/a/908440/473672 - sys.stdout.buffer.write(CSI + b"1J") - sys.stdout.buffer.write(CSI + b"H") - sys.stdout.flush() + writebytes(CSI + b"1J") + writebytes(CSI + b"H") -def redraw(load_bar, baseline, rows, columns): - """ - Refresh display relative to the given baseline. - - The new display will be (at most) rows rows x columns columns. - """ +def get_screen_lines(load_bar, baseline, rows, columns): load = px_load.get_load_values() - loadstring = px_load.get_load_string(load) + loadstring = px_load.get_load_string(load).encode('utf-8') loadbar = load_bar.get_bar(load=load[0], columns=40) lines = [ - "System load: " + loadstring + " |{}|".format(loadbar), - ""] + b"System load: " + loadstring + b" |" + loadbar + b"|", + b""] toplist_table_lines = px_terminal.to_screen_lines(get_toplist(baseline), columns) if toplist_table_lines: @@ -168,12 +165,25 @@ def redraw(load_bar, baseline, rows, columns): heading_line = px_terminal.get_string_of_length(heading_line, columns) heading_line = px_terminal.inverse_video(heading_line) toplist_table_lines[0] = heading_line + + toplist_table_lines = map(lambda s: s.encode('utf-8'), toplist_table_lines) lines += toplist_table_lines - clear_screen() + return lines[0:rows] + + +def redraw(load_bar, baseline, rows, columns, clear=True): + """ + Refresh display relative to the given baseline. + + The new display will be (at most) rows rows x columns columns. + """ + lines = get_screen_lines(load_bar, baseline, rows, columns) + if clear: + clear_screen() # We need both \r and \n when the TTY is in tty.setraw() mode - sys.stdout.write("\r\n".join(lines[0:rows])) + writebytes(b"\r\n".join(lines)) sys.stdout.flush() diff --git a/tests/px_top_test.py b/tests/px_top_test.py index 992bca1..47a54b1 100644 --- a/tests/px_top_test.py +++ b/tests/px_top_test.py @@ -74,4 +74,29 @@ def test_redraw(): # Just make sure it doesn't crash loadbar = px_load_bar.PxLoadBar(1, 1) baseline = px_process.get_all() - px_top.redraw(loadbar, baseline, 100, 10) + px_top.redraw(loadbar, baseline, 100, 10, clear=False) + + +def test_get_screen_lines(): + loadbar = px_load_bar.PxLoadBar(1, 1) + baseline = px_process.get_all() + + SCREEN_ROWS = 10 + SCREEN_COLUMNS = 100 + lines = px_top.get_screen_lines( + loadbar, baseline, SCREEN_ROWS, SCREEN_COLUMNS) + + # Top row should contain ANSI escape codes + CSI = b"\x1b[" + assert b'CSI' in lines[0].replace(CSI, b'CSI') + + assert len(lines) == SCREEN_ROWS + + # This is a bit fluffy, but due to ANSI escape codes the line lengths are + # difficult to measure. So if at least one line is exactly the right length + # then we're probably fine. + found_right_length = False + for line in lines: + if len(line) == SCREEN_COLUMNS: + found_right_length = True + assert found_right_length