在生产环境下使用生成器像服务器传输大文件

Posted 2021-01-18 Richie Wen

废话不多说，直接上代码

  1 # -*- coding: utf-8 -*-
  2 # Created by richie at 2018/10/23
  3 import os
  4 import sys
  5 
  6 _ver = sys.version_info
  7 
  8 #: Python 2.x?
  9 is_py2 = (_ver[0] == 2)
 10 
 11 #: Python 3.x?
 12 is_py3 = (_ver[0] == 3)
 13 
 14 import httplib
 15 import string
 16 import random
 17 from array import array
 18 from urllib3.exceptions import LocationParseError
 19 from urllib3.util import parse_url
 20 
 21 if is_py2:
 22     from collections import Mapping
 23 
 24     range = xrange
 25 elif is_py3:
 26     from collections.abc import Mapping
 27 
 28     range = range
 29 
 30 _BOUNDARY_CHARS = string.digits + string.ascii_letters
 31 
 32 
 33 def to_key_val_list(value):
 34     if value is None:
 35         return None
 36 
 37     if isinstance(value, (str, bytes, bool, int)):
 38         raise ValueError(‘cannot encode objects that are not 2-tuples‘)
 39 
 40     if isinstance(value, Mapping):
 41         value = value.items()
 42 
 43     return list(value)
 44 
 45 
 46 class HTTPConnection(httplib.HTTPConnection):
 47     """
 48     改写send方法  使之能够接受生成器
 49     """
 50     def __enter__(self):
 51         return self
 52 
 53     def __exit__(self, *args):
 54         self.close()
 55 
 56     def send(self, data):
 57         """Send `data‘ to the server."""
 58         if self.sock is None:
 59             if self.auto_open:
 60                 self.connect()
 61             else:
 62                 raise httplib.NotConnected()
 63 
 64         if self.debuglevel > 0:
 65             print "send:", repr(data)
 66         blocksize = 8192
 67         if hasattr(data, ‘next‘):
 68             for value in data:
 69                 self.sock.sendall(value)
 70 
 71         elif hasattr(data, ‘read‘) and not isinstance(data, array):
 72             if self.debuglevel > 0: print "sendIng a read()able"
 73             datablock = data.read(blocksize)
 74             while datablock:
 75                 self.sock.sendall(datablock)
 76                 datablock = data.read(blocksize)
 77         else:
 78             self.sock.sendall(data)
 79 
 80 
 81 class Request(object):
 82 
 83     def __init__(self, method=None, url=None, data=None, files=None, ):
 84         # Default empty dicts for dict params.
 85         data = [] if data is None else data
 86         files = [] if files is None else files
 87 
 88         self.method = method
 89         self.url = url
 90         self.files = files
 91         self.data = data
 92         self.headers = {
 93             ‘Accept‘: ‘*/*‘,
 94             ‘Connection‘: ‘keep-alive‘,
 95         }
 96         self.__boundary = None
 97         self.__item_header = []
 98         self.end_tag = None
 99 
100         self.init_body()
101 
102     def __repr__(self):
103         return ‘<Request [%s]>‘ % (self.method)
104 
105     def get_boundary(self):
106         if self.__boundary is None:
107             self.__boundary = ‘‘.join(random.choice(_BOUNDARY_CHARS) for _ in range(32))
108         return self.__boundary
109 
110     def init_body(self):
111         """
112         组装数据
113         :return:
114         """
115         def escape_quote(s):
116             return s.replace(‘"‘, ‘\"‘)
117 
118         fields = to_key_val_list(self.data or {})
119         files = to_key_val_list(self.files or {})
120         lines = []
121         fpaths = []
122         fheads = []
123         if fields:
124             for name, value in fields:
125                 lines.extend((
126                     ‘--{0}‘.format(self.get_boundary()),
127                     ‘Content-Disposition: form-data; name="{0}"‘.format(escape_quote(name)),
128                     ‘‘,
129                     str(value),
130                 ))
131         if files:
132             __temp = []
133             for name, value in files:
134                 filename = os.path.basename(value)
135                 filename = filename.encode(‘utf-8‘)
136                 __temp.extend((
137                     ‘--{0}‘.format(self.get_boundary()),
138                     ‘Content-Disposition: form-data; name="{0}"; filename="{1}"‘.format(
139                         escape_quote(name), escape_quote(filename)),
140                     ‘
‘,
141                 ))
142                 fpaths.append(value)
143                 fheads.append(‘
‘.join(__temp))
144         fields_heads = ‘
‘.join(lines) + ‘
‘
145         self.end_tag = "--" + self.get_boundary() + "--
"
146 
147         content_type = str(‘multipart/form-data; boundary=%s‘ % self.get_boundary())
148         if content_type and (‘content-type‘ not in self.headers):
149             self.headers[‘Content-Type‘] = content_type
150 
151         self.prepare_content_length(fields_heads, fheads, fpaths)
152         self.data = self.prepare_data(fields_heads, fheads, fpaths)
153 
154     def prepare_data(self, fields_heads, fheads, fpaths):
155         """
156         通过生成器发送数据
157         :param fields_heads:
158         :param fheads:
159         :param fpaths:
160         :return:
161         """
162         b_size = 4096
163         yield fields_heads
164         for i in range(len(fheads)):
165             yield fheads[i]
166             file_obj = open(fpaths[i], ‘rb‘)
167             while True:
168                 block = file_obj.read(b_size)
169                 if block:
170                     yield block
171                 else:
172                     yield ‘
‘
173                     break
174             yield self.end_tag
175 
176     def prepare_content_length(self, fields_heads, fheads=None, fpaths=None):
177         """
178         准备发送内容的大小
179         :param fields_heads:
180         :param fheads:
181         :param fpaths:
182         :return:
183         """
184         if fields_heads is not None:
185             length = len(fields_heads)
186             if fpaths:
187                 for fpath in fpaths:
188                     length += os.path.getsize(fpath) + len(‘
‘)  # file length
189                 for fhead in fheads:
190                     length += len(fhead)
191                 length += len(self.end_tag)
192             if length:
193                 self.headers[‘Content-Length‘] = str(length)
194 
195         elif self.method not in (‘GET‘, ‘HEAD‘) and self.headers.get(‘Content-Length‘) is None:
196             self.headers[‘Content-Length‘] = ‘0‘
197 
198     def send(self):
199         """
200         通过httplib发送请求
201         :return:
202         """
203         # Support for unicode domain names and paths.
204         try:
205             scheme, auth, host, port, path, query, fragment = parse_url(self.url)
206         except LocationParseError as e:
207             raise Exception(*e.args)
208 
209         with HTTPConnection(host, port=port) as conn:
210             conn.request(self.method.upper(), self.url, self.data, self.headers)
211             # Receive the response from the server
212             try:
213                 # For Python 2.7, use buffering of HTTP responses
214                 response = conn.getresponse(buffering=True)
215             except TypeError:
216                 # For compatibility with Python 3.3+
217                 response = conn.getresponse()
218             status = response.status
219             content = response.read()
220         return status, content