hyperboria/nexus/pylon/client.py
the-superpirate 7389f85ef2 - Send Pylon to golden
-  Move fancy naming to Hub
-  Adopt translations
-  Fix long authors list in bot
-  Support CJK in scitech
-  Fixed pylon bug with decoding improper unicode ch...

GitOrigin-RevId: 74f73c44f749a71cb65dd5ddd3416f32a83d329f
2021-01-08 23:09:19 +03:00

99 lines
3.4 KiB
Python

import asyncio
from typing import (
AsyncIterable,
Iterable,
Optional,
)
import aiohttp
import aiohttp.client_exceptions
from aiohttp_socks import (
ProxyConnectionError,
ProxyError,
)
from aiokit import AioThing
from library.logging import error_log
from nexus.pylon.exceptions import (
DownloadError,
NotFoundError,
)
from nexus.pylon.proto.file_pb2 import FileResponse as FileResponsePb
from nexus.pylon.sources import (
BaseSource,
LibgenDoiSource,
LibgenMd5Source,
LibraryLolSource,
SciHubSeSource,
SciHubTwSource,
)
from nexus.pylon.sources.specific import get_specific_sources_for_doi
from python_socks import ProxyTimeoutError
class PylonClient(AioThing):
def __init__(self, proxy: Optional[str] = None, resolve_proxy: Optional[str] = None):
super().__init__()
self.proxy = proxy
self.resolve_proxy = resolve_proxy
async def by_doi(
self,
doi: str,
md5: Optional[str] = None,
) -> AsyncIterable[FileResponsePb]:
sources = []
sources.extend(get_specific_sources_for_doi(doi, proxy=self.proxy, resolve_proxy=self.resolve_proxy))
sources.extend([
SciHubSeSource(doi=doi, md5=md5, proxy=self.proxy, resolve_proxy=self.resolve_proxy),
SciHubTwSource(doi=doi, md5=md5, proxy=self.proxy, resolve_proxy=self.resolve_proxy),
LibgenDoiSource(doi=doi, md5=md5, proxy=self.proxy, resolve_proxy=self.resolve_proxy),
])
sources = filter(lambda x: x.is_enabled, sources)
async for resp in self.download(sources=sources):
yield resp
async def by_md5(
self,
md5: str,
) -> AsyncIterable[FileResponsePb]:
sources = filter(lambda x: x.is_enabled, [
LibraryLolSource(md5=md5, proxy=self.proxy, resolve_proxy=self.resolve_proxy),
LibgenMd5Source(md5=md5, proxy=self.proxy, resolve_proxy=self.resolve_proxy),
])
async for resp in self.download(sources=sources):
yield resp
async def download_source(self, source) -> AsyncIterable[FileResponsePb]:
try:
yield FileResponsePb(status=FileResponsePb.Status.RESOLVING, source=source.base_url)
async for prepared_request in source.resolve():
async for resp in source.execute_prepared_request(prepared_request=prepared_request):
yield resp
return
raise DownloadError(error='not_found', source=str(source))
except (
aiohttp.client_exceptions.ClientConnectionError,
aiohttp.client_exceptions.ClientPayloadError,
aiohttp.client_exceptions.ClientResponseError,
aiohttp.client_exceptions.TooManyRedirects,
asyncio.TimeoutError,
ProxyConnectionError,
ProxyTimeoutError,
ProxyError,
) as e:
raise DownloadError(nested_error=str(e), nested_error_cls=e.__class__.__name__)
async def download(self, sources: Iterable[BaseSource]) -> AsyncIterable[FileResponsePb]:
for source in sources:
try:
await source.start()
async for resp in self.download_source(source):
yield resp
return
except DownloadError as e:
error_log(e)
continue
finally:
await source.stop()
raise NotFoundError()