242 lines
8.0 KiB
Python
242 lines
8.0 KiB
Python
from datetime import timedelta
|
|
import io
|
|
from typing import Dict
|
|
|
|
from aiohttp import ClientSession
|
|
from aiohttp.client import ClientResponse, ClientPayloadError, ClientOSError
|
|
from asyncio.exceptions import TimeoutError
|
|
from collections import deque
|
|
|
|
from ..exceptions import LogicError, ensureKey
|
|
from ..logger import getLogger
|
|
from ..time import Time
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
CONTENT_LENGTH_HEADER = "content-length"
|
|
CONTENT_LENGTH_ERROR = "Content size must be provided if the webserver doesn't provide it"
|
|
SERVER_CONTENT_LENGTH_ERROR = "Server returned a content length that didn't match the requested size"
|
|
POSITION_ERROR_MESSAGE = "AsyncHttpGetter must also be set up at position 0"
|
|
DEFAULT_CHUNK_SIZE = 1024 * 1024
|
|
|
|
|
|
# This class is dumb but it gets around a dumb problem
|
|
class Stupid(io.BytesIO):
|
|
def __len__(self):
|
|
return len(self.getvalue())
|
|
|
|
|
|
class AsyncHttpGetter:
|
|
def __init__(self, url, headers: Dict[str, str], session, size: int = None, timeout=None, timeoutFactory=None, otherErrorFactory=None, time: Time = None):
|
|
self._url: str = url
|
|
|
|
# Current position of the stream
|
|
self._position: int = 0
|
|
|
|
# Total size of the stream
|
|
self._size: int = size
|
|
|
|
# Headers that should get sent with every request
|
|
self._headers: Dict[str, str] = headers
|
|
|
|
# Session used to make http requests.
|
|
self._session: ClientSession = session
|
|
|
|
# Most recent session request
|
|
self._response: ClientResponse = None
|
|
|
|
# Where the resposne currently starts
|
|
self._responseStart = 0
|
|
|
|
self._history = deque()
|
|
self._time = time
|
|
self._startTime = self._time.now()
|
|
self.timeoutFactory = timeoutFactory
|
|
self.otherErrorFactory = otherErrorFactory
|
|
self.timeout = timeout
|
|
|
|
async def setup(self):
|
|
if not self._position == 0:
|
|
raise LogicError(POSITION_ERROR_MESSAGE)
|
|
await self._startReadRemoteAt(0)
|
|
if CONTENT_LENGTH_HEADER in self._response.headers:
|
|
self._size = int(ensureKey(
|
|
CONTENT_LENGTH_HEADER, self._response.headers, "web server get request's headers"))
|
|
elif self._size is None:
|
|
raise LogicError(
|
|
CONTENT_LENGTH_ERROR)
|
|
|
|
self._history.append([self._time.now(), 0])
|
|
return self._size
|
|
|
|
def _ensureSetup(self):
|
|
if self._size is None:
|
|
raise LogicError("AsyncHttpGetter.setup() must be called first")
|
|
|
|
def size(self) -> int:
|
|
self._ensureSetup()
|
|
return self._size
|
|
|
|
def __len__(self):
|
|
return self.size()
|
|
|
|
def position(self, pos=None):
|
|
if pos is not None:
|
|
self._position = pos
|
|
return self._position
|
|
|
|
async def generator(self, chunk_size):
|
|
while True:
|
|
chunk = await self.read(chunk_size)
|
|
if len(chunk.getbuffer()) == 0:
|
|
break
|
|
yield chunk.getbuffer()
|
|
|
|
def progress(self):
|
|
if self._size == 0:
|
|
return 0
|
|
self._ensureSetup()
|
|
if self._size == 0:
|
|
return 0
|
|
return 100 * float(self.position()) / float(self._size)
|
|
|
|
# return the estimated speed of the tranfser in bytes/second
|
|
def speed(self, period: timedelta = timedelta(seconds=10)):
|
|
if len(self._history) < 2:
|
|
return None
|
|
now = self._time.now()
|
|
intervals = []
|
|
current = self._history[0]
|
|
last_speed = 0
|
|
for x in range(1, len(self._history)):
|
|
next = self._history[x]
|
|
seconds = (next[0] - current[0]).total_seconds()
|
|
data = next[1] - current[1]
|
|
if seconds == 0:
|
|
speed = 0 # avoid div by 0
|
|
else:
|
|
speed = data / seconds
|
|
intervals.append([current[0], next[0], speed])
|
|
current = next
|
|
last_speed = speed
|
|
|
|
time_since_last = now - current[0]
|
|
if time_since_last > period:
|
|
# linear decay the "estimated" rate over the period
|
|
diff = time_since_last - period
|
|
if diff > period:
|
|
return 0
|
|
else:
|
|
return last_speed * (1 - diff.total_seconds() / period.total_seconds())
|
|
|
|
# behave as though we continued with the pseed form the last interval
|
|
intervals.append([current[0], now, last_speed])
|
|
|
|
# calculate the time-averaged speed over the given period
|
|
stop = now
|
|
start = now - period
|
|
total = 0
|
|
minimum_time = now
|
|
for interval in intervals:
|
|
if start > interval[1]:
|
|
continue
|
|
if stop < interval[0]:
|
|
continue
|
|
overlap_start = max(start, interval[0])
|
|
overlap_stop = min(stop, interval[1])
|
|
if overlap_stop > overlap_start:
|
|
total += (overlap_stop - overlap_start).total_seconds() * interval[2]
|
|
if overlap_start < minimum_time:
|
|
minimum_time = overlap_start
|
|
return total / (now - minimum_time).total_seconds()
|
|
|
|
def startTime(self):
|
|
return self._startTime
|
|
|
|
def __format__(self, format_spec: str) -> str:
|
|
return str(int(self.progress()))
|
|
|
|
async def _startReadRemoteAt(self, where: int):
|
|
headers = self._headers.copy()
|
|
# request a byte range
|
|
if where != 0:
|
|
headers['range'] = "bytes=%s-%s" % (self._position, self._size - 1)
|
|
if self._response is not None:
|
|
self._response.release()
|
|
try:
|
|
resp = await self._session.get(self._url, headers=headers, timeout=self.timeout)
|
|
except TimeoutError:
|
|
if self.timeoutFactory is not None:
|
|
raise self.timeoutFactory()
|
|
raise
|
|
except ClientPayloadError:
|
|
if self.otherErrorFactory is not None:
|
|
raise self.otherErrorFactory()
|
|
raise
|
|
except ClientOSError:
|
|
if self.otherErrorFactory is not None:
|
|
raise self.otherErrorFactory()
|
|
raise
|
|
resp.raise_for_status()
|
|
if where == 0 and self._size is not None and CONTENT_LENGTH_HEADER in resp.headers and int(resp.headers[CONTENT_LENGTH_HEADER]) != self._size:
|
|
resp.release()
|
|
raise LogicError(SERVER_CONTENT_LENGTH_ERROR)
|
|
self._response = resp
|
|
self._responseStart = where
|
|
|
|
async def read(self, count=DEFAULT_CHUNK_SIZE):
|
|
self._ensureSetup()
|
|
ret = Stupid()
|
|
if self._size is not None and self._position >= self._size:
|
|
return ret
|
|
|
|
# See if we need to move the stream elsewhere
|
|
if self._responseStart != self._position:
|
|
# Reset the stream's position
|
|
await self._startReadRemoteAt(self._position)
|
|
|
|
# Limit by how much we can get from the stream
|
|
needed = min(count, self.size() - self._position)
|
|
|
|
# And then get it
|
|
try:
|
|
data = await self._response.content.readexactly(needed)
|
|
except TimeoutError:
|
|
if self.timeoutFactory is not None:
|
|
raise self.timeoutFactory()
|
|
raise
|
|
except ClientPayloadError:
|
|
if self.otherErrorFactory is not None:
|
|
raise self.otherErrorFactory()
|
|
raise
|
|
except ClientOSError:
|
|
if self.otherErrorFactory is not None:
|
|
raise self.otherErrorFactory()
|
|
raise
|
|
ret.write(data)
|
|
# Keep track of where we are in the stream
|
|
self._responseStart += len(data)
|
|
self._position += len(data)
|
|
self._history.append([self._time.now(), self._position])
|
|
if len(self._history) > 50:
|
|
self._history.popleft()
|
|
|
|
ret.seek(0)
|
|
return ret
|
|
|
|
async def __aenter__(self):
|
|
await self.setup()
|
|
|
|
async def __aexit__(self, type, value, traceback):
|
|
if self._response is not None:
|
|
self._response.release()
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
val = await self.read()
|
|
if len(val) == 0:
|
|
raise StopAsyncIteration
|
|
return val.getbuffer()
|