53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
import logging
|
|
from typing import (
|
|
AsyncIterable,
|
|
Dict,
|
|
)
|
|
|
|
from aiokit import AioThing
|
|
from nexus.pylon.exceptions import (
|
|
DownloadError,
|
|
NotFoundError,
|
|
)
|
|
from nexus.pylon.proto.file_pb2 import FileResponse as FileResponsePb
|
|
from nexus.pylon.proxy_manager import ProxyManager
|
|
from nexus.pylon.source import Source
|
|
|
|
|
|
class PylonClient(AioThing):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.proxy_manager = ProxyManager(config.get('proxies'))
|
|
self.sources = []
|
|
if config.get('webdriver_hub') is None:
|
|
logging.getLogger('nexus_pylon').warning({
|
|
'action': 'missed_webdriver',
|
|
'mode': 'pylon',
|
|
})
|
|
for source_config in config['sources']:
|
|
source = Source.from_config(
|
|
proxy_manager=self.proxy_manager,
|
|
config=self.config,
|
|
source_config=source_config,
|
|
)
|
|
if source:
|
|
self.sources.append(source)
|
|
self.starts.append(source)
|
|
|
|
async def download(self, params: Dict) -> AsyncIterable[FileResponsePb]:
|
|
for source in self.sources:
|
|
if not source.is_match(params):
|
|
continue
|
|
try:
|
|
async for resp in source.download(params):
|
|
yield resp
|
|
return
|
|
except NotFoundError as e:
|
|
logging.getLogger('nexus_pylon').debug(e)
|
|
continue
|
|
except DownloadError as e:
|
|
logging.getLogger('nexus_pylon').warning(e)
|
|
continue
|
|
raise NotFoundError(params=params)
|