中国招标与采购网 - 更新流式下载

zbytb/utils/attachment.py

@@ -1,3 +1,4 @@
+import io
 import traceback
 import uuid
 
@@ -48,7 +49,17 @@ class AttachmentDownloader(AliYunService):
             return "{:.1f} kb".format(_kb)
 
     @staticmethod
+    def get_stream(file, response):
+        stream = io.BytesIO()
+        chunk_size = 1024  # 单次请求最大值
+        with open(file, 'wb') as f:
+            for data in response.iter_content(chunk_size=chunk_size):
+                stream.write(data)
+                f.write(data)
+        return stream.getvalue()
+
     def _download(
+            self,
             url: str,
             file: str,
             enable_proxy=False,
@@ -70,18 +81,12 @@ class AttachmentDownloader(AliYunService):
                         if req_kw in req.text:
                             with requests.post(url, **request_params) as req:
                                 if req.status_code == 200:
-                                    stream = req.content
-                                    with open(file, 'wb') as f:
-                                        f.write(stream)
-                                    return stream
+                                    return self.get_stream(file, req)
                                 else:
                                     retries += 1
                     else:
                         if req.status_code == 200:
-                            stream = req.content
-                            with open(file, 'wb') as f:
-                                f.write(stream)
-                            return stream
+                            return self.get_stream(file, req)
                         else:
                             retries += 1
 
@@ -117,7 +122,6 @@ class AttachmentDownloader(AliYunService):
 
         local_tmp_file = self._create_file(file_name, file_type)
 
-
         file_stream = self._download(
             download_url,
             local_tmp_file,