Merge branch 'master' into python3
[pdm.git] / pdm / cli.py
CommitLineData
ebd7a8ba 1"""Python Daemon Management -- Client functions
7f97a47e 2
ebd7a8ba
FT
3This module implements the client part of the PDM protocols. The
4primary objects of interest are the replclient and perfclient classes,
5which implement support for their respective protocols. See their
6documentation for details.
7f97a47e
FT
7"""
8
9import socket, pickle, struct, select, threading
10
ac4f5166 11__all__ = ["client", "replclient", "perfclient"]
7f97a47e
FT
12
13class protoerr(Exception):
73c1ef8b 14 """Raised on protocol errors"""
7f97a47e
FT
15 pass
16
17def resolve(spec):
18 if isinstance(spec, socket.socket):
19 return spec
20 sk = None
21 try:
c30be9c9
FT
22 if ":" in spec:
23 p = spec.rindex(":")
24 first, second = spec[:p], spec[p + 1:]
25 if "/" in second:
16e7fd3d 26 from . import sshsock
c30be9c9
FT
27 sk = sshsock.sshsocket(first, second)
28 else:
29 sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
30 sk.connect((first, second))
31 elif "/" in spec:
7f97a47e
FT
32 sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
33 sk.connect(spec)
34 elif spec.isdigit():
35 sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
36 sk.connect(("localhost", int(spec)))
7f97a47e
FT
37 else:
38 raise Exception("Unknown target specification %r" % spec)
39 rv = sk
40 sk = None
41 finally:
42 if sk is not None: sk.close()
43 return rv
44
45class client(object):
73c1ef8b
FT
46 """PDM client
47
48 This class provides general facilities to speak to PDM servers,
49 and is mainly intended to be subclassed to provide for the
50 specific protocols, such as replclient and perfclient do.
51
52 `client' instances can be passed as arguments to select.select(),
53 and can be used in `with' statements.
54 """
7f97a47e 55 def __init__(self, sk, proto = None):
73c1ef8b
FT
56 """Create a client object connected to the specified
57 server. `sk' can either be a socket object, which is used as
58 it is, or a string specification very similar to the
57808152
FT
59 specification for L{pdm.srv.listen}, so see its documentation
60 for details. The differences are only that this function does
61 not take arguments specific to socket creation, like the mode
62 and group arguments for Unix sockets. If `proto' is given,
63 that subprotocol will negotiated with the server (by calling
64 the select() method).
73c1ef8b 65 """
7f97a47e 66 self.sk = resolve(sk)
11d50d09 67 self.buf = b""
7f97a47e 68 line = self.readline()
11d50d09 69 if line != b"+PDM1":
7f97a47e
FT
70 raise protoerr("Illegal protocol signature")
71 if proto is not None:
72 self.select(proto)
73
74 def close(self):
73c1ef8b 75 """Close this connection"""
7f97a47e
FT
76 self.sk.close()
77
9928d247 78 def fileno(self):
73c1ef8b 79 """Return the file descriptor of the underlying socket."""
9928d247
FT
80 return self.sk.fileno()
81
7f97a47e 82 def readline(self):
73c1ef8b 83 """Read a single NL-terminated line and return it."""
7f97a47e 84 while True:
11d50d09 85 p = self.buf.find(b"\n")
7f97a47e
FT
86 if p >= 0:
87 ret = self.buf[:p]
88 self.buf = self.buf[p + 1:]
89 return ret
90 ret = self.sk.recv(1024)
11d50d09 91 if ret == b"":
7f97a47e
FT
92 return None
93 self.buf += ret
94
95 def select(self, proto):
73c1ef8b 96 """Negotiate the given subprotocol with the server"""
11d50d09
FT
97 if isinstance(proto, str):
98 proto = proto.encode("ascii")
99 if b"\n" in proto:
7f97a47e 100 raise Exception("Illegal protocol specified: %r" % proto)
11d50d09 101 self.sk.send(proto + b"\n")
7f97a47e 102 rep = self.readline()
11d50d09 103 if len(rep) < 1 or rep[0] != b"+"[0]:
7f97a47e
FT
104 raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
105
106 def __enter__(self):
107 return self
108
109 def __exit__(self, *excinfo):
110 self.close()
111 return False
112
113class replclient(client):
73c1ef8b
FT
114 """REPL protocol client
115
57808152
FT
116 Implements the client side of the REPL protocol; see
117 L{pdm.srv.repl} for details on the protocol and its functionality.
73c1ef8b 118 """
7f97a47e 119 def __init__(self, sk):
73c1ef8b 120 """Create a connected client as documented in the `client' class."""
ed115f48 121 super().__init__(sk, "repl")
7f97a47e
FT
122
123 def run(self, code):
73c1ef8b 124 """Run a single block of Python code on the server. Returns
57808152
FT
125 the output of the command (as documented in L{pdm.srv.repl})
126 as a string.
73c1ef8b 127 """
7f97a47e
FT
128 while True:
129 ncode = code.replace("\n\n", "\n")
130 if ncode == code: break
131 code = ncode
132 while len(code) > 0 and code[-1] == "\n":
133 code = code[:-1]
11d50d09
FT
134 self.sk.send((code + "\n\n").encode("utf-8"))
135 buf = b""
7f97a47e
FT
136 while True:
137 ln = self.readline()
11d50d09
FT
138 if ln[0] == b" "[0]:
139 buf += ln[1:] + b"\n"
140 elif ln[0] == b"+"[0]:
141 return buf.decode("utf-8")
142 elif ln[0] == b"-"[0]:
143 raise protoerr("Error reply: %s" % ln[1:].decode("utf-8"))
7f97a47e
FT
144 else:
145 raise protoerr("Illegal reply: %s" % ln)
146
147class perfproxy(object):
148 def __init__(self, cl, id, proto):
149 self.cl = cl
150 self.id = id
151 self.proto = proto
152 self.subscribers = set()
153
154 def lookup(self, name):
155 self.cl.lock.acquire()
156 try:
157 id = self.cl.nextid
158 self.cl.nextid += 1
159 finally:
160 self.cl.lock.release()
161 (proto,) = self.cl.run("lookup", id, self.id, name)
162 proxy = perfproxy(self.cl, id, proto)
163 self.cl.proxies[id] = proxy
164 return proxy
165
166 def listdir(self):
167 return self.cl.run("ls", self.id)[0]
168
169 def readattr(self):
170 return self.cl.run("readattr", self.id)[0]
171
172 def attrinfo(self):
173 return self.cl.run("attrinfo", self.id)[0]
174
175 def invoke(self, method, *args, **kwargs):
176 return self.cl.run("invoke", self.id, method, args, kwargs)[0]
177
178 def subscribe(self, cb):
179 if cb in self.subscribers:
180 raise ValueError("Already subscribed")
181 if len(self.subscribers) == 0:
182 self.cl.run("subs", self.id)
183 self.subscribers.add(cb)
184
441160a2 185 def unsubscribe(self, cb):
7f97a47e
FT
186 if cb not in self.subscribers:
187 raise ValueError("Not subscribed")
188 self.subscribers.remove(cb)
189 if len(self.subscribers) == 0:
190 self.cl.run("unsubs", self.id)
191
192 def notify(self, ev):
193 for cb in self.subscribers:
194 try:
195 cb(ev)
196 except: pass
197
198 def close(self):
4fcf3c74
FT
199 if self.id is not None:
200 self.cl.run("unbind", self.id)
201 del self.cl.proxies[self.id]
202 self.id = None
203
204 def __del__(self):
205 self.close()
7f97a47e
FT
206
207 def __enter__(self):
208 return self
209
210 def __exit__(self, *excinfo):
211 self.close()
212 return False
213
214class perfclient(client):
73c1ef8b
FT
215 """PERF protocol client
216
57808152
FT
217 Implements the client side of the PERF protocol; see
218 L{pdm.srv.perf} for details on the protocol and its functionality.
73c1ef8b
FT
219
220 This client class implements functions for finding PERF objects on
221 the server, and returns, for each server-side object looked up, a
222 proxy object that mimics exactly the PERF interfaces that the
223 object implements. As the proxy objects reference live objects on
224 the server, they should be released when they are no longer used;
225 they implement a close() method for that purpose, and can also be
226 used in `with' statements.
227
57808152
FT
228 See L{pdm.srv.perf} for details on the various PERF interfaces
229 that the proxy objects might implement.
73c1ef8b 230 """
7f97a47e 231 def __init__(self, sk):
73c1ef8b 232 """Create a connected client as documented in the `client' class."""
ed115f48 233 super().__init__(sk, "perf")
7f97a47e
FT
234 self.nextid = 0
235 self.lock = threading.Lock()
236 self.proxies = {}
237 self.names = {}
238
239 def send(self, ob):
240 buf = pickle.dumps(ob)
241 buf = struct.pack(">l", len(buf)) + buf
242 self.sk.send(buf)
243
244 def recvb(self, num):
11d50d09 245 buf = b""
7f97a47e
FT
246 while len(buf) < num:
247 data = self.sk.recv(num - len(buf))
11d50d09 248 if data == b"":
7f97a47e
FT
249 raise EOFError()
250 buf += data
251 return buf
252
253 def recv(self):
254 return pickle.loads(self.recvb(struct.unpack(">l", self.recvb(4))[0]))
255
256 def event(self, id, ev):
257 proxy = self.proxies.get(id)
258 if proxy is None: return
259 proxy.notify(ev)
260
261 def dispatch(self, timeout = None):
73c1ef8b
FT
262 """Wait for an incoming notification from the server, and
263 dispatch it to the callback functions that have been
264 registered for it. If `timeout' is specified, wait no longer
265 than so many seconds; otherwise, wait forever. This client
266 object may also be used as argument to select.select().
267 """
7f97a47e
FT
268 rfd, wfd, efd = select.select([self.sk], [], [], timeout)
269 if self.sk in rfd:
270 msg = self.recv()
271 if msg[0] == "*":
272 self.event(msg[1], msg[2])
273 else:
274 raise ValueError("Unexpected non-event message: %r" % msg[0])
275
276 def recvreply(self):
277 while True:
278 reply = self.recv()
279 if reply[0] in ("+", "-"):
280 return reply
281 elif reply[0] == "*":
282 self.event(reply[1], reply[2])
283 else:
284 raise ValueError("Illegal reply header: %r" % reply[0])
285
286 def run(self, cmd, *args):
287 self.lock.acquire()
288 try:
289 self.send((cmd,) + args)
290 reply = self.recvreply()
291 if reply[0] == "+":
292 return reply[1:]
293 else:
294 raise reply[1]
295 finally:
296 self.lock.release()
297
298 def lookup(self, module, obnm):
73c1ef8b
FT
299 """Look up a single server-side object by the given name in
300 the given module. Will return a new proxy object for each
301 call when called multiple times for the same name.
302 """
7f97a47e
FT
303 self.lock.acquire()
304 try:
305 id = self.nextid
306 self.nextid += 1
307 finally:
308 self.lock.release()
309 (proto,) = self.run("bind", id, module, obnm)
310 proxy = perfproxy(self, id, proto)
311 self.proxies[id] = proxy
312 return proxy
313
314 def find(self, name):
73c1ef8b
FT
315 """Convenience function for looking up server-side objects
316 through PERF directories and for multiple uses. The object
317 name can be given as "MODULE.OBJECT", which will look up the
318 named OBJECT in the named MODULE, and can be followed by any
319 number of slash-separated names, which will assume that the
320 object to the left of the slash is a PERF directory, and will
321 return the object in that directory by the name to the right
322 of the slash. For instance, find("pdm.perf.sysres/cputime")
323 will return the built-in attribute for reading the CPU time
324 used by the server process.
ebd7a8ba
FT
325
326 The proxy objects returned by this function are cached and the
327 same object are returned the next time the same name is
328 requested, which means that they are kept live until the
329 client connection is closed.
73c1ef8b 330 """
7f97a47e
FT
331 ret = self.names.get(name)
332 if ret is None:
333 if "/" in name:
334 p = name.rindex("/")
335 ret = self.find(name[:p]).lookup(name[p + 1:])
336 else:
337 p = name.rindex(".")
338 ret = self.lookup(name[:p], name[p + 1:])
339 self.names[name] = ret
340 return ret