encoder.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655
  1. # -*- coding: utf-8 -*-
  2. """
  3. requests_toolbelt.multipart.encoder
  4. ===================================
  5. This holds all of the implementation details of the MultipartEncoder
  6. """
  7. import contextlib
  8. import io
  9. import os
  10. from uuid import uuid4
  11. import requests
  12. from .._compat import fields
  13. class FileNotSupportedError(Exception):
  14. """File not supported error."""
  15. class MultipartEncoder(object):
  16. """
  17. The ``MultipartEncoder`` object is a generic interface to the engine that
  18. will create a ``multipart/form-data`` body for you.
  19. The basic usage is:
  20. .. code-block:: python
  21. import requests
  22. from requests_toolbelt import MultipartEncoder
  23. encoder = MultipartEncoder({'field': 'value',
  24. 'other_field', 'other_value'})
  25. r = requests.post('https://httpbin.org/post', data=encoder,
  26. headers={'Content-Type': encoder.content_type})
  27. If you do not need to take advantage of streaming the post body, you can
  28. also do:
  29. .. code-block:: python
  30. r = requests.post('https://httpbin.org/post',
  31. data=encoder.to_string(),
  32. headers={'Content-Type': encoder.content_type})
  33. If you want the encoder to use a specific order, you can use an
  34. OrderedDict or more simply, a list of tuples:
  35. .. code-block:: python
  36. encoder = MultipartEncoder([('field', 'value'),
  37. ('other_field', 'other_value')])
  38. .. versionchanged:: 0.4.0
  39. You can also provide tuples as part values as you would provide them to
  40. requests' ``files`` parameter.
  41. .. code-block:: python
  42. encoder = MultipartEncoder({
  43. 'field': ('file_name', b'{"a": "b"}', 'application/json',
  44. {'X-My-Header': 'my-value'})
  45. ])
  46. .. warning::
  47. This object will end up directly in :mod:`httplib`. Currently,
  48. :mod:`httplib` has a hard-coded read size of **8192 bytes**. This
  49. means that it will loop until the file has been read and your upload
  50. could take a while. This is **not** a bug in requests. A feature is
  51. being considered for this object to allow you, the user, to specify
  52. what size should be returned on a read. If you have opinions on this,
  53. please weigh in on `this issue`_.
  54. .. _this issue:
  55. https://github.com/requests/toolbelt/issues/75
  56. """
  57. def __init__(self, fields, boundary=None, encoding='utf-8'):
  58. #: Boundary value either passed in by the user or created
  59. self.boundary_value = boundary or uuid4().hex
  60. # Computed boundary
  61. self.boundary = '--{0}'.format(self.boundary_value)
  62. #: Encoding of the data being passed in
  63. self.encoding = encoding
  64. # Pre-encoded boundary
  65. self._encoded_boundary = b''.join([
  66. encode_with(self.boundary, self.encoding),
  67. encode_with('\r\n', self.encoding)
  68. ])
  69. #: Fields provided by the user
  70. self.fields = fields
  71. #: Whether or not the encoder is finished
  72. self.finished = False
  73. #: Pre-computed parts of the upload
  74. self.parts = []
  75. # Pre-computed parts iterator
  76. self._iter_parts = iter([])
  77. # The part we're currently working with
  78. self._current_part = None
  79. # Cached computation of the body's length
  80. self._len = None
  81. # Our buffer
  82. self._buffer = CustomBytesIO(encoding=encoding)
  83. # Pre-compute each part's headers
  84. self._prepare_parts()
  85. # Load boundary into buffer
  86. self._write_boundary()
  87. @property
  88. def len(self):
  89. """Length of the multipart/form-data body.
  90. requests will first attempt to get the length of the body by calling
  91. ``len(body)`` and then by checking for the ``len`` attribute.
  92. On 32-bit systems, the ``__len__`` method cannot return anything
  93. larger than an integer (in C) can hold. If the total size of the body
  94. is even slightly larger than 4GB users will see an OverflowError. This
  95. manifested itself in `bug #80`_.
  96. As such, we now calculate the length lazily as a property.
  97. .. _bug #80:
  98. https://github.com/requests/toolbelt/issues/80
  99. """
  100. # If _len isn't already calculated, calculate, return, and set it
  101. return self._len or self._calculate_length()
  102. def __repr__(self):
  103. return '<MultipartEncoder: {0!r}>'.format(self.fields)
  104. def _calculate_length(self):
  105. """
  106. This uses the parts to calculate the length of the body.
  107. This returns the calculated length so __len__ can be lazy.
  108. """
  109. boundary_len = len(self.boundary) # Length of --{boundary}
  110. # boundary length + header length + body length + len('\r\n') * 2
  111. self._len = sum(
  112. (boundary_len + total_len(p) + 4) for p in self.parts
  113. ) + boundary_len + 4
  114. return self._len
  115. def _calculate_load_amount(self, read_size):
  116. """This calculates how many bytes need to be added to the buffer.
  117. When a consumer read's ``x`` from the buffer, there are two cases to
  118. satisfy:
  119. 1. Enough data in the buffer to return the requested amount
  120. 2. Not enough data
  121. This function uses the amount of unread bytes in the buffer and
  122. determines how much the Encoder has to load before it can return the
  123. requested amount of bytes.
  124. :param int read_size: the number of bytes the consumer requests
  125. :returns: int -- the number of bytes that must be loaded into the
  126. buffer before the read can be satisfied. This will be strictly
  127. non-negative
  128. """
  129. amount = read_size - total_len(self._buffer)
  130. return amount if amount > 0 else 0
  131. def _load(self, amount):
  132. """Load ``amount`` number of bytes into the buffer."""
  133. self._buffer.smart_truncate()
  134. part = self._current_part or self._next_part()
  135. while amount == -1 or amount > 0:
  136. written = 0
  137. if part and not part.bytes_left_to_write():
  138. written += self._write(b'\r\n')
  139. written += self._write_boundary()
  140. part = self._next_part()
  141. if not part:
  142. written += self._write_closing_boundary()
  143. self.finished = True
  144. break
  145. written += part.write_to(self._buffer, amount)
  146. if amount != -1:
  147. amount -= written
  148. def _next_part(self):
  149. try:
  150. p = self._current_part = next(self._iter_parts)
  151. except StopIteration:
  152. p = None
  153. return p
  154. def _iter_fields(self):
  155. _fields = self.fields
  156. if hasattr(self.fields, 'items'):
  157. _fields = list(self.fields.items())
  158. for k, v in _fields:
  159. file_name = None
  160. file_type = None
  161. file_headers = None
  162. if isinstance(v, (list, tuple)):
  163. if len(v) == 2:
  164. file_name, file_pointer = v
  165. elif len(v) == 3:
  166. file_name, file_pointer, file_type = v
  167. else:
  168. file_name, file_pointer, file_type, file_headers = v
  169. else:
  170. file_pointer = v
  171. field = fields.RequestField(name=k, data=file_pointer,
  172. filename=file_name,
  173. headers=file_headers)
  174. field.make_multipart(content_type=file_type)
  175. yield field
  176. def _prepare_parts(self):
  177. """This uses the fields provided by the user and creates Part objects.
  178. It populates the `parts` attribute and uses that to create a
  179. generator for iteration.
  180. """
  181. enc = self.encoding
  182. self.parts = [Part.from_field(f, enc) for f in self._iter_fields()]
  183. self._iter_parts = iter(self.parts)
  184. def _write(self, bytes_to_write):
  185. """Write the bytes to the end of the buffer.
  186. :param bytes bytes_to_write: byte-string (or bytearray) to append to
  187. the buffer
  188. :returns: int -- the number of bytes written
  189. """
  190. return self._buffer.append(bytes_to_write)
  191. def _write_boundary(self):
  192. """Write the boundary to the end of the buffer."""
  193. return self._write(self._encoded_boundary)
  194. def _write_closing_boundary(self):
  195. """Write the bytes necessary to finish a multipart/form-data body."""
  196. with reset(self._buffer):
  197. self._buffer.seek(-2, 2)
  198. self._buffer.write(b'--\r\n')
  199. return 2
  200. def _write_headers(self, headers):
  201. """Write the current part's headers to the buffer."""
  202. return self._write(encode_with(headers, self.encoding))
  203. @property
  204. def content_type(self):
  205. return str(
  206. 'multipart/form-data; boundary={0}'.format(self.boundary_value)
  207. )
  208. def to_string(self):
  209. """Return the entirety of the data in the encoder.
  210. .. note::
  211. This simply reads all of the data it can. If you have started
  212. streaming or reading data from the encoder, this method will only
  213. return whatever data is left in the encoder.
  214. .. note::
  215. This method affects the internal state of the encoder. Calling
  216. this method will exhaust the encoder.
  217. :returns: the multipart message
  218. :rtype: bytes
  219. """
  220. return self.read()
  221. def read(self, size=-1):
  222. """Read data from the streaming encoder.
  223. :param int size: (optional), If provided, ``read`` will return exactly
  224. that many bytes. If it is not provided, it will return the
  225. remaining bytes.
  226. :returns: bytes
  227. """
  228. if self.finished:
  229. return self._buffer.read(size)
  230. bytes_to_load = size
  231. if bytes_to_load != -1 and bytes_to_load is not None:
  232. bytes_to_load = self._calculate_load_amount(int(size))
  233. self._load(bytes_to_load)
  234. return self._buffer.read(size)
  235. def IDENTITY(monitor):
  236. return monitor
  237. class MultipartEncoderMonitor(object):
  238. """
  239. An object used to monitor the progress of a :class:`MultipartEncoder`.
  240. The :class:`MultipartEncoder` should only be responsible for preparing and
  241. streaming the data. For anyone who wishes to monitor it, they shouldn't be
  242. using that instance to manage that as well. Using this class, they can
  243. monitor an encoder and register a callback. The callback receives the
  244. instance of the monitor.
  245. To use this monitor, you construct your :class:`MultipartEncoder` as you
  246. normally would.
  247. .. code-block:: python
  248. from requests_toolbelt import (MultipartEncoder,
  249. MultipartEncoderMonitor)
  250. import requests
  251. def callback(monitor):
  252. # Do something with this information
  253. pass
  254. m = MultipartEncoder(fields={'field0': 'value0'})
  255. monitor = MultipartEncoderMonitor(m, callback)
  256. headers = {'Content-Type': monitor.content_type}
  257. r = requests.post('https://httpbin.org/post', data=monitor,
  258. headers=headers)
  259. Alternatively, if your use case is very simple, you can use the following
  260. pattern.
  261. .. code-block:: python
  262. from requests_toolbelt import MultipartEncoderMonitor
  263. import requests
  264. def callback(monitor):
  265. # Do something with this information
  266. pass
  267. monitor = MultipartEncoderMonitor.from_fields(
  268. fields={'field0': 'value0'}, callback
  269. )
  270. headers = {'Content-Type': montior.content_type}
  271. r = requests.post('https://httpbin.org/post', data=monitor,
  272. headers=headers)
  273. """
  274. def __init__(self, encoder, callback=None):
  275. #: Instance of the :class:`MultipartEncoder` being monitored
  276. self.encoder = encoder
  277. #: Optionally function to call after a read
  278. self.callback = callback or IDENTITY
  279. #: Number of bytes already read from the :class:`MultipartEncoder`
  280. #: instance
  281. self.bytes_read = 0
  282. #: Avoid the same problem in bug #80
  283. self.len = self.encoder.len
  284. @classmethod
  285. def from_fields(cls, fields, boundary=None, encoding='utf-8',
  286. callback=None):
  287. encoder = MultipartEncoder(fields, boundary, encoding)
  288. return cls(encoder, callback)
  289. @property
  290. def content_type(self):
  291. return self.encoder.content_type
  292. def to_string(self):
  293. return self.read()
  294. def read(self, size=-1):
  295. string = self.encoder.read(size)
  296. self.bytes_read += len(string)
  297. self.callback(self)
  298. return string
  299. def encode_with(string, encoding):
  300. """Encoding ``string`` with ``encoding`` if necessary.
  301. :param str string: If string is a bytes object, it will not encode it.
  302. Otherwise, this function will encode it with the provided encoding.
  303. :param str encoding: The encoding with which to encode string.
  304. :returns: encoded bytes object
  305. """
  306. if not (string is None or isinstance(string, bytes)):
  307. return string.encode(encoding)
  308. return string
  309. def readable_data(data, encoding):
  310. """Coerce the data to an object with a ``read`` method."""
  311. if hasattr(data, 'read'):
  312. return data
  313. return CustomBytesIO(data, encoding)
  314. def total_len(o):
  315. if hasattr(o, '__len__'):
  316. return len(o)
  317. if hasattr(o, 'len'):
  318. return o.len
  319. if hasattr(o, 'fileno'):
  320. try:
  321. fileno = o.fileno()
  322. except io.UnsupportedOperation:
  323. pass
  324. else:
  325. return os.fstat(fileno).st_size
  326. if hasattr(o, 'getvalue'):
  327. # e.g. BytesIO, cStringIO.StringIO
  328. return len(o.getvalue())
  329. @contextlib.contextmanager
  330. def reset(buffer):
  331. """Keep track of the buffer's current position and write to the end.
  332. This is a context manager meant to be used when adding data to the buffer.
  333. It eliminates the need for every function to be concerned with the
  334. position of the cursor in the buffer.
  335. """
  336. original_position = buffer.tell()
  337. buffer.seek(0, 2)
  338. yield
  339. buffer.seek(original_position, 0)
  340. def coerce_data(data, encoding):
  341. """Ensure that every object's __len__ behaves uniformly."""
  342. if not isinstance(data, CustomBytesIO):
  343. if hasattr(data, 'getvalue'):
  344. return CustomBytesIO(data.getvalue(), encoding)
  345. if hasattr(data, 'fileno'):
  346. return FileWrapper(data)
  347. if not hasattr(data, 'read'):
  348. return CustomBytesIO(data, encoding)
  349. return data
  350. def to_list(fields):
  351. if hasattr(fields, 'items'):
  352. return list(fields.items())
  353. return list(fields)
  354. class Part(object):
  355. def __init__(self, headers, body):
  356. self.headers = headers
  357. self.body = body
  358. self.headers_unread = True
  359. self.len = len(self.headers) + total_len(self.body)
  360. @classmethod
  361. def from_field(cls, field, encoding):
  362. """Create a part from a Request Field generated by urllib3."""
  363. headers = encode_with(field.render_headers(), encoding)
  364. body = coerce_data(field.data, encoding)
  365. return cls(headers, body)
  366. def bytes_left_to_write(self):
  367. """Determine if there are bytes left to write.
  368. :returns: bool -- ``True`` if there are bytes left to write, otherwise
  369. ``False``
  370. """
  371. to_read = 0
  372. if self.headers_unread:
  373. to_read += len(self.headers)
  374. return (to_read + total_len(self.body)) > 0
  375. def write_to(self, buffer, size):
  376. """Write the requested amount of bytes to the buffer provided.
  377. The number of bytes written may exceed size on the first read since we
  378. load the headers ambitiously.
  379. :param CustomBytesIO buffer: buffer we want to write bytes to
  380. :param int size: number of bytes requested to be written to the buffer
  381. :returns: int -- number of bytes actually written
  382. """
  383. written = 0
  384. if self.headers_unread:
  385. written += buffer.append(self.headers)
  386. self.headers_unread = False
  387. while total_len(self.body) > 0 and (size == -1 or written < size):
  388. amount_to_read = size
  389. if size != -1:
  390. amount_to_read = size - written
  391. written += buffer.append(self.body.read(amount_to_read))
  392. return written
  393. class CustomBytesIO(io.BytesIO):
  394. def __init__(self, buffer=None, encoding='utf-8'):
  395. buffer = encode_with(buffer, encoding)
  396. super(CustomBytesIO, self).__init__(buffer)
  397. def _get_end(self):
  398. current_pos = self.tell()
  399. self.seek(0, 2)
  400. length = self.tell()
  401. self.seek(current_pos, 0)
  402. return length
  403. @property
  404. def len(self):
  405. length = self._get_end()
  406. return length - self.tell()
  407. def append(self, bytes):
  408. with reset(self):
  409. written = self.write(bytes)
  410. return written
  411. def smart_truncate(self):
  412. to_be_read = total_len(self)
  413. already_read = self._get_end() - to_be_read
  414. if already_read >= to_be_read:
  415. old_bytes = self.read()
  416. self.seek(0, 0)
  417. self.truncate()
  418. self.write(old_bytes)
  419. self.seek(0, 0) # We want to be at the beginning
  420. class FileWrapper(object):
  421. def __init__(self, file_object):
  422. self.fd = file_object
  423. @property
  424. def len(self):
  425. return total_len(self.fd) - self.fd.tell()
  426. def read(self, length=-1):
  427. return self.fd.read(length)
  428. class FileFromURLWrapper(object):
  429. """File from URL wrapper.
  430. The :class:`FileFromURLWrapper` object gives you the ability to stream file
  431. from provided URL in chunks by :class:`MultipartEncoder`.
  432. Provide a stateless solution for streaming file from one server to another.
  433. You can use the :class:`FileFromURLWrapper` without a session or with
  434. a session as demonstated by the examples below:
  435. .. code-block:: python
  436. # no session
  437. import requests
  438. from requests_toolbelt import MultipartEncoder, FileFromURLWrapper
  439. url = 'https://httpbin.org/image/png'
  440. streaming_encoder = MultipartEncoder(
  441. fields={
  442. 'file': FileFromURLWrapper(url)
  443. }
  444. )
  445. r = requests.post(
  446. 'https://httpbin.org/post', data=streaming_encoder,
  447. headers={'Content-Type': streaming_encoder.content_type}
  448. )
  449. .. code-block:: python
  450. # using a session
  451. import requests
  452. from requests_toolbelt import MultipartEncoder, FileFromURLWrapper
  453. session = requests.Session()
  454. url = 'https://httpbin.org/image/png'
  455. streaming_encoder = MultipartEncoder(
  456. fields={
  457. 'file': FileFromURLWrapper(url, session=session)
  458. }
  459. )
  460. r = session.post(
  461. 'https://httpbin.org/post', data=streaming_encoder,
  462. headers={'Content-Type': streaming_encoder.content_type}
  463. )
  464. """
  465. def __init__(self, file_url, session=None):
  466. self.session = session or requests.Session()
  467. requested_file = self._request_for_file(file_url)
  468. self.len = int(requested_file.headers['content-length'])
  469. self.raw_data = requested_file.raw
  470. def _request_for_file(self, file_url):
  471. """Make call for file under provided URL."""
  472. response = self.session.get(file_url, stream=True)
  473. content_length = response.headers.get('content-length', None)
  474. if content_length is None:
  475. error_msg = (
  476. "Data from provided URL {url} is not supported. Lack of "
  477. "content-length Header in requested file response.".format(
  478. url=file_url)
  479. )
  480. raise FileNotSupportedError(error_msg)
  481. elif not content_length.isdigit():
  482. error_msg = (
  483. "Data from provided URL {url} is not supported. content-length"
  484. " header value is not a digit.".format(url=file_url)
  485. )
  486. raise FileNotSupportedError(error_msg)
  487. return response
  488. def read(self, chunk_size):
  489. """Read file in chunks."""
  490. chunk_size = chunk_size if chunk_size >= 0 else self.len
  491. chunk = self.raw_data.read(chunk_size) or b''
  492. self.len -= len(chunk) if chunk else 0 # left to read
  493. return chunk