88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
# -*- coding=utf-8
|
|
import os
|
|
import uuid
|
|
|
|
|
|
class StreamBody(object):
|
|
def __init__(self, rt):
|
|
self._rt = rt
|
|
self._read_len = 0
|
|
self._content_len = 0
|
|
self._use_chunked = False
|
|
self._use_encoding = False
|
|
if 'Content-Length' in self._rt.headers:
|
|
self._content_len = int(self._rt.headers['Content-Length'])
|
|
elif 'Transfer-Encoding' in self._rt.headers and self._rt.headers['Transfer-Encoding'] == "chunked":
|
|
self._use_chunked = True
|
|
else:
|
|
raise IOError("create StreamBody failed without Content-Length header or Transfer-Encoding header")
|
|
|
|
if 'Content-Encoding' in self._rt.headers:
|
|
self._use_encoding = True
|
|
|
|
def __iter__(self):
|
|
"""提供一个默认的迭代器"""
|
|
return self._rt.iter_content(1024)
|
|
|
|
def __len__(self):
|
|
return self._content_len
|
|
|
|
def get_raw_stream(self):
|
|
"""提供原始流"""
|
|
return self._rt.raw
|
|
|
|
def get_stream(self, chunk_size=1024):
|
|
"""提供一个chunk可变的迭代器"""
|
|
return self._rt.iter_content(chunk_size=chunk_size)
|
|
|
|
def read(self, chunk_size=1024, auto_decompress=False):
|
|
chunk = None
|
|
if self._use_encoding and not auto_decompress:
|
|
chunk = self._rt.raw.read(chunk_size)
|
|
else:
|
|
try:
|
|
chunk = next(self._rt.iter_content(chunk_size))
|
|
except StopIteration:
|
|
return ''
|
|
return chunk
|
|
|
|
def get_stream_to_file(self, file_name, disable_tmp_file=False, auto_decompress=False):
|
|
"""保存流到本地文件"""
|
|
self._read_len = 0
|
|
tmp_file_name = "{file_name}_{uuid}".format(file_name=file_name, uuid=uuid.uuid4().hex)
|
|
if disable_tmp_file:
|
|
tmp_file_name = file_name
|
|
chunk_size = 1024 * 1024
|
|
with open(tmp_file_name, 'wb') as fp:
|
|
while True:
|
|
chunk = self.read(chunk_size, auto_decompress)
|
|
if not chunk:
|
|
break
|
|
self._read_len += len(chunk)
|
|
fp.write(chunk)
|
|
|
|
if not self._use_chunked and not (
|
|
self._use_encoding and auto_decompress) and self._read_len != self._content_len:
|
|
if os.path.exists(tmp_file_name):
|
|
os.remove(tmp_file_name)
|
|
raise IOError("download failed with incomplete file")
|
|
if file_name != tmp_file_name:
|
|
if os.path.exists(file_name):
|
|
os.remove(file_name)
|
|
os.rename(tmp_file_name, file_name)
|
|
|
|
def pget_stream_to_file(self, fdst, offset, expected_len, auto_decompress=False):
|
|
"""保存流到本地文件的offset偏移"""
|
|
self._read_len = 0
|
|
fdst.seek(offset, 0)
|
|
chunk_size = 1024 * 1024
|
|
while True:
|
|
chunk = self.read(chunk_size, auto_decompress)
|
|
if not chunk:
|
|
break
|
|
self._read_len += len(chunk)
|
|
fdst.write(chunk)
|
|
|
|
if not self._use_chunked and not (self._use_encoding and auto_decompress) and self._read_len != expected_len:
|
|
raise IOError("download failed with incomplete file")
|