Save ('s') and restore ('u') support added (2nd try)

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

@@ -118,6 +118,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
@@ -255,6 +258,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

@@ -2,7 +2,7 @@
 import sys
 from unittest import TestCase, main
 
-from ..ansi import Back, Fore, Style
+from ..ansi import Back, Cursor, Fore, Style
 from ..ansitowin32 import AnsiToWin32
 
 stdout_orig = sys.stdout
@@ -72,5 +72,10 @@ 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

@@ -83,11 +83,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,
@@ -96,31 +98,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
@@ -150,9 +159,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-cmdline-args.py

@@ -0,0 +1,21 @@
+# https://www.youtube.com/watch?v=F5a8RLY2N8M&list=PL1_riyn9sOjcKIAYzo7f8drxD-Yg9La-D&index=61
+# Generic colorama demo using command line arguments
+# By George Ogden
+from colorama import Fore, Back, Style, init
+import argparse
+parser = argparse.ArgumentParser("colorama demo")
+
+def format(module):
+    return list(map(lambda x: x.lower(),module.__dict__.keys()))
+
+def find(module,item):
+    return module.__dict__[item.upper()]
+
+parser.add_argument("-c","--colour",choices=format(Fore),default="RESET")
+parser.add_argument("-b","--background",choices=format(Back),default="RESET")
+parser.add_argument("-s","--style",choices=format(Style),default="RESET_ALL")
+parser.add_argument("-t","--text",default="Lorem ipsum dolor sit amet")
+
+args = parser.parse_args()
+
+print(find(Style,args.style) + find(Fore,args.colour) + find(Back,args.background) + args.text + Style.RESET_ALL)

demos/demo.bat

@@ -34,3 +34,7 @@ python demo07.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/demo.sh

@@ -36,3 +36,6 @@ python demo07.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

@@ -1,21 +1,35 @@
-# https://www.youtube.com/watch?v=F5a8RLY2N8M&list=PL1_riyn9sOjcKIAYzo7f8drxD-Yg9La-D&index=61
-# Generic colorama demo using command line arguments
-# By George Ogden
-from colorama import Fore, Back, Style, init
-import argparse
-parser = argparse.ArgumentParser("colorama demo")
-
-def format(module):
-    return list(map(lambda x: x.lower(),module.__dict__.keys()))
-
-def find(module,item):
-    return module.__dict__[item.upper()]
-
-parser.add_argument("-c","--colour",choices=format(Fore),default="RESET")
-parser.add_argument("-b","--background",choices=format(Back),default="RESET")
-parser.add_argument("-s","--style",choices=format(Style),default="RESET_ALL")
-parser.add_argument("-t","--text",default="Lorem ipsum dolor sit amet")
-
-args = parser.parse_args()
-
-print(find(Style,args.style) + find(Fore,args.colour) + find(Back,args.background) + args.text + Style.RESET_ALL)
+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()
+