Save ('s') and restore ('u') support added

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):

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):

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()
 

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

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

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

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()