Merge remote-tracking branch 'bancal/master' into support-method-as-run-event-callback

2 files changed, 55 insertions, 2 deletions

diff --git a/pexpect/__init__.py b/pexpect/__init__.py
index c906e89..8346b5f 100644
--- a/pexpect/__init__.py
+++ b/pexpect/__init__.py
@@ -206,7 +206,8 @@ def _run(command, timeout, withexitstatus, events, extra_args, logfile, cwd,
                 child_result_list.append(child.before)
             if isinstance(responses[index], child.allowed_string_types):
                 child.send(responses[index])
-            elif isinstance(responses[index], types.FunctionType):
+            elif isinstance(responses[index], types.FunctionType) or \
+                 isinstance(responses[index], types.MethodType):
                 callback_result = responses[index](locals())
                 sys.stdout.flush()
                 if isinstance(callback_result, child.allowed_string_types):
@@ -214,7 +215,7 @@ def _run(command, timeout, withexitstatus, events, extra_args, logfile, cwd,
                 elif callback_result:
                     break
             else:
-                raise TypeError('The callback must be a string or function.')
+                raise TypeError('The callback must be a string, function or method.')
             event_count = event_count + 1
         except TIMEOUT:
             child_result_list.append(child.before)
diff --git a/tests/test_run.py b/tests/test_run.py
index c018b4d..286bf69 100755
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -35,6 +35,20 @@ def timeout_callback (d):
         return 1
     return 0
 
+def function_events_callback (d):
+    try:
+        previous_echoed = d["child_result_list"][-1].decode().split("\n")[-2].strip()
+        if previous_echoed.endswith("foo1"):
+            return "echo foo2\n"
+        elif previous_echoed.endswith("foo2"):
+            return "echo foo3\n"
+        elif previous_echoed.endswith("foo3"):
+            return "exit\n"
+        else:
+            raise Exception("Unexpected output {0}".format(previous_echoed))
+    except IndexError:
+        return "echo foo1\n"
+
 class RunFuncTestCase(PexpectTestCase.PexpectTestCase):
     runfunc = staticmethod(pexpect.run)
     cr = b'\r'
@@ -88,6 +102,44 @@ class RunFuncTestCase(PexpectTestCase.PexpectTestCase):
             timeout=10)
         assert exitstatus == 0
 
+    def test_run_function (self):
+        events = [
+            ('GO:', function_events_callback)
+        ]
+    
+        (data, exitstatus) = pexpect.run(
+            'bash --rcfile {0}'.format(self.rcfile),
+            withexitstatus=True,
+            events=events,
+            timeout=10)
+        assert exitstatus == 0
+    
+    def test_run_method (self):
+        events = [
+            ('GO:', self.method_events_callback)
+        ]
+    
+        (data, exitstatus) = pexpect.run(
+            'bash --rcfile {0}'.format(self.rcfile),
+            withexitstatus=True,
+            events=events,
+            timeout=10)
+        assert exitstatus == 0
+    
+    def method_events_callback (self, d):
+        try:
+            previous_echoed = d["child_result_list"][-1].decode().split("\n")[-2].strip()
+            if previous_echoed.endswith("foo1"):
+                return "echo foo2\n"
+            elif previous_echoed.endswith("foo2"):
+                return "echo foo3\n"
+            elif previous_echoed.endswith("foo3"):
+                return "exit\n"
+            else:
+                raise Exception("Unexpected output {0}".format(previous_echoed))
+        except IndexError:
+            return "echo foo1\n"
+
 class RunUnicodeFuncTestCase(RunFuncTestCase):
     runfunc = staticmethod(pexpect.runu)
     cr = b'\r'.decode('ascii')