SumProject/rt-thread/tools/win32spawn.py

#
# File      : win32spawn.py
# This file is part of RT-Thread RTOS
# COPYRIGHT (C) 2006 - 2015, RT-Thread Development Team
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License along
#  with this program; if not, write to the Free Software Foundation, Inc.,
#  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# Change Logs:
# Date           Author       Notes
# 2015-01-20     Bernard      Add copyright information
#

import os
import threading
import sys

_PY2 = sys.version_info[0] < 3
if _PY2:
    import Queue
else:
    import queue as Queue

# Windows import
import win32file
import win32pipe
import win32api
import win32con
import win32security
import win32process
import win32event

class Win32Spawn(object):
    def __init__(self, cmd, shell=False):
        self.queue = Queue.Queue()
        self.is_terminated = False
        self.wake_up_event = win32event.CreateEvent(None, 0, 0, None)

        exec_dir = os.getcwd()
        comspec = os.environ.get("COMSPEC", "cmd.exe")
        cmd = comspec + ' /c ' + cmd

        win32event.ResetEvent(self.wake_up_event)

        currproc = win32api.GetCurrentProcess()

        sa = win32security.SECURITY_ATTRIBUTES()
        sa.bInheritHandle = 1

        child_stdout_rd, child_stdout_wr = win32pipe.CreatePipe(sa, 0)
        child_stdout_rd_dup = win32api.DuplicateHandle(currproc, child_stdout_rd, currproc, 0, 0, win32con.DUPLICATE_SAME_ACCESS)
        win32file.CloseHandle(child_stdout_rd)

        child_stderr_rd, child_stderr_wr = win32pipe.CreatePipe(sa, 0)
        child_stderr_rd_dup = win32api.DuplicateHandle(currproc, child_stderr_rd, currproc, 0, 0, win32con.DUPLICATE_SAME_ACCESS)
        win32file.CloseHandle(child_stderr_rd)

        child_stdin_rd, child_stdin_wr = win32pipe.CreatePipe(sa, 0)
        child_stdin_wr_dup = win32api.DuplicateHandle(currproc, child_stdin_wr, currproc, 0, 0, win32con.DUPLICATE_SAME_ACCESS)
        win32file.CloseHandle(child_stdin_wr)

        startup_info = win32process.STARTUPINFO()
        startup_info.hStdInput = child_stdin_rd
        startup_info.hStdOutput = child_stdout_wr
        startup_info.hStdError = child_stderr_wr
        startup_info.dwFlags = win32process.STARTF_USESTDHANDLES

        cr_flags = 0
        cr_flags = win32process.CREATE_NEW_PROCESS_GROUP

        env = os.environ.copy()
        self.h_process, h_thread, dw_pid, dw_tid = win32process.CreateProcess(None, cmd, None, None, 1,
                                                                              cr_flags, env, os.path.abspath(exec_dir),
                                                                              startup_info)

        win32api.CloseHandle(h_thread)

        win32file.CloseHandle(child_stdin_rd)
        win32file.CloseHandle(child_stdout_wr)
        win32file.CloseHandle(child_stderr_wr)

        self.__child_stdout = child_stdout_rd_dup
        self.__child_stderr = child_stderr_rd_dup
        self.__child_stdin = child_stdin_wr_dup

        self.exit_code = -1

    def close(self):
        win32file.CloseHandle(self.__child_stdout)
        win32file.CloseHandle(self.__child_stderr)
        win32file.CloseHandle(self.__child_stdin)
        win32api.CloseHandle(self.h_process)
        win32api.CloseHandle(self.wake_up_event)

    def kill_subprocess():
        win32event.SetEvent(self.wake_up_event)

    def sleep(secs):
        win32event.ResetEvent(self.wake_up_event)
        timeout = int(1000 * secs)
        val = win32event.WaitForSingleObject(self.wake_up_event, timeout)
        if val == win32event.WAIT_TIMEOUT:
            return True
        else:
            # The wake_up_event must have been signalled
            return False

    def get(self, block=True, timeout=None):
        return self.queue.get(block=block, timeout=timeout)

    def qsize(self):
        return self.queue.qsize()

    def __wait_for_child(self):
        # kick off threads to read from stdout and stderr of the child process
        threading.Thread(target=self.__do_read, args=(self.__child_stdout, )).start()
        threading.Thread(target=self.__do_read, args=(self.__child_stderr, )).start()

        while True:
            # block waiting for the process to finish or the interrupt to happen
            handles = (self.wake_up_event, self.h_process)
            val = win32event.WaitForMultipleObjects(handles, 0, win32event.INFINITE)

            if val >= win32event.WAIT_OBJECT_0 and val < win32event.WAIT_OBJECT_0 + len(handles):
                handle = handles[val - win32event.WAIT_OBJECT_0]
                if handle == self.wake_up_event:
                    win32api.TerminateProcess(self.h_process, 1)
                    win32event.ResetEvent(self.wake_up_event)
                    return False
                elif handle == self.h_process:
                    # the process has ended naturally
                    return True
                else:
                    assert False, "Unknown handle fired"
            else:
                assert False, "Unexpected return from WaitForMultipleObjects"

    # Wait for job to finish. Since this method blocks, it can to be called from another thread.
    # If the application wants to kill the process, it should call kill_subprocess().
    def wait(self):
        if not self.__wait_for_child():
            # it's been killed
            result = False
        else:
            # normal termination
            self.exit_code = win32process.GetExitCodeProcess(self.h_process)
            result = self.exit_code == 0
        self.close()
        self.is_terminated = True

        return result

    # This method gets called on a worker thread to read from either a stderr
    # or stdout thread from the child process.
    def __do_read(self, handle):
        bytesToRead = 1024
        while 1:
            try:
                finished = 0
                hr, data = win32file.ReadFile(handle, bytesToRead, None)
                if data:
                    self.queue.put_nowait(data)
            except win32api.error:
                finished = 1

            if finished:
                return

    def start_pipe(self):
        def worker(pipe):
            return pipe.wait()

        thrd = threading.Thread(target=worker, args=(self, ))
        thrd.start()