form.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208
  1. import os
  2. import sys
  3. from _io import BytesIO
  4. import random
  5. from darabonba.utils.stream import BaseStream, READABLE
  6. from darabonba.core import DaraModel
  7. from urllib.parse import urlencode
  8. class Form:
  9. @staticmethod
  10. def to_form_string(
  11. val: dict,
  12. ) -> str:
  13. """
  14. Format a map to form string, like a=a%20b%20c
  15. @return: the form string
  16. """
  17. if not val:
  18. return ""
  19. keys = sorted(list(val))
  20. dic = {k: val[k] for k in keys if not isinstance(val[k], READABLE)}
  21. return urlencode(dic)
  22. @staticmethod
  23. def get_boundary():
  24. result = '%s' % int(random.random() * 100000000000000)
  25. return result.zfill(14)
  26. @staticmethod
  27. def to_file_form(form, boundary):
  28. return FileFormInputStream(form, boundary)
  29. def _length(o):
  30. if hasattr(o, 'len'):
  31. return o.len
  32. elif isinstance(o, BytesIO):
  33. return o.getbuffer().nbytes
  34. elif hasattr(o, 'fileno'):
  35. return os.path.getsize(o.name)
  36. return len(o)
  37. class FileFormInputStream(BaseStream):
  38. def __init__(self, form, boundary, size=1024):
  39. super().__init__(size)
  40. self.form = form
  41. self.boundary = boundary
  42. self.file_size_left = 0
  43. self.forms = {}
  44. self.files = {}
  45. self.files_keys = []
  46. self._to_map()
  47. self.form_str = b''
  48. self._build_str_forms()
  49. self.str_length = len(self.form_str)
  50. def _to_map(self):
  51. for k, v in self.form.items():
  52. if isinstance(v, FileField):
  53. self.files[k] = v
  54. self.files_keys.append(k)
  55. else:
  56. self.forms[k] = v
  57. def _build_str_forms(self):
  58. form_str = ''
  59. str_fmt = '--%s\r\nContent-Disposition: form-data; name="%s"\r\n\r\n%s\r\n'
  60. forms_list = sorted(list(self.forms))
  61. for key in forms_list:
  62. value = self.forms[key]
  63. form_str += str_fmt % (self.boundary, key, value)
  64. self.form_str = form_str.encode('utf-8')
  65. def _get_stream_length(self):
  66. file_length = 0
  67. for k, ff in self.files.items():
  68. field_length = len(ff.filename.encode('utf-8')) + len(ff.content_type) +\
  69. len(k.encode('utf-8')) + len(self.boundary) + 78
  70. file_length += _length(ff.content) + field_length
  71. stream_length = self.str_length + file_length + len(self.boundary) + 6
  72. return stream_length
  73. def __len__(self):
  74. return self._get_stream_length()
  75. def __iter__(self):
  76. return self
  77. def __next__(self):
  78. return self.read(self.size, loop=True)
  79. def file_str(self, size):
  80. # handle file object
  81. form_str = b''
  82. start_fmt = '--%s\r\nContent-Disposition: form-data; name="%s";'
  83. content_fmt = b' filename="%s"\r\nContent-Type: %s\r\n\r\n%s'
  84. if self.file_size_left:
  85. for key in self.files_keys[:]:
  86. if size <= 0:
  87. break
  88. file_field = self.files[key]
  89. file_content = file_field.content.read(size)
  90. if isinstance(file_content, str):
  91. file_content = file_content.encode('utf-8')
  92. if self.file_size_left <= size:
  93. form_str += b'%s\r\n' % file_content
  94. self.file_size_left = 0
  95. size -= len(file_content)
  96. self.files_keys.remove(key)
  97. else:
  98. form_str += file_content
  99. self.file_size_left -= size
  100. size -= len(file_content)
  101. else:
  102. for key in self.files_keys[:]:
  103. if size <= 0:
  104. break
  105. file_field = self.files[key]
  106. file_size = _length(file_field.content)
  107. self.file_size_left = file_size
  108. file_content = file_field.content.read(size)
  109. if isinstance(file_content, str):
  110. file_content = file_content.encode('utf-8')
  111. # build form_str
  112. start = start_fmt % (self.boundary, key)
  113. content = content_fmt % (
  114. file_field.filename.encode('utf-8'),
  115. file_field.content_type.encode('utf-8'),
  116. file_content
  117. )
  118. if self.file_size_left < size:
  119. form_str += b'%s%s\r\n' % (start.encode('utf-8'), content)
  120. self.file_size_left = 0
  121. size -= len(file_content)
  122. self.files_keys.remove(key)
  123. else:
  124. form_str += b'%s%s' % (start.encode('utf-8'), content)
  125. self.file_size_left -= size
  126. size -= len(file_content)
  127. return form_str
  128. def read(self, size=None, loop=False):
  129. if not self.files_keys and not self.form_str:
  130. self.refresh()
  131. if loop:
  132. raise StopIteration
  133. else:
  134. return b''
  135. if size is None:
  136. size = sys.maxsize
  137. if self.form_str:
  138. form_str = self.form_str[:size]
  139. self.form_str = self.form_str[size:]
  140. if len(form_str) < size:
  141. form_str += self.file_str(size)
  142. else:
  143. form_str = self.file_str(size)
  144. if not self.form_str and not self.files_keys:
  145. form_str += b'--%s--\r\n' % self.boundary.encode('utf-8')
  146. return form_str
  147. def refresh_cursor(self):
  148. for ff in self.files.values():
  149. if hasattr(ff.content, 'seek'):
  150. ff.content.seek(0, 0)
  151. def refresh(self):
  152. self.file_size_left = 0
  153. self._to_map()
  154. self._build_str_forms()
  155. self.refresh_cursor()
  156. class FileField(DaraModel):
  157. def __init__(self, filename=None, content_type=None, content=None):
  158. self.filename = filename
  159. self.content_type = content_type
  160. self.content = content
  161. def validate(self):
  162. self.validate_required(self.filename, 'filename')
  163. self.validate_required(self.content_type, 'content_type')
  164. self.validate_required(self.content, 'content')
  165. def to_map(self):
  166. result = {}
  167. result['filename'] = self.filename
  168. result['contentType'] = self.content_type
  169. result['content'] = self.content
  170. return result
  171. def from_map(self, map={}):
  172. self.filename = map.get('filename')
  173. self.content_type = map.get('contentType')
  174. self.content = map.get('content')
  175. return self