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