New Addon
Google BK
This commit is contained in:
@@ -0,0 +1,241 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user