Properly split exception reports into lines.
[pdm.git] / pdm / cli.py
1 """Python Daemon Management -- Client functions
2
3 This module implements the client part of the PDM protocols. The
4 primary objects of interest are the replclient and perfclient classes,
5 which implement support for their respective protocols. See their
6 documentation for details.
7 """
8
9 import socket, pickle, struct, select, threading
10
11 __all__ = ["client", "replclient", "perfclient"]
12
13 class protoerr(Exception):
14     """Raised on protocol errors"""
15     pass
16
17 def resolve(spec):
18     if isinstance(spec, socket.socket):
19         return spec
20     sk = None
21     try:
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:
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)))
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
45 class client(object):
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     """
55     def __init__(self, sk, proto = None):
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
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).
65         """
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):
75         """Close this connection"""
76         self.sk.close()
77
78     def fileno(self):
79         """Return the file descriptor of the underlying socket."""
80         return self.sk.fileno()
81
82     def readline(self):
83         """Read a single NL-terminated line and return it."""
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):
96         """Negotiate the given subprotocol with the server"""
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
111 class replclient(client):
112     """REPL protocol client
113     
114     Implements the client side of the REPL protocol; see
115     L{pdm.srv.repl} for details on the protocol and its functionality.
116     """
117     def __init__(self, sk):
118         """Create a connected client as documented in the `client' class."""
119         super(replclient, self).__init__(sk, "repl")
120
121     def run(self, code):
122         """Run a single block of Python code on the server. Returns
123         the output of the command (as documented in L{pdm.srv.repl})
124         as a string.
125         """
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
145 class 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
183     def unsubscribe(self, cb):
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):
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()
204
205     def __enter__(self):
206         return self
207
208     def __exit__(self, *excinfo):
209         self.close()
210         return False
211
212 class perfclient(client):
213     """PERF protocol client
214     
215     Implements the client side of the PERF protocol; see
216     L{pdm.srv.perf} for details on the protocol and its functionality.
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
226     See L{pdm.srv.perf} for details on the various PERF interfaces
227     that the proxy objects might implement.
228     """
229     def __init__(self, sk):
230         """Create a connected client as documented in the `client' class."""
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):
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         """
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):
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         """
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):
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.
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.
328         """
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