#919: psutil.Popen: add support for ctx manager protocol

4 files changed, 26 insertions, 12 deletions

diff --git a/CREDITS b/CREDITS
index 7be5f200..4dcb888b 100644
--- a/CREDITS
+++ b/CREDITS
@@ -412,3 +412,7 @@ I: 880
 N: ewedlund
 W: https://github.com/ewedlund
 I: 874
+
+N: Arcadiy Ivanov
+W: https://github.com/arcivanov
+I: 919
diff --git a/HISTORY.rst b/HISTORY.rst
index 26fcc010..6de6703f 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -11,6 +11,8 @@ Bug tracker at https://github.com/giampaolo/psutil/issues
   precise and match "free" cmdline utility.  "available" also takes into
   account LCX containers preventing "available" to overflow "total".
 - #891: procinfo.py script has been updated and provides a lot more info.
+- #919: psutil.Popen() now supports the ctx manager protocol and can be used
+  with the "with" statement.
 
 **Bug fixes**
 
diff --git a/psutil/__init__.py b/psutil/__init__.py
index 8978546d..d748c64d 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -1264,6 +1264,24 @@ class Popen(Process):
     def __dir__(self):
         return sorted(set(dir(Popen) + dir(subprocess.Popen)))
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        if hasattr(self.__subproc, '__exit__'):
+            return self.__subproc.__exit__(*args, **kwargs)
+        else:
+            if self.stdout:
+                self.stdout.close()
+            if self.stderr:
+                self.stderr.close()
+            try:  # Flushing a BufferedWriter may raise an error
+                if self.stdin:
+                    self.stdin.close()
+            finally:
+                # Wait for the process to terminate, to avoid zombies.
+                self.wait()
+
     def __getattribute__(self, name):
         try:
             return object.__getattribute__(self, name)
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index de16d265..dfca1b6c 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -1450,24 +1450,14 @@ class TestProcess(unittest.TestCase):
             self.assertTrue(psutil.pid_exists(0))
 
     def test_Popen(self):
-        # Popen class test
-        # XXX this test causes a ResourceWarning on Python 3 because
-        # psutil.__subproc instance doesn't get propertly freed.
-        # Not sure what to do though.
-        cmd = [PYTHON, "-c", "import time; time.sleep(60);"]
-        proc = psutil.Popen(cmd, stdout=subprocess.PIPE,
-                            stderr=subprocess.PIPE)
-        try:
+        with psutil.Popen([PYTHON, "V"], stdout=subprocess.PIPE,
+                          stderr=subprocess.PIPE) as proc:
             proc.name()
             proc.stdin
             self.assertTrue(hasattr(proc, 'name'))
             self.assertTrue(hasattr(proc, 'stdin'))
             self.assertTrue(dir(proc))
             self.assertRaises(AttributeError, getattr, proc, 'foo')
-        finally:
-            proc.kill()
-            proc.wait()
-            self.assertIsNotNone(proc.returncode)
 
     @unittest.skipUnless(hasattr(psutil.Process, "environ"),
                          "platform not supported")