Add timeouts into the ExecutionModifiers

plumbum/commands/modifiers.py

@@ -2,6 +2,7 @@
 from select import select
 from subprocess import PIPE
 import sys
+import time
 from itertools import chain
 
 from plumbum.commands.processes import run_proc, ProcessExecutionError
@@ -97,14 +98,15 @@ class BG(ExecutionModifier):
        every once in a while, using a monitoring thread/reactor in the background.
        For more info, see `#48 <https://github.com/tomerfiliba/plumbum/issues/48>`_
     """
-    __slots__ = ("retcode", "kargs")
+    __slots__ = ("retcode", "kargs", "timeout")
 
-    def __init__(self, retcode=0, **kargs):
+    def __init__(self, retcode=0, timeout=None, **kargs):
         self.retcode = retcode
         self.kargs = kargs
+        self.timeout = timeout
 
     def __rand__(self, cmd):
-        return Future(cmd.popen(**self.kargs), self.retcode)
+        return Future(cmd.popen(**self.kargs), self.retcode, timeout=self.timeout)
 
 BG = BG()
 
@@ -139,13 +141,15 @@ class FG(ExecutionModifier):
         vim & FG       # run vim in the foreground, expecting an exit code of 0
         vim & FG(7)    # run vim in the foreground, expecting an exit code of 7
     """
-    __slots__ = ("retcode",)
+    __slots__ = ("retcode", "timeout")
 
-    def __init__(self, retcode=0):
+    def __init__(self, retcode=0, timeout=None):
         self.retcode = retcode
+        self.timeout = timeout
 
     def __rand__(self, cmd):
-        cmd(retcode = self.retcode, stdin = None, stdout = None, stderr = None)
+        cmd(retcode = self.retcode, stdin = None, stdout = None, stderr = None,
+            timeout = self.timeout)
 
 FG = FG()
 
@@ -162,19 +166,21 @@ class TEE(ExecutionModifier):
     Returns a tuple of (return code, stdout, stderr), just like ``run()``.
     """
 
-    __slots__ = ("retcode", "buffered")
+    __slots__ = ("retcode", "buffered", "timeout")
 
-    def __init__(self, retcode=0, buffered=True):
+    def __init__(self, retcode=0, buffered=True, timeout=None):
         """`retcode` is the return code to expect to mean "success".  Set
         `buffered` to False to disable line-buffering the output, which may
         cause stdout and stderr to become more entangled than usual.
         """
         self.retcode = retcode
         self.buffered = buffered
+        self.timeout = timeout
 
 
     def __rand__(self, cmd):
-        with cmd.bgrun(retcode=self.retcode, stdin=None, stdout=PIPE, stderr=PIPE) as p:
+        with cmd.bgrun(retcode=self.retcode, stdin=None, stdout=PIPE, stderr=PIPE,
+                       timeout=self.timeout) as p:
             outbuf = []
             errbuf = []
             out = p.stdout
@@ -224,14 +230,15 @@ class TF(ExecutionModifier):
         local['touch']['/root/test'] & TF(FG=True) * Returns False, will show error message
     """
 
-    __slots__ = ("retcode", "FG")
+    __slots__ = ("retcode", "FG", "timeout")
 
-    def __init__(self, retcode=0, FG=False):
+    def __init__(self, retcode=0, FG=False, timeout=None):
         """`retcode` is the return code to expect to mean "success".  Set
         `FG` to True to run in the foreground.
         """
         self.retcode = retcode
         self.FG = FG
+        self.timeout = timeout
 
     @classmethod
     def __call__(cls, *args, **kwargs):
@@ -240,9 +247,10 @@ def __call__(cls, *args, **kwargs):
     def __rand__(self, cmd):
         try:
             if self.FG:
-                cmd(retcode = self.retcode, stdin = None, stdout = None, stderr = None)
+                cmd(retcode = self.retcode, stdin = None, stdout = None, stderr = None,
+                    timeout = self.timeout)
             else:
-                cmd(retcode = self.retcode)
+                cmd(retcode = self.retcode, timeout = self.timeout)
             return True
         except ProcessExecutionError:
             return False
@@ -264,22 +272,24 @@ class RETCODE(ExecutionModifier):
         local['touch']['/root/test'] & RETCODE(FG=True) * Returns 1, will show error message
     """
 
-    __slots__ = ("foreground",)
+    __slots__ = ("foreground", "timeout")
 
-    def __init__(self,  FG=False):
+    def __init__(self,  FG=False, timeout=None):
         """`FG` to True to run in the foreground.
         """
         self.foreground = FG
+        self.timeout = timeout
 
     @classmethod
     def __call__(cls, *args, **kwargs):
         return cls(*args, **kwargs)
 
     def __rand__(self, cmd):
             if self.foreground:
-                return cmd.run(retcode = None, stdin = None, stdout = None, stderr = None)[0]
+                return cmd.run(retcode = None, stdin = None, stdout = None, stderr = None,
+                               timeout = self.timeout)[0]
             else:
-                return cmd.run(retcode = None)[0]
+                return cmd.run(retcode = None, timeout = self.timeout)[0]
 
 RETCODE = RETCODE()
 
@@ -317,7 +327,6 @@ def __init__(self, cwd='.', stdout='nohup.out', stderr=None, append=True):
         self.stderr = stderr
         self.append = append
 
-
     def __rand__(self, cmd):
         if isinstance(cmd, StdoutRedirection):
             stdout = cmd.file

tests/test_validate.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division
+import sys
 import pytest
 from plumbum import cli
 from plumbum.lib import captured_stdout
@@ -41,13 +42,13 @@ def main(selfy, x, y, *g):
 
         assert Try.main.positional == [abs, str]
         assert Try.main.positional_varargs == int
-        
+
     def test_defaults(self):
         class Try(object):
             @cli.positional(abs, str)
             def main(selfy, x, y = 'hello'):
                 pass
-            
+
         assert Try.main.positional == [abs, str]
 
 
@@ -56,8 +57,8 @@ def test_prog(self, capsys):
         class MainValidator(cli.Application):
             @cli.positional(int, int, int)
             def main(self, myint, myint2, *mylist):
-                print(repr(myint), myint2, mylist)        
-        
+                print(repr(myint), myint2, mylist)
+
         _, rc = MainValidator.run(["prog", "1", "2", '3', '4', '5'], exit = False)
         assert rc == 0
         assert "1 2 (3, 4, 5)" == capsys.readouterr()[0].strip()
@@ -69,28 +70,33 @@ class MainValidator(cli.Application):
             def main(self, myint, myint2, *mylist):
                 print(myint, myint2, mylist)
         _, rc = MainValidator.run(["prog", "1.2", "2", '3', '4', '5'], exit = False)
-        
+
         assert rc == 2
         value = capsys.readouterr()[0].strip()
-        assert "'int'>, not '1.2':" in value
-        assert " 'int'>, not '1.2':" in value
+        # In Python > 2.6, sys.version_info is a namedtuple.
+        if sys.version_info[0] == 3 and sys.version_info[1] >= 6:
+            assert "<class 'int' " in value
+            assert ">, not '1.2':" in value
+        else:
+            assert "'int'>, not '1.2':" in value
+            assert " 'int'>, not '1.2':" in value
         assert '''ValueError("invalid literal for int() with base 10: '1.2'"''' in value
 
     def test_defaults(self, capsys):
         class MainValidator(cli.Application):
             @cli.positional(int, int)
             def main(self, myint, myint2=2):
-                print(repr(myint), repr(myint2))        
-        
+                print(repr(myint), repr(myint2))
+
         _, rc = MainValidator.run(["prog", "1"], exit = False)
         assert rc == 0
         assert "1 2" == capsys.readouterr()[0].strip()
-        
+
         _, rc = MainValidator.run(["prog", "1", "3"], exit = False)
         assert rc == 0
         assert "1 3" ==  capsys.readouterr()[0].strip()
-        
-        
+
+
 # Unfortionatly, Py3 anotations are a syntax error in Py2, so using exec to add test for Py3
 if six.PY3:
     exec("""