Fixed copyrequest missing input.
[wrw.git] / wrw / req.py
1 import io
2
3 __all__ = ["request"]
4
5 class headdict(object):
6     def __init__(self):
7         self.dict = {}
8
9     def __getitem__(self, key):
10         return self.dict[key.lower()][1]
11
12     def __setitem__(self, key, val):
13         self.dict[key.lower()] = [key, val]
14
15     def __contains__(self, key):
16         return key.lower() in self.dict
17
18     def __delitem__(self, key):
19         del self.dict[key.lower()]
20
21     def __iter__(self):
22         return iter((list[0] for list in self.dict.values()))
23     
24     def get(self, key, default=""):
25         if key.lower() in self.dict:
26             return self.dict[key.lower()][1]
27         return default
28
29     def getlist(self, key):
30         return self.dict.setdefault(key.lower(), [key])[1:]
31
32     def add(self, key, val):
33         self.dict.setdefault(key.lower(), [key]).append(val)
34
35     def __repr__(self):
36         return repr(self.dict)
37
38     def __str__(self):
39         return str(self.dict)
40
41 def fixcase(str):
42     str = str.lower()
43     i = 0
44     b = True
45     while i < len(str):
46         if b:
47             str = str[:i] + str[i].upper() + str[i + 1:]
48         b = False
49         if str[i] == '-':
50             b = True
51         i += 1
52     return str
53
54 class limitreader(object):
55     def __init__(self, back, limit):
56         self.bk = back
57         self.limit = limit
58         self.rb = 0
59         self.buf = bytearray()
60
61     def close(self):
62         pass
63
64     def read(self, size=-1):
65         ra = self.limit - self.rb
66         if size >= 0:
67             ra = min(ra, size)
68         while len(self.buf) < ra:
69             ret = self.bk.read(ra - len(self.buf))
70             if ret == b"":
71                 raise IOError("Unexpected EOF")
72             self.buf.extend(ret)
73             self.rb += len(ret)
74         ret = bytes(self.buf[:ra])
75         self.buf = self.buf[ra:]
76         return ret
77
78     def readline(self, size=-1):
79         off = 0
80         while True:
81             p = self.buf.find(b'\n', off)
82             if p >= 0:
83                 ret = bytes(self.buf[:p + 1])
84                 self.buf = self.buf[p + 1:]
85                 return ret
86             off = len(self.buf)
87             if size >= 0 and len(self.buf) >= size:
88                 ret = bytes(self.buf[:size])
89                 self.buf = self.buf[size:]
90                 return ret
91             if self.rb == self.limit:
92                 ret = bytes(self.buf)
93                 self.buf = bytearray()
94                 return ret
95             ra = self.limit - self.rb
96             if size >= 0:
97                 ra = min(ra, size)
98             ra = min(ra, 1024)
99             ret = self.bk.read(ra)
100             if ret == b"":
101                 raise IOError("Unpexpected EOF")
102             self.buf.extend(ret)
103             self.rb += len(ret)
104
105     def readlines(self, hint=None):
106         return list(self)
107
108     def __iter__(rd):
109         class lineiter(object):
110             def __iter__(self):
111                 return self
112             def __next__(self):
113                 ret = rd.readline()
114                 if ret == b"":
115                     raise StopIteration()
116                 return ret
117         return lineiter()
118
119 class request(object):
120     def copy(self):
121         return copyrequest(self)
122
123     def shift(self, n):
124         new = self.copy()
125         new.uriname = self.uriname + self.pathinfo[:n]
126         new.pathinfo = self.pathinfo[n:]
127         return new
128
129 class origrequest(request):
130     def __init__(self, env):
131         self.env = env
132         self.method = env["REQUEST_METHOD"].upper()
133         self.uriname = env["SCRIPT_NAME"]
134         self.filename = env.get("SCRIPT_FILENAME")
135         self.uri = env["REQUEST_URI"]
136         self.pathinfo = env["PATH_INFO"]
137         self.query = env["QUERY_STRING"]
138         self.remoteaddr = env["REMOTE_ADDR"]
139         self.serverport = env["SERVER_PORT"]
140         self.servername = env["SERVER_NAME"]
141         self.https = "HTTPS" in env
142         self.ihead = headdict()
143         if "CONTENT_TYPE" in env:
144             self.ihead["Content-Type"] = env["CONTENT_TYPE"]
145             if "CONTENT_LENGTH" in env:
146                 clen = self.ihead["Content-Length"] = env["CONTENT_LENGTH"]
147                 if clen.isdigit():
148                     self.input = limitreader(env["wsgi.input"], int(clen))
149                 else:
150                     # XXX: What to do?
151                     self.input = io.BytesIO(b"")
152             else:
153                 # Assume input is chunked and read until ordinary EOF.
154                 self.input = env["wsgi.input"]
155         else:
156             self.input = None
157         self.ohead = headdict()
158         for k, v in env.items():
159             if k[:5] == "HTTP_":
160                 self.ihead.add(fixcase(k[5:].replace("_", "-")), v)
161         self.items = {}
162         self.statuscode = (200, "OK")
163         self.ohead["Content-Type"] = "text/html"
164         self.resources = set()
165         self.clean = set()
166         self.commitfuns = []
167
168     def status(self, code, msg):
169         self.statuscode = code, msg
170
171     def item(self, id):
172         if id in self.items:
173             return self.items[id]
174         self.items[id] = new = id(self)
175         if hasattr(new, "__enter__") and hasattr(new, "__exit__"):
176             self.withres(new)
177         return new
178
179     def withres(self, res):
180         if res not in self.resources:
181             done = False
182             res.__enter__()
183             try:
184                 self.resources.add(res)
185                 self.clean.add(res.__exit__)
186                 done = True
187             finally:
188                 if not done:
189                     res.__exit__(None, None, None)
190                     self.resources.discard(res)
191
192     def cleanup(self):
193         def clean1(list):
194             if len(list) > 0:
195                 try:
196                     list[0]()
197                 finally:
198                     clean1(list[1:])
199         clean1(list(self.clean))
200
201     def oncommit(self, fn):
202         if fn not in self.commitfuns:
203             self.commitfuns.append(fn)
204
205     def commit(self, startreq):
206         for fun in reversed(self.commitfuns):
207             fun(self)
208         hdrs = []
209         for nm in self.ohead:
210             for val in self.ohead.getlist(nm):
211                 hdrs.append((nm, val))
212         startreq("%s %s" % self.statuscode, hdrs)
213
214     def topreq(self):
215         return self
216
217 class copyrequest(request):
218     def __init__(self, p):
219         self.parent = p
220         self.top = p.topreq()
221         self.env = p.env
222         self.method = p.method
223         self.uriname = p.uriname
224         self.filename = p.filename
225         self.uri = p.uri
226         self.pathinfo = p.pathinfo
227         self.query = p.query
228         self.remoteaddr = p.remoteaddr
229         self.serverport = p.serverport
230         self.https = p.https
231         self.ihead = p.ihead
232         self.ohead = p.ohead
233         self.input = p.input
234
235     def status(self, code, msg):
236         return self.parent.status(code, msg)
237
238     def item(self, id):
239         return self.top.item(id)
240
241     def withres(self, res):
242         return self.top.withres(res)
243
244     def oncommit(self, fn):
245         return self.top.oncommit(fn)
246
247     def topreq(self):
248         return self.parent.topreq()