Properly split exception reports into lines.
[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:
26 import sshsock
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
FT
66 self.sk = resolve(sk)
67 self.buf = ""
68 line = self.readline()
69 if line != "+PDM1":
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
FT
84 while True:
85 p = self.buf.find("\n")
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)
91 if ret == "":
92 return None
93 self.buf += ret
94
95 def select(self, proto):
73c1ef8b 96 """Negotiate the given subprotocol with the server"""
7f97a47e
FT
97 if "\n" in proto:
98 raise Exception("Illegal protocol specified: %r" % proto)
99 self.sk.send(proto + "\n")
100 rep = self.readline()
101 if len(rep) < 1 or rep[0] != "+":
102 raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
103
104 def __enter__(self):
105 return self
106
107 def __exit__(self, *excinfo):
108 self.close()
109 return False
110
111class replclient(client):
73c1ef8b
FT
112 """REPL protocol client
113
57808152
FT
114 Implements the client side of the REPL protocol; see
115 L{pdm.srv.repl} for details on the protocol and its functionality.
73c1ef8b 116 """
7f97a47e 117 def __init__(self, sk):
73c1ef8b 118 """Create a connected client as documented in the `client' class."""
7f97a47e
FT
119 super(replclient, self).__init__(sk, "repl")
120
121 def run(self, code):
73c1ef8b 122 """Run a single block of Python code on the server. Returns
57808152
FT
123 the output of the command (as documented in L{pdm.srv.repl})
124 as a string.
73c1ef8b 125 """
7f97a47e
FT
126 while True:
127 ncode = code.replace("\n\n", "\n")
128 if ncode == code: break
129 code = ncode
130 while len(code) > 0 and code[-1] == "\n":
131 code = code[:-1]
132 self.sk.send(code + "\n\n")
133 buf = ""
134 while True:
135 ln = self.readline()
136 if ln[0] == " ":
137 buf += ln[1:] + "\n"
138 elif ln[0] == "+":
139 return buf
140 elif ln[0] == "-":
141 raise protoerr("Error reply: %s" % ln[1:])
142 else:
143 raise protoerr("Illegal reply: %s" % ln)
144
145class perfproxy(object):
146 def __init__(self, cl, id, proto):
147 self.cl = cl
148 self.id = id
149 self.proto = proto
150 self.subscribers = set()
151
152 def lookup(self, name):
153 self.cl.lock.acquire()
154 try:
155 id = self.cl.nextid
156 self.cl.nextid += 1
157 finally:
158 self.cl.lock.release()
159 (proto,) = self.cl.run("lookup", id, self.id, name)
160 proxy = perfproxy(self.cl, id, proto)
161 self.cl.proxies[id] = proxy
162 return proxy
163
164 def listdir(self):
165 return self.cl.run("ls", self.id)[0]
166
167 def readattr(self):
168 return self.cl.run("readattr", self.id)[0]
169
170 def attrinfo(self):
171 return self.cl.run("attrinfo", self.id)[0]
172
173 def invoke(self, method, *args, **kwargs):
174 return self.cl.run("invoke", self.id, method, args, kwargs)[0]
175
176 def subscribe(self, cb):
177 if cb in self.subscribers:
178 raise ValueError("Already subscribed")
179 if len(self.subscribers) == 0:
180 self.cl.run("subs", self.id)
181 self.subscribers.add(cb)
182
441160a2 183 def unsubscribe(self, cb):
7f97a47e
FT
184 if cb not in self.subscribers:
185 raise ValueError("Not subscribed")
186 self.subscribers.remove(cb)
187 if len(self.subscribers) == 0:
188 self.cl.run("unsubs", self.id)
189
190 def notify(self, ev):
191 for cb in self.subscribers:
192 try:
193 cb(ev)
194 except: pass
195
196 def close(self):
4fcf3c74
FT
197 if self.id is not None:
198 self.cl.run("unbind", self.id)
199 del self.cl.proxies[self.id]
200 self.id = None
201
202 def __del__(self):
203 self.close()
7f97a47e
FT
204
205 def __enter__(self):
206 return self
207
208 def __exit__(self, *excinfo):
209 self.close()
210 return False
211
212class perfclient(client):
73c1ef8b
FT
213 """PERF protocol client
214
57808152
FT
215 Implements the client side of the PERF protocol; see
216 L{pdm.srv.perf} for details on the protocol and its functionality.
73c1ef8b
FT
217
218 This client class implements functions for finding PERF objects on
219 the server, and returns, for each server-side object looked up, a
220 proxy object that mimics exactly the PERF interfaces that the
221 object implements. As the proxy objects reference live objects on
222 the server, they should be released when they are no longer used;
223 they implement a close() method for that purpose, and can also be
224 used in `with' statements.
225
57808152
FT
226 See L{pdm.srv.perf} for details on the various PERF interfaces
227 that the proxy objects might implement.
73c1ef8b 228 """
7f97a47e 229 def __init__(self, sk):
73c1ef8b 230 """Create a connected client as documented in the `client' class."""
7f97a47e
FT
231 super(perfclient, self).__init__(sk, "perf")
232 self.nextid = 0
233 self.lock = threading.Lock()
234 self.proxies = {}
235 self.names = {}
236
237 def send(self, ob):
238 buf = pickle.dumps(ob)
239 buf = struct.pack(">l", len(buf)) + buf
240 self.sk.send(buf)
241
242 def recvb(self, num):
243 buf = ""
244 while len(buf) < num:
245 data = self.sk.recv(num - len(buf))
246 if data == "":
247 raise EOFError()
248 buf += data
249 return buf
250
251 def recv(self):
252 return pickle.loads(self.recvb(struct.unpack(">l", self.recvb(4))[0]))
253
254 def event(self, id, ev):
255 proxy = self.proxies.get(id)
256 if proxy is None: return
257 proxy.notify(ev)
258
259 def dispatch(self, timeout = None):
73c1ef8b
FT
260 """Wait for an incoming notification from the server, and
261 dispatch it to the callback functions that have been
262 registered for it. If `timeout' is specified, wait no longer
263 than so many seconds; otherwise, wait forever. This client
264 object may also be used as argument to select.select().
265 """
7f97a47e
FT
266 rfd, wfd, efd = select.select([self.sk], [], [], timeout)
267 if self.sk in rfd:
268 msg = self.recv()
269 if msg[0] == "*":
270 self.event(msg[1], msg[2])
271 else:
272 raise ValueError("Unexpected non-event message: %r" % msg[0])
273
274 def recvreply(self):
275 while True:
276 reply = self.recv()
277 if reply[0] in ("+", "-"):
278 return reply
279 elif reply[0] == "*":
280 self.event(reply[1], reply[2])
281 else:
282 raise ValueError("Illegal reply header: %r" % reply[0])
283
284 def run(self, cmd, *args):
285 self.lock.acquire()
286 try:
287 self.send((cmd,) + args)
288 reply = self.recvreply()
289 if reply[0] == "+":
290 return reply[1:]
291 else:
292 raise reply[1]
293 finally:
294 self.lock.release()
295
296 def lookup(self, module, obnm):
73c1ef8b
FT
297 """Look up a single server-side object by the given name in
298 the given module. Will return a new proxy object for each
299 call when called multiple times for the same name.
300 """
7f97a47e
FT
301 self.lock.acquire()
302 try:
303 id = self.nextid
304 self.nextid += 1
305 finally:
306 self.lock.release()
307 (proto,) = self.run("bind", id, module, obnm)
308 proxy = perfproxy(self, id, proto)
309 self.proxies[id] = proxy
310 return proxy
311
312 def find(self, name):
73c1ef8b
FT
313 """Convenience function for looking up server-side objects
314 through PERF directories and for multiple uses. The object
315 name can be given as "MODULE.OBJECT", which will look up the
316 named OBJECT in the named MODULE, and can be followed by any
317 number of slash-separated names, which will assume that the
318 object to the left of the slash is a PERF directory, and will
319 return the object in that directory by the name to the right
320 of the slash. For instance, find("pdm.perf.sysres/cputime")
321 will return the built-in attribute for reading the CPU time
322 used by the server process.
ebd7a8ba
FT
323
324 The proxy objects returned by this function are cached and the
325 same object are returned the next time the same name is
326 requested, which means that they are kept live until the
327 client connection is closed.
73c1ef8b 328 """
7f97a47e
FT
329 ret = self.names.get(name)
330 if ret is None:
331 if "/" in name:
332 p = name.rindex("/")
333 ret = self.find(name[:p]).lookup(name[p + 1:])
334 else:
335 p = name.rindex(".")
336 ret = self.lookup(name[:p], name[p + 1:])
337 self.names[name] = ret
338 return ret