Merge pull request #89 from dcoshea/issue-85-try2

Issue #85: Don't leave unwanted numbers in ANSI FSM stack/memory.

2 files changed, 43 insertions, 11 deletions

diff --git a/pexpect/ANSI.py b/pexpect/ANSI.py
index 83ea6a8..3460114 100644
--- a/pexpect/ANSI.py
+++ b/pexpect/ANSI.py
@@ -227,7 +227,7 @@ class ANSI (term):
         self.state.add_transition ('J', 'ELB', DoEraseDown, 'INIT')
         self.state.add_transition ('K', 'ELB', DoEraseEndOfLine, 'INIT')
         self.state.add_transition ('r', 'ELB', DoEnableScroll, 'INIT')
-        self.state.add_transition ('m', 'ELB', None, 'INIT')
+        self.state.add_transition ('m', 'ELB', self.do_sgr, 'INIT')
         self.state.add_transition ('?', 'ELB', None, 'MODECRAP')
         self.state.add_transition_list (string.digits, 'ELB', DoStartNumber, 'NUMBER_1')
         self.state.add_transition_list (string.digits, 'NUMBER_1', DoBuildNumber, 'NUMBER_1')
@@ -241,16 +241,16 @@ class ANSI (term):
         ### It gets worse... the 'm' code can have infinite number of
         ### number;number;number before it. I've never seen more than two,
         ### but the specs say it's allowed. crap!
-        self.state.add_transition ('m', 'NUMBER_1', None, 'INIT')
+        self.state.add_transition ('m', 'NUMBER_1', self.do_sgr, 'INIT')
         ### LED control. Same implementation problem as 'm' code.
-        self.state.add_transition ('q', 'NUMBER_1', None, 'INIT')
+        self.state.add_transition ('q', 'NUMBER_1', self.do_decsca, 'INIT')
 
         # \E[?47h switch to alternate screen
         # \E[?47l restores to normal screen from alternate screen.
         self.state.add_transition_list (string.digits, 'MODECRAP', DoStartNumber, 'MODECRAP_NUM')
         self.state.add_transition_list (string.digits, 'MODECRAP_NUM', DoBuildNumber, 'MODECRAP_NUM')
-        self.state.add_transition ('l', 'MODECRAP_NUM', None, 'INIT')
-        self.state.add_transition ('h', 'MODECRAP_NUM', None, 'INIT')
+        self.state.add_transition ('l', 'MODECRAP_NUM', self.do_modecrap, 'INIT')
+        self.state.add_transition ('h', 'MODECRAP_NUM', self.do_modecrap, 'INIT')
 
 #RM   Reset Mode                Esc [ Ps l                   none
         self.state.add_transition (';', 'NUMBER_1', None, 'SEMICOLON')
@@ -264,9 +264,9 @@ class ANSI (term):
         ### It gets worse... the 'm' code can have infinite number of
         ### number;number;number before it. I've never seen more than two,
         ### but the specs say it's allowed. crap!
-        self.state.add_transition ('m', 'NUMBER_2', None, 'INIT')
+        self.state.add_transition ('m', 'NUMBER_2', self.do_sgr, 'INIT')
         ### LED control. Same problem as 'm' code.
-        self.state.add_transition ('q', 'NUMBER_2', None, 'INIT')
+        self.state.add_transition ('q', 'NUMBER_2', self.do_decsca, 'INIT')
         self.state.add_transition (';', 'NUMBER_2', None, 'SEMICOLON_X')
 
         # Create a state for 'q' and 'm' which allows an infinite number of ignored numbers
@@ -274,8 +274,8 @@ class ANSI (term):
         self.state.add_transition_list (string.digits, 'SEMICOLON_X', DoStartNumber, 'NUMBER_X')
         self.state.add_transition_list (string.digits, 'NUMBER_X', DoBuildNumber, 'NUMBER_X')
         self.state.add_transition_any ('NUMBER_X', DoLog, 'INIT')
-        self.state.add_transition ('m', 'NUMBER_X', None, 'INIT')
-        self.state.add_transition ('q', 'NUMBER_X', None, 'INIT')
+        self.state.add_transition ('m', 'NUMBER_X', self.do_sgr, 'INIT')
+        self.state.add_transition ('q', 'NUMBER_X', self.do_decsca, 'INIT')
         self.state.add_transition (';', 'NUMBER_X', None, 'SEMICOLON_X')
 
     def process (self, c):
@@ -330,3 +330,20 @@ class ANSI (term):
                 self.scroll_up ()
                 self.cursor_home (self.cur_r, 1)
                 self.erase_line()
+
+    def do_sgr (self, fsm):
+        '''Select Graphic Rendition, e.g. color. '''
+        screen = fsm.memory[0]
+        fsm.memory = [screen]
+
+    def do_decsca (self, fsm):
+        '''Select character protection attribute. '''
+        screen = fsm.memory[0]
+        fsm.memory = [screen]
+
+    def do_modecrap (self, fsm):
+        '''Handler for \x1b[?<number>h and \x1b[?<number>l. If anyone
+        wanted to actually use these, they'd need to add more states to the
+        FSM rather than just improve or override this method. '''
+        screen = fsm.memory[0]
+        fsm.memory = [screen]
diff --git a/tests/test_ansi.py b/tests/test_ansi.py
index 3b8d6a9..516509c 100755
--- a/tests/test_ansi.py
+++ b/tests/test_ansi.py
@@ -142,10 +142,25 @@ class ansiTestCase (PexpectTestCase.PexpectTestCase):
 
     def test_number_x(self):
         """Test the FSM state used to handle more than 2 numeric parameters."""
-        s = ANSI.ANSI(1, 20)
+        class TestANSI(ANSI.ANSI):
+            captured_memory = None
+            def do_sgr(self, fsm):
+                assert self.captured_memory is None
+                self.captured_memory = fsm.memory
+
+        s = TestANSI(1, 20)
         s.write('\x1b[0;1;32;45mtest')
         assert str(s) == ('test                ')
-        assert(s.state.memory == [s, '0', '1', '32', '45'])
+        assert s.captured_memory is not None
+        assert s.captured_memory == [s, '0', '1', '32', '45']
+
+    def test_fsm_memory(self):
+        """Test the FSM stack/memory does not have numbers left on it
+        after some sequences with numbers are passed in."""
+        s = ANSI.ANSI(1, 20)
+        s.write('\x1b[0;1;2;3m\x1b[4;5;6;7q\x1b[?8h\x1b[?9ltest')
+        assert str(s) == ('test                ')
+        assert s.state.memory == [s]
 
 if __name__ == '__main__':
     unittest.main()