Allow spawn() and friends to be used as context managers

Closes gh-111

2 files changed, 19 insertions, 0 deletions

diff --git a/pexpect/spawnbase.py b/pexpect/spawnbase.py
index b1bc2f3..d79c5c0 100644
--- a/pexpect/spawnbase.py
+++ b/pexpect/spawnbase.py
@@ -443,6 +443,15 @@ class SpawnBase(object):
         """Overridden in subclass using tty"""
         return False
 
+    # For 'with spawn(...) as child:'
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, etype, evalue, tb):
+        # We rely on subclasses to implement close(). If they don't, it's not
+        # clear what a context manager should do.
+        self.close()
+
 class SpawnBaseUnicode(SpawnBase):
     if PY3:
         string_type = str
diff --git a/tests/test_misc.py b/tests/test_misc.py
index 28df570..792afa2 100755
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -141,6 +141,16 @@ class TestCaseMisc(PexpectTestCase.PexpectTestCase):
         with self.assertRaises(pexpect.EOF):
             child.expect('the unexpected')
 
+    def test_with(self):
+        "spawn can be used as a context manager"
+        with pexpect.spawn(sys.executable + ' echo_w_prompt.py') as p:
+            p.expect('<in >')
+            p.sendline(b'alpha')
+            p.expect(b'<out>alpha')
+            assert p.isalive()
+        
+        assert not p.isalive()
+
     def test_terminate(self):
         " test force terminate always succeeds (SIGKILL). "
         child = pexpect.spawn('cat')