diff --git a/colorama/ansi.py b/colorama/ansi.py index 7877658..4c6895d 100644 --- a/colorama/ansi.py +++ b/colorama/ansi.py @@ -44,6 +44,10 @@ def BACK(self, n=1): return CSI + str(n) + 'D' def POS(self, x=1, y=1): return CSI + str(y) + ';' + str(x) + 'H' + def SAVE(self): + return CSI + 's' + def RESTORE(self): + return CSI + 'u' class AnsiFore(AnsiCodes): diff --git a/colorama/ansitowin32.py b/colorama/ansitowin32.py index 1d6e605..b556b70 100644 --- a/colorama/ansitowin32.py +++ b/colorama/ansitowin32.py @@ -82,6 +82,9 @@ def __init__(self, wrapped, convert=None, strip=None, autoreset=False): # are we wrapping stderr? self.on_stderr = self.wrapped is sys.stderr + # saved cursor positions + self.saved_positions = [] + def should_wrap(self): ''' True if this class is actually needed. If false, then the output @@ -219,6 +222,12 @@ def call_win32(self, command, params): # A - up, B - down, C - forward, D - back x, y = {'A': (0, -n), 'B': (0, n), 'C': (n, 0), 'D': (-n, 0)}[command] winterm.cursor_adjust(x, y, on_stderr=self.on_stderr) + elif command in 's': # save cursor position + self.saved_positions.append(winterm.get_cursor_position(on_stderr=self.on_stderr)) + elif command in 'u': # restore cursor position + if self.saved_positions: + position = self.saved_positions.pop() + winterm.set_cursor_position(position, on_stderr=self.on_stderr) def convert_osc(self, text): diff --git a/colorama/tests/ansi_test.py b/colorama/tests/ansi_test.py index f4adf4e..25ec04d 100644 --- a/colorama/tests/ansi_test.py +++ b/colorama/tests/ansi_test.py @@ -5,7 +5,7 @@ except ImportError: from unittest import TestCase, main -from ..ansi import Fore, Back, Style +from ..ansi import Fore, Back, Style, Cursor from ..ansitowin32 import AnsiToWin32 @@ -76,6 +76,11 @@ def testStyleAttributes(self): self.assertEqual(Style.BRIGHT, '\033[1m') + def testCursorMethods(self): + self.assertEqual(Cursor.SAVE(), '\033[s') + self.assertEqual(Cursor.RESTORE(), '\033[u') + + if __name__ == '__main__': main() diff --git a/colorama/winterm.py b/colorama/winterm.py index 60309d3..c8f6e65 100644 --- a/colorama/winterm.py +++ b/colorama/winterm.py @@ -76,11 +76,13 @@ def style(self, style=None, on_stderr=False): def set_console(self, attrs=None, on_stderr=False): if attrs is None: attrs = self.get_attrs() - handle = win32.STDOUT - if on_stderr: - handle = win32.STDERR + handle = self.get_handle(on_stderr) win32.SetConsoleTextAttribute(handle, attrs) + @staticmethod + def get_handle(on_stderr=False): + return win32.STDERR if on_stderr else win32.STDOUT + def get_position(self, handle): position = win32.GetConsoleScreenBufferInfo(handle).dwCursorPosition # Because Windows coordinates are 0-based, @@ -89,31 +91,38 @@ def get_position(self, handle): position.Y += 1 return position + def get_cursor_position(self, on_stderr=False, adjust=True): + handle = self.get_handle(on_stderr) + info = win32.GetConsoleScreenBufferInfo(handle) + position = info.dwCursorPosition + # Because Windows coordinates are 0-based, + # and win32.SetConsoleCursorPosition expects 1-based. + y, x = position.Y + 1, position.X + 1 + if adjust: + window = info.srWindow + y -= window.Top + x -= window.Left + return y, x + def set_cursor_position(self, position=None, on_stderr=False): if position is None: # I'm not currently tracking the position, so there is no default. # position = self.get_position() return - handle = win32.STDOUT - if on_stderr: - handle = win32.STDERR + handle = self.get_handle(on_stderr) win32.SetConsoleCursorPosition(handle, position) def cursor_adjust(self, x, y, on_stderr=False): - handle = win32.STDOUT - if on_stderr: - handle = win32.STDERR - position = self.get_position(handle) - adjusted_position = (position.Y + y, position.X + x) + (cy, cx) = self.get_cursor_position(on_stderr, adjust=False) + adjusted_position = (cy + y, cx + x) + handle = self.get_handle(on_stderr) win32.SetConsoleCursorPosition(handle, adjusted_position, adjust=False) def erase_screen(self, mode=0, on_stderr=False): # 0 should clear from the cursor to the end of the screen. # 1 should clear from the cursor to the beginning of the screen. # 2 should clear the entire screen, and move cursor to (1,1) - handle = win32.STDOUT - if on_stderr: - handle = win32.STDERR + handle = self.get_handle(on_stderr) csbi = win32.GetConsoleScreenBufferInfo(handle) # get the number of character cells in the current buffer cells_in_screen = csbi.dwSize.X * csbi.dwSize.Y @@ -140,9 +149,7 @@ def erase_line(self, mode=0, on_stderr=False): # 0 should clear from the cursor to the end of the line. # 1 should clear from the cursor to the beginning of the line. # 2 should clear the entire line. - handle = win32.STDOUT - if on_stderr: - handle = win32.STDERR + handle = self.get_handle(on_stderr) csbi = win32.GetConsoleScreenBufferInfo(handle) if mode == 0: from_coord = csbi.dwCursorPosition diff --git a/demos/demo.bat b/demos/demo.bat index 8386a7f..56183b5 100644 --- a/demos/demo.bat +++ b/demos/demo.bat @@ -30,3 +30,6 @@ python demo06.py :: demo07.py not shown :: Demonstrate cursor relative movement: UP, DOWN, FORWARD, and BACK in colorama.CURSOR + +:: Demonstrate cursor saving, loading and positioning: SAVE, LOAD and POS in colorama.Cursor +python demo09.py diff --git a/demos/demo.sh b/demos/demo.sh index 07f556e..b449130 100644 --- a/demos/demo.sh +++ b/demos/demo.sh @@ -34,3 +34,6 @@ python demo06.py #Demonstrate the use of a context manager instead of manually using init and deinit python demo08.py + +# Demonstrate cursor saving, loading and positioning: SAVE, LOAD and POS in colorama.Cursor +python demo09.py diff --git a/demos/demo09.py b/demos/demo09.py new file mode 100644 index 0000000..7d9a9f4 --- /dev/null +++ b/demos/demo09.py @@ -0,0 +1,34 @@ +from __future__ import print_function +import fixpath +import colorama +import sys +import time + +# Demonstrate cursor saving, restoring and positioning: SAVE, RESTORE and POS in colorama.Cursor + +save = colorama.Cursor.SAVE +restore = colorama.Cursor.RESTORE +pos = colorama.Cursor.POS + +blue = colorama.Back.BLUE +reset = colorama.Back.RESET + +def main(): + """ + expected output: + Current state is shown at top + Progress is shown at the current cursor position + """ + colorama.init() + for i in range(1, 10): + sys.stdout.write("Step {}: ".format(i)) + for j in range(1, 10): + sys.stdout.write(str(j)) + sys.stdout.write(save() + pos(10, 1) + blue + " State {}.{} ".format(i, j) + restore() + reset) + sys.stdout.flush() + time.sleep(0.02) + print() + + +if __name__ == '__main__': + main()