Include perfclient in pdm.cli's exports.
[pdm.git] / pdm / cli.py
CommitLineData
7f97a47e
FT
1"""Management for daemon processes
2
3This module provides some client support for the daemon management
4provided in the pdm.srv module.
5"""
6
7import socket, pickle, struct, select, threading
8
ac4f5166 9__all__ = ["client", "replclient", "perfclient"]
7f97a47e
FT
10
11class protoerr(Exception):
12 pass
13
14def 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
37class 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
9928d247
FT
50 def fileno(self):
51 return self.sk.fileno()
52
7f97a47e
FT
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
80class 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
104class 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
441160a2 142 def unsubscribe(self, cb):
7f97a47e
FT
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):
4fcf3c74
FT
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()
7f97a47e
FT
163
164 def __enter__(self):
165 return self
166
167 def __exit__(self, *excinfo):
168 self.close()
169 return False
170
171class 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