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