/
ShellSubprocess.py
112 lines (86 loc) · 2.94 KB
/
ShellSubprocess.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
"""
ShellSubprocess
This module is used to create and manage a bash shell in a subprocess.
Once instantiated, the shell can execute commands that will write to stdout and
stderr. Both buffers can be read with the read_stdout and read_stderr methods.
The subprocess should be terminated with the close method or used in a context
manager.
"""
import subprocess
from time import sleep
from tempfile import TemporaryFile, NamedTemporaryFile
class ShellSubprocess(object):
"""
Interactive shell running in a persistent process
"""
def __init__(self):
self.stdout = TemporaryFile()
self.stderr = TemporaryFile()
self.process = subprocess.Popen(
"/bin/bash", stdin=subprocess.PIPE, stdout=self.stdout, stderr=self.stderr
)
self.stdin = self.process.stdin
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
self.close()
def execute(self, cmd):
"""Execute a command in the shell.
Args:
cmd (string, list): command to execute
"""
if isinstance(cmd, list):
cmd = subprocess.list2cmdline(cmd)
self.stdin.write(cmd.encode() + b"\n")
self.stdin.flush()
def run(self, cmd):
"""Execute a command in the shell an return the stdout, stderr and exit code.
NOTE: This method does not write to stdout or stderr buffers.
Args:
cmd (string, list): command to execute
"""
with NamedTemporaryFile() as outfile, NamedTemporaryFile() as errfile, NamedTemporaryFile() as exitfile:
self.execute(
f"({cmd}) 1> {outfile.name} 2> {errfile.name}; echo $? > {exitfile.name}"
)
wait_time = 0.0001
while True:
exitfile.seek(0)
exit_code = exitfile.read()
if exit_code:
break
sleep(wait_time)
wait_time *= 2
outfile.seek(0)
out = outfile.read()
errfile.seek(0)
err = errfile.read()
exit_code = int(exit_code.strip())
return out, err, exit_code
def read_stdout(self):
"""Read stdout from the shell
Returns:
bytes: stdout
"""
self.stdout.seek(0)
return self.stdout.read()
def read_stderr(self):
"""Read stderr from the shell
Returns:
bytes: stderr
"""
self.stderr.seek(0)
return self.stderr.read()
def clear_stdout(self):
"""Clear stdout from the shell"""
self.stdout.seek(0)
self.stdout.truncate()
def clear_stderr(self):
"""Clear stderr from the shell"""
self.stderr.seek(0)
self.stderr.truncate()
def close(self):
"""Close the shell"""
self.process.terminate()
self.stdout.close()
self.stderr.close()