forked from jupyterlab/jupyterlab
/
process.py
253 lines (203 loc) · 6.82 KB
/
process.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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
# coding: utf-8
"""JupyterLab command handler"""
# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
from __future__ import print_function
import atexit
import logging
import os
import re
import signal
import sys
import threading
import time
import weakref
from tornado import gen
from .jlpmapp import which, subprocess
try:
import pty
except ImportError:
pty = False
if sys.platform == 'win32':
list2cmdline = subprocess.list2cmdline
else:
def list2cmdline(cmd_list):
import pipes
return ' '.join(map(pipes.quote, cmd_list))
logging.basicConfig(format='%(message)s', level=logging.INFO)
class Process(object):
"""A wrapper for a child process.
"""
_procs = weakref.WeakSet()
_pool = None
def __init__(self, cmd, logger=None, cwd=None, kill_event=None,
env=None, quiet=False):
"""Start a subprocess that can be run asynchronously.
Parameters
----------
cmd: list
The command to run.
logger: :class:`~logger.Logger`, optional
The logger instance.
cwd: string, optional
The cwd of the process.
env: dict, optional
The environment for the process.
kill_event: :class:`~threading.Event`, optional
An event used to kill the process operation.
"""
if not isinstance(cmd, (list, tuple)):
raise ValueError('Command must be given as a list')
if kill_event and kill_event.is_set():
raise ValueError('Process aborted')
self.logger = logger = logger or logging.getLogger('jupyterlab')
self._last_line = ''
if not quiet:
self.logger.info('> ' + list2cmdline(cmd))
self.cmd = cmd
self.proc = self._create_process(cwd=cwd, env=env)
self._kill_event = kill_event or threading.Event()
Process._procs.add(self)
def terminate(self):
"""Terminate the process and return the exit code.
"""
proc = self.proc
# Kill the process.
if proc.poll() is None:
os.kill(proc.pid, signal.SIGTERM)
# Wait for the process to close.
try:
proc.wait()
finally:
Process._procs.remove(self)
return proc.returncode
def wait(self):
"""Wait for the process to finish.
Returns
-------
The process exit code.
"""
proc = self.proc
kill_event = self._kill_event
while proc.poll() is None:
if kill_event.is_set():
self.terminate()
raise ValueError('Process was aborted')
time.sleep(1.)
return self.terminate()
@gen.coroutine
def wait_async(self):
"""Asynchronously wait for the process to finish.
"""
proc = self.proc
kill_event = self._kill_event
while proc.poll() is None:
if kill_event.is_set():
self.terminate()
raise ValueError('Process was aborted')
yield gen.sleep(1.)
raise gen.Return(self.terminate())
def _create_process(self, **kwargs):
"""Create the process.
"""
cmd = self.cmd
kwargs.setdefault('stderr', subprocess.STDOUT)
cmd[0] = which(cmd[0], kwargs.get('env'))
if os.name == 'nt':
kwargs['shell'] = True
proc = subprocess.Popen(cmd, **kwargs)
return proc
@classmethod
def _cleanup(cls):
"""Clean up the started subprocesses at exit.
"""
for proc in list(cls._procs):
proc.terminate()
class WatchHelper(Process):
"""A process helper for a watch process.
"""
def __init__(self, cmd, startup_regex, logger=None, cwd=None,
kill_event=None, env=None):
"""Initialize the process helper.
Parameters
----------
cmd: list
The command to run.
startup_regex: string
The regex to wait for at startup.
logger: :class:`~logger.Logger`, optional
The logger instance.
cwd: string, optional
The cwd of the process.
env: dict, optional
The environment for the process.
kill_event: callable, optional
A function to call to check if we should abort.
"""
super(WatchHelper, self).__init__(cmd, logger=logger,
cwd=cwd, kill_event=kill_event, env=env)
if not pty:
self._stdout = self.proc.stdout
while 1:
line = self._stdout.readline().decode('utf-8')
if not line:
raise RuntimeError('Process ended improperly')
print(line.rstrip())
if re.match(startup_regex, line):
break
self._read_thread = threading.Thread(target=self._read_incoming)
self._read_thread.setDaemon(True)
self._read_thread.start()
def terminate(self):
"""Terminate the process.
"""
proc = self.proc
if proc.poll() is None:
if os.name != 'nt':
# Kill the process group if we started a new session.
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
else:
os.kill(proc.pid, signal.SIGTERM)
# Close stdout.
try:
self._stdout.close()
except Exception as e:
pass
# Wait for the process to close.
try:
proc.wait()
finally:
Process._procs.remove(self)
return proc.returncode
def _read_incoming(self):
"""Run in a thread to read stdout and print"""
fileno = self._stdout.fileno()
while 1:
try:
buf = os.read(fileno, 1024)
except OSError as e:
self.logger.debug('Read incoming error %s', e)
return
if not buf:
return
print(buf.decode('utf-8'), end='')
def _create_process(self, **kwargs):
"""Create the watcher helper process.
"""
kwargs['bufsize'] = 0
if pty:
master, slave = pty.openpty()
kwargs['stderr'] = kwargs['stdout'] = slave
kwargs['start_new_session'] = True
self._stdout = os.fdopen(master, 'rb')
else:
kwargs['stdout'] = subprocess.PIPE
if os.name == 'nt':
startupinfo = subprocess.STARTUPINFO()
startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
kwargs['startupinfo'] = startupinfo
kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP
kwargs['shell'] = True
return super(WatchHelper, self)._create_process(**kwargs)
# Register the cleanup handler.
atexit.register(Process._cleanup)