149 lines
No EOL
3.4 KiB
Python
149 lines
No EOL
3.4 KiB
Python
#! /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()
|
|
|
|
def char_in_str(chars, string):
|
|
for char in chars:
|
|
if char in string:
|
|
return True
|
|
|
|
return False
|
|
|
|
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 char_in_str(" ()[];\\", 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 |