Parse the socket mode explicitly in octal.
[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
9__all__ = ["client", "replclient"]
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
50 def readline(self):
51 while True:
52 p = self.buf.find("\n")
53 if p >= 0:
54 ret = self.buf[:p]
55 self.buf = self.buf[p + 1:]
56 return ret
57 ret = self.sk.recv(1024)
58 if ret == "":
59 return None
60 self.buf += ret
61
62 def select(self, proto):
63 if "\n" in proto:
64 raise Exception("Illegal protocol specified: %r" % proto)
65 self.sk.send(proto + "\n")
66 rep = self.readline()
67 if len(rep) < 1 or rep[0] != "+":
68 raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
69
70 def __enter__(self):
71 return self
72
73 def __exit__(self, *excinfo):
74 self.close()
75 return False
76
77class replclient(client):
78 def __init__(self, sk):
79 super(replclient, self).__init__(sk, "repl")
80
81 def run(self, code):
82 while True:
83 ncode = code.replace("\n\n", "\n")
84 if ncode == code: break
85 code = ncode
86 while len(code) > 0 and code[-1] == "\n":
87 code = code[:-1]
88 self.sk.send(code + "\n\n")
89 buf = ""
90 while True:
91 ln = self.readline()
92 if ln[0] == " ":
93 buf += ln[1:] + "\n"
94 elif ln[0] == "+":
95 return buf
96 elif ln[0] == "-":
97 raise protoerr("Error reply: %s" % ln[1:])
98 else:
99 raise protoerr("Illegal reply: %s" % ln)
100
101class perfproxy(object):
102 def __init__(self, cl, id, proto):
103 self.cl = cl
104 self.id = id
105 self.proto = proto
106 self.subscribers = set()
107
108 def lookup(self, name):
109 self.cl.lock.acquire()
110 try:
111 id = self.cl.nextid
112 self.cl.nextid += 1
113 finally:
114 self.cl.lock.release()
115 (proto,) = self.cl.run("lookup", id, self.id, name)
116 proxy = perfproxy(self.cl, id, proto)
117 self.cl.proxies[id] = proxy
118 return proxy
119
120 def listdir(self):
121 return self.cl.run("ls", self.id)[0]
122
123 def readattr(self):
124 return self.cl.run("readattr", self.id)[0]
125
126 def attrinfo(self):
127 return self.cl.run("attrinfo", self.id)[0]
128
129 def invoke(self, method, *args, **kwargs):
130 return self.cl.run("invoke", self.id, method, args, kwargs)[0]
131
132 def subscribe(self, cb):
133 if cb in self.subscribers:
134 raise ValueError("Already subscribed")
135 if len(self.subscribers) == 0:
136 self.cl.run("subs", self.id)
137 self.subscribers.add(cb)
138
441160a2 139 def unsubscribe(self, cb):
7f97a47e
FT
140 if cb not in self.subscribers:
141 raise ValueError("Not subscribed")
142 self.subscribers.remove(cb)
143 if len(self.subscribers) == 0:
144 self.cl.run("unsubs", self.id)
145
146 def notify(self, ev):
147 for cb in self.subscribers:
148 try:
149 cb(ev)
150 except: pass
151
152 def close(self):
153 self.cl.run("unbind", self.id)
154 del self.cl.proxies[self.id]
155
156 def __enter__(self):
157 return self
158
159 def __exit__(self, *excinfo):
160 self.close()
161 return False
162
163class perfclient(client):
164 def __init__(self, sk):
165 super(perfclient, self).__init__(sk, "perf")
166 self.nextid = 0
167 self.lock = threading.Lock()
168 self.proxies = {}
169 self.names = {}
170
171 def send(self, ob):
172 buf = pickle.dumps(ob)
173 buf = struct.pack(">l", len(buf)) + buf
174 self.sk.send(buf)
175
176 def recvb(self, num):
177 buf = ""
178 while len(buf) < num:
179 data = self.sk.recv(num - len(buf))
180 if data == "":
181 raise EOFError()
182 buf += data
183 return buf
184
185 def recv(self):
186 return pickle.loads(self.recvb(struct.unpack(">l", self.recvb(4))[0]))
187
188 def event(self, id, ev):
189 proxy = self.proxies.get(id)
190 if proxy is None: return
191 proxy.notify(ev)
192
193 def dispatch(self, timeout = None):
194 rfd, wfd, efd = select.select([self.sk], [], [], timeout)
195 if self.sk in rfd:
196 msg = self.recv()
197 if msg[0] == "*":
198 self.event(msg[1], msg[2])
199 else:
200 raise ValueError("Unexpected non-event message: %r" % msg[0])
201
202 def recvreply(self):
203 while True:
204 reply = self.recv()
205 if reply[0] in ("+", "-"):
206 return reply
207 elif reply[0] == "*":
208 self.event(reply[1], reply[2])
209 else:
210 raise ValueError("Illegal reply header: %r" % reply[0])
211
212 def run(self, cmd, *args):
213 self.lock.acquire()
214 try:
215 self.send((cmd,) + args)
216 reply = self.recvreply()
217 if reply[0] == "+":
218 return reply[1:]
219 else:
220 raise reply[1]
221 finally:
222 self.lock.release()
223
224 def lookup(self, module, obnm):
225 self.lock.acquire()
226 try:
227 id = self.nextid
228 self.nextid += 1
229 finally:
230 self.lock.release()
231 (proto,) = self.run("bind", id, module, obnm)
232 proxy = perfproxy(self, id, proto)
233 self.proxies[id] = proxy
234 return proxy
235
236 def find(self, name):
237 ret = self.names.get(name)
238 if ret is None:
239 if "/" in name:
240 p = name.rindex("/")
241 ret = self.find(name[:p]).lookup(name[p + 1:])
242 else:
243 p = name.rindex(".")
244 ret = self.lookup(name[:p], name[p + 1:])
245 self.names[name] = ret
246 return ret