diff --git a/core/__init__.py b/core/__init__.py index d0eff64..1b97b3f 100644 --- a/core/__init__.py +++ b/core/__init__.py @@ -3,6 +3,7 @@ from . import app from . import aux +from . import cmd from . import file from . import misc from . import textui diff --git a/core/cmd.py b/core/cmd.py new file mode 100644 index 0000000..f3ef5e0 --- /dev/null +++ b/core/cmd.py @@ -0,0 +1,142 @@ +#! /bin/env python3 +# encoding: utf-8 + +import queue +import subprocess +import threading + +# -------------------------------------------------- + +def capture_output(stream, fifo): + while True: + chunk = stream.read1(128) + + if chunk: + fifo.put(chunk) + else: + fifo.put(None) # indicates a process has terminated + break + + stream.close() + +class Command: + ''' + A shell command with argument substitution and output capture. + + >>> c = Command(['process.sh', '-x', 'ARGUMENT']) + >>> c.ARGUMENT = 'file.txt' + >>> c.dump() + ['process.sh', '-x', 'file.txt'] + >>> c.run(stderr=False) # capture stdout + >>> c.get_output() + b'some output' + >>> c.get_output() + b'' # there was no output since the last call + >>> c.get_output() + b'more of it\nand some more' + >>> c.get_output() + None # indicates the process ended and there is no more output + + ''' + + def __init__(self, *sequence): + self.__dict__['sequence_original'] = list(sequence) + self.reset() + + def __setattr__(self, name, value): + found = False + + for i, word in enumerate(self.sequence): + if word == name: + self.sequence[i] = value + found = True + + #if not found: + #raise AttributeError('Command has no attribute \'%s\'' %(name)) + + def reset(self): + self.__dict__['sequence'] = list(self.sequence_original) + self.__dict__['thread'] = None + self.__dict__['fifo'] = None + self.__dict__['terminated'] = False + + def dump(self, sequence=None, pretty=False): + out = [] + + if not sequence: + sequence = self.sequence + + for item in sequence: + if type(item) is list: + out += self.dump(item) + elif type(item) is Command: + out += item.dump() + elif item is not None: + out.append(str(item)) + + if pretty: + return [('"%s"' %(a) if ' ' in a else a) for a in out] + else: + return out + + def run(self, stderr=False): + ''' + Executes the command in the current environment. + + stderr capture stderr instead of stdout (can't do both yet) + ''' + + self.__dict__['process'] = subprocess.Popen(self.dump(), stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=1) + self.__dict__['fifo'] = queue.Queue() + self.__dict__['thread'] = threading.Thread( + target=capture_output, + args=(self.process.stderr if stderr else self.process.stdout, self.fifo) + ) + self.__dict__['terminated'] = False + self.thread.daemon = True # thread dies with the program + self.thread.start() + + def get_output(self, blocking=False): + ''' + Returns the output of a currently running process and clears the buffer. + + Returns None if no process is running and no more output is available. + + blocking block until some output is available or the process terminates + ''' + + buffer = [] + + if self.terminated: + return None + + if blocking: + buffer.append(self.fifo.get()) + + while not self.fifo.empty(): + buffer.append(self.fifo.get_nowait()) # FIXME nowait needed? + + if None in buffer: + self.__dict__['terminated'] = True + self.__dict__['returncode'] = self.process.poll() + + if len(buffer) == 1: + return None + else: + assert(buffer[-1] is None) + del buffer[-1] + + return b''.join(buffer) + + def get_all_output(self): + buffer = [] + + while True: + chunk = self.get_output(blocking=True) + + if chunk is None: + break + else: + buffer.append(chunk) + + return b''.join(buffer) if buffer else None \ No newline at end of file