Merge branch 'master' into python3
[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 = b""
63         line = self.readline()
64         if line != b"+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(b"\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 == b"":
87                 return None
88             self.buf += ret
89
90     def select(self, proto):
91         """Negotiate the given subprotocol with the server"""
92         if isinstance(proto, str):
93             proto = proto.encode("ascii")
94         if b"\n" in proto:
95             raise Exception("Illegal protocol specified: %r" % proto)
96         self.sk.send(proto + b"\n")
97         rep = self.readline()
98         if len(rep) < 1 or rep[0] != b"+"[0]:
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
108 class replclient(client):
109     """REPL protocol client
110     
111     Implements the client side of the REPL protocol; see pdm.srv.repl
112     for details on the protocol and its functionality.
113     """
114     def __init__(self, sk):
115         """Create a connected client as documented in the `client' class."""
116         super().__init__(sk, "repl")
117
118     def run(self, code):
119         """Run a single block of Python code on the server. Returns
120         the output of the command (as documented in pdm.srv.repl) as a
121         string.
122         """
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]
129         self.sk.send((code + "\n\n").encode("utf-8"))
130         buf = b""
131         while True:
132             ln = self.readline()
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"))
139             else:
140                 raise protoerr("Illegal reply: %s" % ln)
141
142 class 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
180     def unsubscribe(self, cb):
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):
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()
201
202     def __enter__(self):
203         return self
204
205     def __exit__(self, *excinfo):
206         self.close()
207         return False
208
209 class perfclient(client):
210     """PERF protocol client
211     
212     Implements the client side of the PERF protocol; see pdm.srv.perf
213     for details on the protocol and its functionality.
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
223     See pdm.srv.perf for details on the various PERF interfaces that
224     the proxy objects might implement.
225     """
226     def __init__(self, sk):
227         """Create a connected client as documented in the `client' class."""
228         super().__init__(sk, "perf")
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):
240         buf = b""
241         while len(buf) < num:
242             data = self.sk.recv(num - len(buf))
243             if data == b"":
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):
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         """
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):
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         """
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):
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.
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.
325         """
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