execution.py

Go to the documentation of this file.

# Copyright 2005, 2006  Timo Savola
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.

'''Asynchronous command execution.  Integrated into the GLib main loop.'''

import sys
import os
import fcntl
import signal
import gobject

import sets
set = sets.Set

##
## Spawn

def spawn(args, vars=None, workdir=None, channels=[]):
        assert type(args) == list and len(args) > 0
        assert not vars or type(vars) == dict

        pipes = [os.pipe() for entry in channels]

        pid = os.fork()
        if pid == 0:
                for fd in xrange(1024):
                        try:
                                fcntl.fcntl(fd, fcntl.F_SETFD, fcntl.FD_CLOEXEC)
                        except:
                                pass

                try:
                        if vars:
                                os.environ.update(vars)

                        if workdir:
                                os.chdir(workdir)

                        for fileno, channel in channels:
                                fd = pipes.pop(0)

                                if channel.output:
                                        master, slave = fd
                                else:
                                        slave, master = fd

                                os.close(master)
                                if slave != fileno:
                                        os.dup2(slave, fileno)
                                        os.close(slave)

                        os.execvp(args[0], args)

                except Exception, e:
                        try:
                                print >>sys.stderr, '%s:' % args[0],
                        except:
                                pass

                        try:
                                print >>sys.stderr, e
                        except:
                                pass

                os._exit(127)

        for fileno, channel in channels:
                fd = pipes.pop(0)

                if channel.output:
                        master, slave = fd
                else:
                        slave, master = fd

                os.close(slave)
                channel.open(master)

        return pid

##
## Channel

class Channel (gobject.GObject):
        def __init__(self, condition=0, callback=None):
                gobject.GObject.__init__(self)

                self.file = None
                self.condition = condition
                self.callback = callback

        def open(self, file):
                assert not self.file

                self.file = file

                mask = self.condition | gobject.IO_ERR | gobject.IO_HUP
                self.id = gobject.io_add_watch(self.file, mask, self._watch)

                self.emit('opened')

        def _watch(self, file, condition):
                if condition & gobject.IO_ERR:
                        close = True
                elif condition & gobject.IO_HUP:
                        if condition & self.condition:
                                self.callback(file, flush=True)

                        close = True
                else:
                        if condition & self.condition:
                                self.callback(file, flush=False)

                        close = False

                if close:
                        self.file = None
                        self.id = None

                        try:
                                file.close()
                        except IOError:
                                pass

                        self.emit('closed')

                        return False
                else:
                        return True

        def close(self):
                if self.id:
                        try:
                                gobject.remove_source(self.id)
                        except:
                                pass

                        try:
                                self.file.close()
                        except:
                                pass

                        self.id = None
                        self.file = None

                        self.emit('closed')

gobject.type_register(Channel)
gobject.signal_new('opened', Channel, 0, gobject.TYPE_NONE, [])
gobject.signal_new('closed', Channel, 0, gobject.TYPE_NONE, [])

class WriteChannel (Channel):
        output = False

        def __init__(self):
                Channel.__init__(self)

        def open(self, fd):
                file = os.fdopen(fd, 'w')
                Channel.open(self, file)

        def write(self, data):
                assert self.file
                self.file.write(data)

        def flush(self):
                assert self.file
                self.file.flush()

gobject.type_register(WriteChannel)

class ReadChannel (Channel):
        output = True

        def __init__(self, consumers=[]):
                Channel.__init__(self, gobject.IO_IN, self._read)
                self.consumers = set(consumers)

        def open(self, fd):
                fcntl.fcntl(fd, fcntl.F_SETFL, os.O_NONBLOCK)
                file = os.fdopen(fd, 'r')
                Channel.open(self, file)

        def _read(self, file, flush):
                while True:
                        data = file.read()
                        if data:
                                for callback in self.consumers:
                                        callback(data)

                        if not flush or not data:
                                break

        def write(self, data):
                for callback in self.consumers:
                        callback(data)


gobject.type_register(ReadChannel)

##
## Executor

class Executor (gobject.GObject):
        def __init__(self):
                gobject.GObject.__init__(self)

                self.channels = []
                self.open_channels = set()

                self.pid = None
                self.watch = None
                self.status = None

        def redirect_input(self, fileno):
                return self.redirect(fileno, WriteChannel())

        def redirect_output(self, fileno):
                return self.redirect(fileno, ReadChannel())

        def redirect(self, fileno, channel):
                channel.connect('opened', self.on_channel_opened)
                channel.connect('closed', self.on_channel_closed)

                self.channels.append((fileno, channel))
                return channel

        def start(self, job):
                assert not self.is_alive()

                self.status = None
                self.pid = job(self.channels)
                self.watch = gobject.child_watch_add(self.pid, self.on_child_exited)

        def stop(self):
                try:
                        self.kill(signal.SIGTERM)
                except:
                        pass

        def kill(self, signal):
                if self.pid is not None:
                        os.kill(self.pid, signal)

        def is_alive(self):
                return (self.pid is not None) or self.open_channels

        def on_channel_opened(self, channel):
                self.open_channels.add(channel)

        def on_channel_closed(self, channel):
                self.open_channels.remove(channel)

                self._maybe_emit_stopped()

        def on_child_exited(self, pid, status):
                self.pid = None
                self.watch = None
                self.status = status

                self._maybe_emit_stopped()

        def _maybe_emit_stopped(self):
                if (self.pid is None) and (not self.open_channels):
                        self.emit('stopped', self.status)

gobject.type_register(Executor)
gobject.signal_new('stopped', Executor, 0, gobject.TYPE_NONE, [gobject.TYPE_INT])

##
## Job

class Job (gobject.GObject):
        '''Describes command parameters.  The 'args', 'vars' and 'workdir' attributes are publicly
           accessible.'''

        def __init__(self, args, vars=None, workdir=None):
                '''ARGS is the argument vector starting with the program name or path.
                   VARS is a dictionary of additional environment variables.
                   WORKDIR is the current working directory of the command.'''

                gobject.GObject.__init__(self)

                self.args = args
                self.vars = vars
                self.workdir = workdir

        def __iter__(self):
                return iter((self.args, self.vars, self.workdir))

        def __repr__(self):
                return '%s(%s, %s, %s)' % (Job.__name__, repr(self.args), repr(self.vars), repr(self.workdir))

        def __call__(self, channels=[]):
                return spawn(self.args, self.vars, self.workdir, channels)

gobject.type_register(Job)

##
## Hook

class Hook (gobject.GObject):
        def __init__(self, func):
                gobject.GObject.__init__(self)
                self.func = func

        def __repr__(self):
                return '%s(%s)' % (Job.__name__, repr(self.func))

        def __str__(self):
                return self.func.func_name

        def __call__(self):
                return self.func()

gobject.type_register(Hook)

##
## Batch

class Batch (gobject.GObject):
        '''Sequence of jobs that can be monitored via signals.  The "job-started" signal is emitted
           when a new job starts executing (callbacks are invoked with the Job as a parameter).
           "finished" is emitted when a job fails or the whole batch finishes (callbacks are
           invoked with the exit status as a parameter).'''

        def __init__(self, jobs=None):
                '''Initialise an empty batch.'''

                gobject.GObject.__init__(self)
                self.connect('job-started', self.on_started)

                self.executor = Executor()
                self.executor.connect('stopped', self.on_executor_stopped)

                self.channels = {}

                if jobs:
                        self.jobs = jobs[:]
                else:
                        self.jobs = []

                self.position = 0

        def get_input(self, fileno=0):
                '''Return (and possibly initialize) input channel.'''

                channel = self.channels.get(fileno)
                if channel is None:
                        assert not self.executor.is_alive()

                        channel = self.executor.redirect_input(fileno)
                        self.channels[fileno] = channel
                else:
                        assert not channel.output

                return channel

        def get_output(self, fileno=1):
                '''Return (and possibly initialize) output channel.'''

                channel = self.channels.get(fileno)
                if channel is None:
                        assert not self.executor.is_alive()

                        channel = self.executor.redirect_output(fileno)
                        self.channels[fileno] = channel
                else:
                        assert channel.output

                return channel

        def get_error(self):
                '''Return (and possibly initialize) error output channel.'''

                return self.get_output(2)

        def add(self, job):
                '''Add JOB to be executed in this batch.'''

                self.jobs.append(job)

        def extend(self, jobs):
                '''Add JOBS to be executed in this batch.'''

                self.jobs.extend(jobs)

        def remove(self, job):
                '''Remove JOB from this batch.  The batch must not be running.'''

                assert not self.executor.is_alive()
                self.jobs.remove(job)

        def clear(self):
                '''Remove all jobs from this batch.  The batch must not be running.'''

                assert not self.executor.is_alive()
                del self.jobs[:]

        def rewind(self):
                '''Batch position will be reset.  The batch must not be running.'''

                assert not self.executor.is_alive()
                self.position = 0

        def start(self):
                '''Execute the remaining jobs.'''

                while self.position < len(self.jobs):
                        job = self.jobs[self.position]

                        if isinstance(job, Job):
                                self.emit('job-started', job)
                                return
                        else:
                                original = sys.stdout, sys.stderr
                                try:
                                        sys.stdout = self.get_output()
                                        sys.stderr = self.get_error()

                                        job()
                                finally:
                                        sys.stdout, sys.stderr = original

                                self.position += 1

                self.emit('finished', 0)

        def stop(self):
                '''Stop the execution.'''

                self.executor.stop()

        def is_alive(self):
                '''Check if one of the jobs is currently running.'''

                return self.executor.is_alive()

        def on_started(self, batch, job):
                self.executor.start(job)

        def on_executor_stopped(self, executor, status):
                if status == 0:
                        self.position += 1
                        self.start()
                else:
                        self.emit('finished', status)

gobject.type_register(Batch)
gobject.signal_new('job-started', Batch, 0, gobject.TYPE_NONE, [gobject.GObject])
gobject.signal_new('finished', Batch, 0, gobject.TYPE_NONE, [gobject.TYPE_INT])

---