You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

143 lines
5.2 KiB

4 years ago
  1. import io
  2. from urllib.parse import urlencode
  3. from multidict import MultiDict, MultiDictProxy
  4. from . import hdrs, multipart, payload
  5. from .helpers import guess_filename
  6. __all__ = ('FormData',)
  7. class FormData:
  8. """Helper class for multipart/form-data and
  9. application/x-www-form-urlencoded body generation."""
  10. def __init__(self, fields=(), quote_fields=True, charset=None):
  11. self._writer = multipart.MultipartWriter('form-data')
  12. self._fields = []
  13. self._is_multipart = False
  14. self._quote_fields = quote_fields
  15. self._charset = charset
  16. if isinstance(fields, dict):
  17. fields = list(fields.items())
  18. elif not isinstance(fields, (list, tuple)):
  19. fields = (fields,)
  20. self.add_fields(*fields)
  21. @property
  22. def is_multipart(self):
  23. return self._is_multipart
  24. def add_field(self, name, value, *, content_type=None, filename=None,
  25. content_transfer_encoding=None):
  26. if isinstance(value, io.IOBase):
  27. self._is_multipart = True
  28. elif isinstance(value, (bytes, bytearray, memoryview)):
  29. if filename is None and content_transfer_encoding is None:
  30. filename = name
  31. type_options = MultiDict({'name': name})
  32. if filename is not None and not isinstance(filename, str):
  33. raise TypeError('filename must be an instance of str. '
  34. 'Got: %s' % filename)
  35. if filename is None and isinstance(value, io.IOBase):
  36. filename = guess_filename(value, name)
  37. if filename is not None:
  38. type_options['filename'] = filename
  39. self._is_multipart = True
  40. headers = {}
  41. if content_type is not None:
  42. if not isinstance(content_type, str):
  43. raise TypeError('content_type must be an instance of str. '
  44. 'Got: %s' % content_type)
  45. headers[hdrs.CONTENT_TYPE] = content_type
  46. self._is_multipart = True
  47. if content_transfer_encoding is not None:
  48. if not isinstance(content_transfer_encoding, str):
  49. raise TypeError('content_transfer_encoding must be an instance'
  50. ' of str. Got: %s' % content_transfer_encoding)
  51. headers[hdrs.CONTENT_TRANSFER_ENCODING] = content_transfer_encoding
  52. self._is_multipart = True
  53. self._fields.append((type_options, headers, value))
  54. def add_fields(self, *fields):
  55. to_add = list(fields)
  56. while to_add:
  57. rec = to_add.pop(0)
  58. if isinstance(rec, io.IOBase):
  59. k = guess_filename(rec, 'unknown')
  60. self.add_field(k, rec)
  61. elif isinstance(rec, (MultiDictProxy, MultiDict)):
  62. to_add.extend(rec.items())
  63. elif isinstance(rec, (list, tuple)) and len(rec) == 2:
  64. k, fp = rec
  65. self.add_field(k, fp)
  66. else:
  67. raise TypeError('Only io.IOBase, multidict and (name, file) '
  68. 'pairs allowed, use .add_field() for passing '
  69. 'more complex parameters, got {!r}'
  70. .format(rec))
  71. def _gen_form_urlencoded(self):
  72. # form data (x-www-form-urlencoded)
  73. data = []
  74. for type_options, _, value in self._fields:
  75. data.append((type_options['name'], value))
  76. charset = self._charset if self._charset is not None else 'utf-8'
  77. if charset == 'utf-8':
  78. content_type = 'application/x-www-form-urlencoded'
  79. else:
  80. content_type = ('application/x-www-form-urlencoded; '
  81. 'charset=%s' % charset)
  82. return payload.BytesPayload(
  83. urlencode(data, doseq=True, encoding=charset).encode(),
  84. content_type=content_type)
  85. def _gen_form_data(self):
  86. """Encode a list of fields using the multipart/form-data MIME format"""
  87. for dispparams, headers, value in self._fields:
  88. try:
  89. if hdrs.CONTENT_TYPE in headers:
  90. part = payload.get_payload(
  91. value, content_type=headers[hdrs.CONTENT_TYPE],
  92. headers=headers, encoding=self._charset)
  93. else:
  94. part = payload.get_payload(
  95. value, headers=headers, encoding=self._charset)
  96. except Exception as exc:
  97. raise TypeError(
  98. 'Can not serialize value type: %r\n '
  99. 'headers: %r\n value: %r' % (
  100. type(value), headers, value)) from exc
  101. if dispparams:
  102. part.set_content_disposition(
  103. 'form-data', quote_fields=self._quote_fields, **dispparams
  104. )
  105. # FIXME cgi.FieldStorage doesn't likes body parts with
  106. # Content-Length which were sent via chunked transfer encoding
  107. part.headers.popall(hdrs.CONTENT_LENGTH, None)
  108. self._writer.append_payload(part)
  109. return self._writer
  110. def __call__(self):
  111. if self._is_multipart:
  112. return self._gen_form_data()
  113. else:
  114. return self._gen_form_urlencoded()