_multipart.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269
  1. import io
  2. import os
  3. import typing
  4. from pathlib import Path
  5. from ._types import (
  6. AsyncByteStream,
  7. FileContent,
  8. FileTypes,
  9. RequestData,
  10. RequestFiles,
  11. SyncByteStream,
  12. )
  13. from ._utils import (
  14. format_form_param,
  15. guess_content_type,
  16. peek_filelike_length,
  17. primitive_value_to_str,
  18. to_bytes,
  19. )
  20. def get_multipart_boundary_from_content_type(
  21. content_type: typing.Optional[bytes],
  22. ) -> typing.Optional[bytes]:
  23. if not content_type or not content_type.startswith(b"multipart/form-data"):
  24. return None
  25. # parse boundary according to
  26. # https://www.rfc-editor.org/rfc/rfc2046#section-5.1.1
  27. if b";" in content_type:
  28. for section in content_type.split(b";"):
  29. if section.strip().lower().startswith(b"boundary="):
  30. return section.strip()[len(b"boundary=") :].strip(b'"')
  31. return None
  32. class DataField:
  33. """
  34. A single form field item, within a multipart form field.
  35. """
  36. def __init__(
  37. self, name: str, value: typing.Union[str, bytes, int, float, None]
  38. ) -> None:
  39. if not isinstance(name, str):
  40. raise TypeError(
  41. f"Invalid type for name. Expected str, got {type(name)}: {name!r}"
  42. )
  43. if value is not None and not isinstance(value, (str, bytes, int, float)):
  44. raise TypeError(
  45. "Invalid type for value. Expected primitive type,"
  46. f" got {type(value)}: {value!r}"
  47. )
  48. self.name = name
  49. self.value: typing.Union[str, bytes] = (
  50. value if isinstance(value, bytes) else primitive_value_to_str(value)
  51. )
  52. def render_headers(self) -> bytes:
  53. if not hasattr(self, "_headers"):
  54. name = format_form_param("name", self.name)
  55. self._headers = b"".join(
  56. [b"Content-Disposition: form-data; ", name, b"\r\n\r\n"]
  57. )
  58. return self._headers
  59. def render_data(self) -> bytes:
  60. if not hasattr(self, "_data"):
  61. self._data = to_bytes(self.value)
  62. return self._data
  63. def get_length(self) -> int:
  64. headers = self.render_headers()
  65. data = self.render_data()
  66. return len(headers) + len(data)
  67. def render(self) -> typing.Iterator[bytes]:
  68. yield self.render_headers()
  69. yield self.render_data()
  70. class FileField:
  71. """
  72. A single file field item, within a multipart form field.
  73. """
  74. CHUNK_SIZE = 64 * 1024
  75. def __init__(self, name: str, value: FileTypes) -> None:
  76. self.name = name
  77. fileobj: FileContent
  78. headers: typing.Dict[str, str] = {}
  79. content_type: typing.Optional[str] = None
  80. # This large tuple based API largely mirror's requests' API
  81. # It would be good to think of better APIs for this that we could
  82. # include in httpx 2.0 since variable length tuples(especially of 4 elements)
  83. # are quite unwieldly
  84. if isinstance(value, tuple):
  85. if len(value) == 2:
  86. # neither the 3rd parameter (content_type) nor the 4th (headers)
  87. # was included
  88. filename, fileobj = value # type: ignore
  89. elif len(value) == 3:
  90. filename, fileobj, content_type = value # type: ignore
  91. else:
  92. # all 4 parameters included
  93. filename, fileobj, content_type, headers = value # type: ignore
  94. else:
  95. filename = Path(str(getattr(value, "name", "upload"))).name
  96. fileobj = value
  97. if content_type is None:
  98. content_type = guess_content_type(filename)
  99. has_content_type_header = any("content-type" in key.lower() for key in headers)
  100. if content_type is not None and not has_content_type_header:
  101. # note that unlike requests, we ignore the content_type provided in the 3rd
  102. # tuple element if it is also included in the headers requests does
  103. # the opposite (it overwrites the headerwith the 3rd tuple element)
  104. headers["Content-Type"] = content_type
  105. if isinstance(fileobj, io.StringIO):
  106. raise TypeError(
  107. "Multipart file uploads require 'io.BytesIO', not 'io.StringIO'."
  108. )
  109. if isinstance(fileobj, io.TextIOBase):
  110. raise TypeError(
  111. "Multipart file uploads must be opened in binary mode, not text mode."
  112. )
  113. self.filename = filename
  114. self.file = fileobj
  115. self.headers = headers
  116. def get_length(self) -> typing.Optional[int]:
  117. headers = self.render_headers()
  118. if isinstance(self.file, (str, bytes)):
  119. return len(headers) + len(to_bytes(self.file))
  120. file_length = peek_filelike_length(self.file)
  121. # If we can't determine the filesize without reading it into memory,
  122. # then return `None` here, to indicate an unknown file length.
  123. if file_length is None:
  124. return None
  125. return len(headers) + file_length
  126. def render_headers(self) -> bytes:
  127. if not hasattr(self, "_headers"):
  128. parts = [
  129. b"Content-Disposition: form-data; ",
  130. format_form_param("name", self.name),
  131. ]
  132. if self.filename:
  133. filename = format_form_param("filename", self.filename)
  134. parts.extend([b"; ", filename])
  135. for header_name, header_value in self.headers.items():
  136. key, val = f"\r\n{header_name}: ".encode(), header_value.encode()
  137. parts.extend([key, val])
  138. parts.append(b"\r\n\r\n")
  139. self._headers = b"".join(parts)
  140. return self._headers
  141. def render_data(self) -> typing.Iterator[bytes]:
  142. if isinstance(self.file, (str, bytes)):
  143. yield to_bytes(self.file)
  144. return
  145. if hasattr(self.file, "seek"):
  146. try:
  147. self.file.seek(0)
  148. except io.UnsupportedOperation:
  149. pass
  150. chunk = self.file.read(self.CHUNK_SIZE)
  151. while chunk:
  152. yield to_bytes(chunk)
  153. chunk = self.file.read(self.CHUNK_SIZE)
  154. def render(self) -> typing.Iterator[bytes]:
  155. yield self.render_headers()
  156. yield from self.render_data()
  157. class MultipartStream(SyncByteStream, AsyncByteStream):
  158. """
  159. Request content as streaming multipart encoded form data.
  160. """
  161. def __init__(
  162. self,
  163. data: RequestData,
  164. files: RequestFiles,
  165. boundary: typing.Optional[bytes] = None,
  166. ) -> None:
  167. if boundary is None:
  168. boundary = os.urandom(16).hex().encode("ascii")
  169. self.boundary = boundary
  170. self.content_type = "multipart/form-data; boundary=%s" % boundary.decode(
  171. "ascii"
  172. )
  173. self.fields = list(self._iter_fields(data, files))
  174. def _iter_fields(
  175. self, data: RequestData, files: RequestFiles
  176. ) -> typing.Iterator[typing.Union[FileField, DataField]]:
  177. for name, value in data.items():
  178. if isinstance(value, (tuple, list)):
  179. for item in value:
  180. yield DataField(name=name, value=item)
  181. else:
  182. yield DataField(name=name, value=value)
  183. file_items = files.items() if isinstance(files, typing.Mapping) else files
  184. for name, value in file_items:
  185. yield FileField(name=name, value=value)
  186. def iter_chunks(self) -> typing.Iterator[bytes]:
  187. for field in self.fields:
  188. yield b"--%s\r\n" % self.boundary
  189. yield from field.render()
  190. yield b"\r\n"
  191. yield b"--%s--\r\n" % self.boundary
  192. def get_content_length(self) -> typing.Optional[int]:
  193. """
  194. Return the length of the multipart encoded content, or `None` if
  195. any of the files have a length that cannot be determined upfront.
  196. """
  197. boundary_length = len(self.boundary)
  198. length = 0
  199. for field in self.fields:
  200. field_length = field.get_length()
  201. if field_length is None:
  202. return None
  203. length += 2 + boundary_length + 2 # b"--{boundary}\r\n"
  204. length += field_length
  205. length += 2 # b"\r\n"
  206. length += 2 + boundary_length + 4 # b"--{boundary}--\r\n"
  207. return length
  208. # Content stream interface.
  209. def get_headers(self) -> typing.Dict[str, str]:
  210. content_length = self.get_content_length()
  211. content_type = self.content_type
  212. if content_length is None:
  213. return {"Transfer-Encoding": "chunked", "Content-Type": content_type}
  214. return {"Content-Length": str(content_length), "Content-Type": content_type}
  215. def __iter__(self) -> typing.Iterator[bytes]:
  216. for chunk in self.iter_chunks():
  217. yield chunk
  218. async def __aiter__(self) -> typing.AsyncIterator[bytes]:
  219. for chunk in self.iter_chunks():
  220. yield chunk