buffers.py 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330
  1. ##############################################################################
  2. #
  3. # Copyright (c) 2001-2004 Zope Foundation and Contributors.
  4. # All Rights Reserved.
  5. #
  6. # This software is subject to the provisions of the Zope Public License,
  7. # Version 2.1 (ZPL). A copy of the ZPL should accompany this distribution.
  8. # THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
  9. # WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  10. # WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
  11. # FOR A PARTICULAR PURPOSE.
  12. #
  13. ##############################################################################
  14. """Buffers
  15. """
  16. from io import BytesIO
  17. # copy_bytes controls the size of temp. strings for shuffling data around.
  18. COPY_BYTES = 1 << 18 # 256K
  19. # The maximum number of bytes to buffer in a simple string.
  20. STRBUF_LIMIT = 8192
  21. class FileBasedBuffer:
  22. remain = 0
  23. def __init__(self, file, from_buffer=None):
  24. self.file = file
  25. if from_buffer is not None:
  26. from_file = from_buffer.getfile()
  27. read_pos = from_file.tell()
  28. from_file.seek(0)
  29. while True:
  30. data = from_file.read(COPY_BYTES)
  31. if not data:
  32. break
  33. file.write(data)
  34. self.remain = int(file.tell() - read_pos)
  35. from_file.seek(read_pos)
  36. file.seek(read_pos)
  37. def __len__(self):
  38. return self.remain
  39. def __nonzero__(self):
  40. return True
  41. __bool__ = __nonzero__ # py3
  42. def append(self, s):
  43. file = self.file
  44. read_pos = file.tell()
  45. file.seek(0, 2)
  46. file.write(s)
  47. file.seek(read_pos)
  48. self.remain = self.remain + len(s)
  49. def get(self, numbytes=-1, skip=False):
  50. file = self.file
  51. if not skip:
  52. read_pos = file.tell()
  53. if numbytes < 0:
  54. # Read all
  55. res = file.read()
  56. else:
  57. res = file.read(numbytes)
  58. if skip:
  59. self.remain -= len(res)
  60. else:
  61. file.seek(read_pos)
  62. return res
  63. def skip(self, numbytes, allow_prune=0):
  64. if self.remain < numbytes:
  65. raise ValueError(
  66. "Can't skip %d bytes in buffer of %d bytes" % (numbytes, self.remain)
  67. )
  68. self.file.seek(numbytes, 1)
  69. self.remain = self.remain - numbytes
  70. def newfile(self):
  71. raise NotImplementedError()
  72. def prune(self):
  73. file = self.file
  74. if self.remain == 0:
  75. read_pos = file.tell()
  76. file.seek(0, 2)
  77. sz = file.tell()
  78. file.seek(read_pos)
  79. if sz == 0:
  80. # Nothing to prune.
  81. return
  82. nf = self.newfile()
  83. while True:
  84. data = file.read(COPY_BYTES)
  85. if not data:
  86. break
  87. nf.write(data)
  88. self.file = nf
  89. def getfile(self):
  90. return self.file
  91. def close(self):
  92. if hasattr(self.file, "close"):
  93. self.file.close()
  94. self.remain = 0
  95. class TempfileBasedBuffer(FileBasedBuffer):
  96. def __init__(self, from_buffer=None):
  97. FileBasedBuffer.__init__(self, self.newfile(), from_buffer)
  98. def newfile(self):
  99. from tempfile import TemporaryFile
  100. return TemporaryFile("w+b")
  101. class BytesIOBasedBuffer(FileBasedBuffer):
  102. def __init__(self, from_buffer=None):
  103. if from_buffer is not None:
  104. FileBasedBuffer.__init__(self, BytesIO(), from_buffer)
  105. else:
  106. # Shortcut. :-)
  107. self.file = BytesIO()
  108. def newfile(self):
  109. return BytesIO()
  110. def _is_seekable(fp):
  111. if hasattr(fp, "seekable"):
  112. return fp.seekable()
  113. return hasattr(fp, "seek") and hasattr(fp, "tell")
  114. class ReadOnlyFileBasedBuffer(FileBasedBuffer):
  115. # used as wsgi.file_wrapper
  116. def __init__(self, file, block_size=32768):
  117. self.file = file
  118. self.block_size = block_size # for __iter__
  119. # This is for the benefit of anyone that is attempting to wrap this
  120. # wsgi.file_wrapper in a WSGI middleware and wants to seek, this is
  121. # useful for instance for support Range requests
  122. if _is_seekable(self.file):
  123. if hasattr(self.file, "seekable"):
  124. self.seekable = self.file.seekable
  125. self.seek = self.file.seek
  126. self.tell = self.file.tell
  127. def prepare(self, size=None):
  128. if _is_seekable(self.file):
  129. start_pos = self.file.tell()
  130. self.file.seek(0, 2)
  131. end_pos = self.file.tell()
  132. self.file.seek(start_pos)
  133. fsize = end_pos - start_pos
  134. if size is None:
  135. self.remain = fsize
  136. else:
  137. self.remain = min(fsize, size)
  138. return self.remain
  139. def get(self, numbytes=-1, skip=False):
  140. # never read more than self.remain (it can be user-specified)
  141. if numbytes == -1 or numbytes > self.remain:
  142. numbytes = self.remain
  143. file = self.file
  144. if not skip:
  145. read_pos = file.tell()
  146. res = file.read(numbytes)
  147. if skip:
  148. self.remain -= len(res)
  149. else:
  150. file.seek(read_pos)
  151. return res
  152. def __iter__(self): # called by task if self.filelike has no seek/tell
  153. return self
  154. def next(self):
  155. val = self.file.read(self.block_size)
  156. if not val:
  157. raise StopIteration
  158. return val
  159. __next__ = next # py3
  160. def append(self, s):
  161. raise NotImplementedError
  162. class OverflowableBuffer:
  163. """
  164. This buffer implementation has four stages:
  165. - No data
  166. - Bytes-based buffer
  167. - BytesIO-based buffer
  168. - Temporary file storage
  169. The first two stages are fastest for simple transfers.
  170. """
  171. overflowed = False
  172. buf = None
  173. strbuf = b"" # Bytes-based buffer.
  174. def __init__(self, overflow):
  175. # overflow is the maximum to be stored in a StringIO buffer.
  176. self.overflow = overflow
  177. def __len__(self):
  178. buf = self.buf
  179. if buf is not None:
  180. # use buf.__len__ rather than len(buf) FBO of not getting
  181. # OverflowError on Python 2
  182. return buf.__len__()
  183. else:
  184. return self.strbuf.__len__()
  185. def __nonzero__(self):
  186. # use self.__len__ rather than len(self) FBO of not getting
  187. # OverflowError on Python 2
  188. return self.__len__() > 0
  189. __bool__ = __nonzero__ # py3
  190. def _create_buffer(self):
  191. strbuf = self.strbuf
  192. if len(strbuf) >= self.overflow:
  193. self._set_large_buffer()
  194. else:
  195. self._set_small_buffer()
  196. buf = self.buf
  197. if strbuf:
  198. buf.append(self.strbuf)
  199. self.strbuf = b""
  200. return buf
  201. def _set_small_buffer(self):
  202. oldbuf = self.buf
  203. self.buf = BytesIOBasedBuffer(oldbuf)
  204. # Attempt to close the old buffer
  205. if hasattr(oldbuf, "close"):
  206. oldbuf.close()
  207. self.overflowed = False
  208. def _set_large_buffer(self):
  209. oldbuf = self.buf
  210. self.buf = TempfileBasedBuffer(oldbuf)
  211. # Attempt to close the old buffer
  212. if hasattr(oldbuf, "close"):
  213. oldbuf.close()
  214. self.overflowed = True
  215. def append(self, s):
  216. buf = self.buf
  217. if buf is None:
  218. strbuf = self.strbuf
  219. if len(strbuf) + len(s) < STRBUF_LIMIT:
  220. self.strbuf = strbuf + s
  221. return
  222. buf = self._create_buffer()
  223. buf.append(s)
  224. # use buf.__len__ rather than len(buf) FBO of not getting
  225. # OverflowError on Python 2
  226. sz = buf.__len__()
  227. if not self.overflowed:
  228. if sz >= self.overflow:
  229. self._set_large_buffer()
  230. def get(self, numbytes=-1, skip=False):
  231. buf = self.buf
  232. if buf is None:
  233. strbuf = self.strbuf
  234. if not skip:
  235. return strbuf
  236. buf = self._create_buffer()
  237. return buf.get(numbytes, skip)
  238. def skip(self, numbytes, allow_prune=False):
  239. buf = self.buf
  240. if buf is None:
  241. if allow_prune and numbytes == len(self.strbuf):
  242. # We could slice instead of converting to
  243. # a buffer, but that would eat up memory in
  244. # large transfers.
  245. self.strbuf = b""
  246. return
  247. buf = self._create_buffer()
  248. buf.skip(numbytes, allow_prune)
  249. def prune(self):
  250. """
  251. A potentially expensive operation that removes all data
  252. already retrieved from the buffer.
  253. """
  254. buf = self.buf
  255. if buf is None:
  256. self.strbuf = b""
  257. return
  258. buf.prune()
  259. if self.overflowed:
  260. # use buf.__len__ rather than len(buf) FBO of not getting
  261. # OverflowError on Python 2
  262. sz = buf.__len__()
  263. if sz < self.overflow:
  264. # Revert to a faster buffer.
  265. self._set_small_buffer()
  266. def getfile(self):
  267. buf = self.buf
  268. if buf is None:
  269. buf = self._create_buffer()
  270. return buf.getfile()
  271. def close(self):
  272. buf = self.buf
  273. if buf is not None:
  274. buf.close()