pygame / async_sub.py

################################################################################
"""

Modification of http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/440554

"""

#################################### IMPORTS ###################################

import os
import subprocess
import errno
import time
import sys
import unittest
import tempfile

if subprocess.mswindows:
    # sys.path.append('async_libs.zip')
    from win32file import ReadFile, WriteFile
    from win32pipe import PeekNamedPipe

    import msvcrt

    def _call_proc(cmd):
        return subprocess.Popen (
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell = 1,
        )

    win32_kill_commands = (
        ('pskill', 'pskill -t %s'),
        ('taskkill /?', 'taskkill /F /T /PID %s'),  # /? so no err code
    )

    for test_cmd, kill_cmd in win32_kill_commands:
        if _call_proc(test_cmd).wait() is not 1:
            os.kill = lambda pid: _call_proc(kill_cmd % pid)
            break

        else: os.kill = None

    if os.kill is None:
        raise SystemExit('No way of killing unruly processes. Try installing '
                         'sysinternals pskill and placing on %PATH%.')

else:
    import select
    import fcntl

################################### CONSTANTS ##################################

PIPE = subprocess.PIPE

################################################################################


class Popen(subprocess.Popen):
    def recv(self, maxsize=None):
        return self._recv('stdout', maxsize)
    
    def recv_err(self, maxsize=None):
        return self._recv('stderr', maxsize)

    def send_recv(self, input='', maxsize=None):
        return self.send(input), self.recv(maxsize), self.recv_err(maxsize)
    
    def read_async(self,  wait=.1, e=1, tr=5, stderr=0):
        if tr < 1:
            tr = 1
        x = time.time()+ wait
        y = []
        r = ''
        pr = self.recv
        if stderr:
            pr = self.recv_err
        while time.time() < x or r:
            r = pr()
            if r is None:
                if e:
                    raise Exception("Other end disconnected!")
                else:
                    break
            elif r:
                y.append(r)
            else:
                time.sleep(max((x-time.time())/tr, 0))
        return ''.join(y)
        
    def send_all(self, data):
        data += '\r\n'
        while len(data):
            sent = self.send(data)
            if sent is None:
                raise Exception("Other end disconnected!")
            data = buffer(data, sent)
    
    def get_conn_maxsize(self, which, maxsize):
        if maxsize is None:
            maxsize = 1024
        elif maxsize < 1:
            maxsize = 1
        return getattr(self, which), maxsize
    
    def _close(self, which):
        getattr(self, which).close()
        setattr(self, which, None)
    
    if subprocess.mswindows:
        def send(self, input):
            if not self.stdin:
                return None

            try:
                x = msvcrt.get_osfhandle(self.stdin.fileno())
                (errCode, written) = WriteFile(x, input)
            except ValueError:
                return self._close('stdin')
            except (subprocess.pywintypes.error, Exception), why:
                if why[0] in (109, errno.ESHUTDOWN):
                    return self._close('stdin')
                raise

            return written

        def _recv(self, which, maxsize):
            conn, maxsize = self.get_conn_maxsize(which, maxsize)
            if conn is None:
                return None
            
            try:
                x = msvcrt.get_osfhandle(conn.fileno())
                (read, nAvail, nMessage) = PeekNamedPipe(x, 0)
                if maxsize < nAvail:
                    nAvail = maxsize
                if nAvail > 0:
                    (errCode, read) = ReadFile(x, nAvail, None)
            except ValueError:
                return self._close(which)
            except (subprocess.pywintypes.error, Exception), why:
                if why[0] in (109, errno.ESHUTDOWN):
                    return self._close(which)
                raise
            
            if self.universal_newlines:
                read = self._translate_newlines(read)
            return read

    else:
        def send(self, input):
            if not self.stdin:
                return None

            if not select.select([], [self.stdin], [], 0)[1]:
                return 0

            try:
                written = os.write(self.stdin.fileno(), input)
            except OSError, why:
                if why[0] == errno.EPIPE: #broken pipe
                    return self._close('stdin')
                raise

            return written

        def _recv(self, which, maxsize):
            conn, maxsize = self.get_conn_maxsize(which, maxsize)
            if conn is None:
                return None
            
            flags = fcntl.fcntl(conn, fcntl.F_GETFL)
            if not conn.closed:
                fcntl.fcntl(conn, fcntl.F_SETFL, flags| os.O_NONBLOCK)
            
            try:
                if not select.select([conn], [], [], 0)[0]:
                    return ''
                
                r = conn.read(maxsize)
                if not r:
                    return self._close(which)
    
                if self.universal_newlines:
                    r = self._translate_newlines(r)
                return r
            finally:
                if not conn.closed:
                    fcntl.fcntl(conn, fcntl.F_SETFL, flags)

################################################################################

def proc_in_time_or_kill(cmd, time_out):
    proc = Popen (
        cmd, shell = True, bufsize = -1,
        stdin = subprocess.PIPE, stdout = subprocess.PIPE, 
        stderr = subprocess.STDOUT, universal_newlines = 1
    )

    ret_code = None
    response = []

    t = time.time()
    while ret_code is None and ((time.time() -t) < time_out):
        ret_code = proc.poll()
        response += [proc.read_async(wait=0.1, e=0)]

    if ret_code is None:
        os.kill(proc.pid)
        ret_code = '"Process timed out (time_out = %s secs)"' % time_out

    return ret_code, ''.join(response)

################################################################################

class AsyncTest(unittest.TestCase):
    def test_proc_in_time_or_kill(self):
        temp_dir   = tempfile.mkdtemp()
        temp_file  = os.path.join(temp_dir, 'xxxxxx.py')
        fh = open(temp_file, 'w')
        try:
            fh.write('while True: print "GST"')
        finally:
            fh.close()

        ret_code, response = proc_in_time_or_kill(temp_file, time_out=1)

        self.assert_(
            ret_code.startswith('"Process timed out') and
            "GST" in response
        )

################################################################################

def _example():
    if sys.platform == 'win32':
        shell, commands, tail = ('cmd', ('echo "hello"', 'echo "HELLO WORLD"'), '\r\n')
    else:
        shell, commands, tail = ('sh', ('ls', 'echo HELLO WORLD'), '\n')
    
    a = Popen(shell, stdin=PIPE, stdout=PIPE)
    print a.read_async(),
    for cmd in commands:
        a.send_all(cmd + tail)
        print a.read_async(),
    a.send_all('exit' + tail)
    print a.read_async(e=0)
    a.wait()

################################################################################
    
if __name__ == '__main__':
    if 1:
        unittest.main()
    else:
        _example()
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.