d16cbbdc92bc316f370f819df789f6f18c369ab4
[pdm.git] / pdm / cli.py
1 """Management for daemon processes
2
3 This module provides some client support for the daemon management
4 provided in the pdm.srv module.
5 """
6
7 import socket, pickle, struct, select, threading
8
9 __all__ = ["client", "replclient"]
10
11 class protoerr(Exception):
12     pass
13
14 def resolve(spec):
15     if isinstance(spec, socket.socket):
16         return spec
17     sk = None
18     try:
19         if "/" in spec:
20             sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
21             sk.connect(spec)
22         elif spec.isdigit():
23             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
24             sk.connect(("localhost", int(spec)))
25         elif ":" in spec:
26             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
27             p = spec.rindex(":")
28             sk.connect((spec[:p], int(spec[p + 1:])))
29         else:
30             raise Exception("Unknown target specification %r" % spec)
31         rv = sk
32         sk = None
33     finally:
34         if sk is not None: sk.close()
35     return rv
36
37 class client(object):
38     def __init__(self, sk, proto = None):
39         self.sk = resolve(sk)
40         self.buf = ""
41         line = self.readline()
42         if line != "+PDM1":
43             raise protoerr("Illegal protocol signature")
44         if proto is not None:
45             self.select(proto)
46
47     def close(self):
48         self.sk.close()
49
50     def fileno(self):
51         return self.sk.fileno()
52
53     def readline(self):
54         while True:
55             p = self.buf.find("\n")
56             if p >= 0:
57                 ret = self.buf[:p]
58                 self.buf = self.buf[p + 1:]
59                 return ret
60             ret = self.sk.recv(1024)
61             if ret == "":
62                 return None
63             self.buf += ret
64
65     def select(self, proto):
66         if "\n" in proto:
67             raise Exception("Illegal protocol specified: %r" % proto)
68         self.sk.send(proto + "\n")
69         rep = self.readline()
70         if len(rep) < 1 or rep[0] != "+":
71             raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
72
73     def __enter__(self):
74         return self
75
76     def __exit__(self, *excinfo):
77         self.close()
78         return False
79
80 class replclient(client):
81     def __init__(self, sk):
82         super(replclient, self).__init__(sk, "repl")
83
84     def run(self, code):
85         while True:
86             ncode = code.replace("\n\n", "\n")
87             if ncode == code: break
88             code = ncode
89         while len(code) > 0 and code[-1] == "\n":
90             code = code[:-1]
91         self.sk.send(code + "\n\n")
92         buf = ""
93         while True:
94             ln = self.readline()
95             if ln[0] == " ":
96                 buf += ln[1:] + "\n"
97             elif ln[0] == "+":
98                 return buf
99             elif ln[0] == "-":
100                 raise protoerr("Error reply: %s" % ln[1:])
101             else:
102                 raise protoerr("Illegal reply: %s" % ln)
103
104 class perfproxy(object):
105     def __init__(self, cl, id, proto):
106         self.cl = cl
107         self.id = id
108         self.proto = proto
109         self.subscribers = set()
110
111     def lookup(self, name):
112         self.cl.lock.acquire()
113         try:
114             id = self.cl.nextid
115             self.cl.nextid += 1
116         finally:
117             self.cl.lock.release()
118         (proto,) = self.cl.run("lookup", id, self.id, name)
119         proxy = perfproxy(self.cl, id, proto)
120         self.cl.proxies[id] = proxy
121         return proxy
122
123     def listdir(self):
124         return self.cl.run("ls", self.id)[0]
125
126     def readattr(self):
127         return self.cl.run("readattr", self.id)[0]
128
129     def attrinfo(self):
130         return self.cl.run("attrinfo", self.id)[0]
131
132     def invoke(self, method, *args, **kwargs):
133         return self.cl.run("invoke", self.id, method, args, kwargs)[0]
134
135     def subscribe(self, cb):
136         if cb in self.subscribers:
137             raise ValueError("Already subscribed")
138         if len(self.subscribers) == 0:
139             self.cl.run("subs", self.id)
140         self.subscribers.add(cb)
141
142     def unsubscribe(self, cb):
143         if cb not in self.subscribers:
144             raise ValueError("Not subscribed")
145         self.subscribers.remove(cb)
146         if len(self.subscribers) == 0:
147             self.cl.run("unsubs", self.id)
148
149     def notify(self, ev):
150         for cb in self.subscribers:
151             try:
152                 cb(ev)
153             except: pass
154
155     def close(self):
156         if self.id is not None:
157             self.cl.run("unbind", self.id)
158             del self.cl.proxies[self.id]
159             self.id = None
160
161     def __del__(self):
162         self.close()
163
164     def __enter__(self):
165         return self
166
167     def __exit__(self, *excinfo):
168         self.close()
169         return False
170
171 class perfclient(client):
172     def __init__(self, sk):
173         super(perfclient, self).__init__(sk, "perf")
174         self.nextid = 0
175         self.lock = threading.Lock()
176         self.proxies = {}
177         self.names = {}
178
179     def send(self, ob):
180         buf = pickle.dumps(ob)
181         buf = struct.pack(">l", len(buf)) + buf
182         self.sk.send(buf)
183
184     def recvb(self, num):
185         buf = ""
186         while len(buf) < num:
187             data = self.sk.recv(num - len(buf))
188             if data == "":
189                 raise EOFError()
190             buf += data
191         return buf
192
193     def recv(self):
194         return pickle.loads(self.recvb(struct.unpack(">l", self.recvb(4))[0]))
195
196     def event(self, id, ev):
197         proxy = self.proxies.get(id)
198         if proxy is None: return
199         proxy.notify(ev)
200
201     def dispatch(self, timeout = None):
202         rfd, wfd, efd = select.select([self.sk], [], [], timeout)
203         if self.sk in rfd:
204             msg = self.recv()
205             if msg[0] == "*":
206                 self.event(msg[1], msg[2])
207             else:
208                 raise ValueError("Unexpected non-event message: %r" % msg[0])
209
210     def recvreply(self):
211         while True:
212             reply = self.recv()
213             if reply[0] in ("+", "-"):
214                 return reply
215             elif reply[0] == "*":
216                 self.event(reply[1], reply[2])
217             else:
218                 raise ValueError("Illegal reply header: %r" % reply[0])
219
220     def run(self, cmd, *args):
221         self.lock.acquire()
222         try:
223             self.send((cmd,) + args)
224             reply = self.recvreply()
225             if reply[0] == "+":
226                 return reply[1:]
227             else:
228                 raise reply[1]
229         finally:
230             self.lock.release()
231
232     def lookup(self, module, obnm):
233         self.lock.acquire()
234         try:
235             id = self.nextid
236             self.nextid += 1
237         finally:
238             self.lock.release()
239         (proto,) = self.run("bind", id, module, obnm)
240         proxy = perfproxy(self, id, proto)
241         self.proxies[id] = proxy
242         return proxy
243
244     def find(self, name):
245         ret = self.names.get(name)
246         if ret is None:
247             if "/" in name:
248                 p = name.rindex("/")
249                 ret = self.find(name[:p]).lookup(name[p + 1:])
250             else:
251                 p = name.rindex(".")
252                 ret = self.lookup(name[:p], name[p + 1:])
253             self.names[name] = ret
254         return ret