diff --git a/.env.example b/.env.example index 8ffd6373..01d6e503 100644 --- a/.env.example +++ b/.env.example @@ -50,6 +50,9 @@ IS_PUBLIC = false # Retry attempts RETRIES = 5 +# Request timeout +TIMEOUT = 10 + # Proxy for requests # PROXY = diff --git a/nazurin/__main__.py b/nazurin/__main__.py index 41c6958f..04e5e34f 100644 --- a/nazurin/__main__.py +++ b/nazurin/__main__.py @@ -6,6 +6,7 @@ from aiogram.dispatcher.filters import IDFilter from aiogram.types import ChatActions, Message, Update from aiogram.utils.exceptions import TelegramAPIError +from aiohttp import ClientResponseError from nazurin import config, dp from nazurin.utils import logger @@ -52,6 +53,9 @@ async def clear_cache(message: Message): async def on_error(update: Update, exception: Exception): try: raise exception + except ClientResponseError as error: + await update.message.reply( + f'Response Error: {error.status} {error.message}') except NazurinError as error: await update.message.reply(error.msg) except Exception as error: diff --git a/nazurin/config.py b/nazurin/config.py index 678e9b58..c4fc382e 100644 --- a/nazurin/config.py +++ b/nazurin/config.py @@ -38,6 +38,7 @@ ALLOW_GROUP = env.list('ALLOW_GROUP', subcast=int, default=[]) RETRIES = env.int('RETRIES', default=5) +TIMEOUT = env.int('TIMEOUT', default=10) PROXY = env.str('HTTP_PROXY', default=None) UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) \ AppleWebKit/537.36 (KHTML, like Gecko) \ diff --git a/nazurin/sites/Artstation/api.py b/nazurin/sites/Artstation/api.py index 058ba55c..696feb08 100644 --- a/nazurin/sites/Artstation/api.py +++ b/nazurin/sites/Artstation/api.py @@ -3,16 +3,19 @@ from nazurin.models import Caption, Illust, Image from nazurin.utils import Request +from nazurin.utils.decorators import network_retry from nazurin.utils.exceptions import NazurinError class Artstation(object): + @network_retry async def getPost(self, post_id: str): - """Fetch an post.""" + """Fetch a post.""" api = f"https://www.artstation.com/projects/{post_id}.json" async with Request() as request: async with request.get(api) as response: - if not response.status == 200: + if response.status == 404: raise NazurinError('Post not found') + response.raise_for_status() post = await response.json() return post diff --git a/nazurin/sites/Bilibili/api.py b/nazurin/sites/Bilibili/api.py index aa523383..67472e46 100644 --- a/nazurin/sites/Bilibili/api.py +++ b/nazurin/sites/Bilibili/api.py @@ -4,16 +4,27 @@ from nazurin.models import Caption, Illust, Image from nazurin.utils import Request +from nazurin.utils.decorators import network_retry +from nazurin.utils.exceptions import NazurinError class Bilibili(object): + @network_retry async def getDynamic(self, dynamic_id: int): """Get dynamic data from API.""" api = 'https://api.vc.bilibili.com/dynamic_svr/v1/dynamic_svr/get_dynamic_detail?dynamic_id=' + str( dynamic_id) async with Request() as request: async with request.get(api) as response: - source = await response.json() - card = json.loads(source['data']['card']['card']) + response.raise_for_status() + data = await response.json() + # For some IDs, the API returns code 0 but empty content + if data['code'] == 500207 or (data['code'] == 0 and 'card' + not in data['data'].keys()): + raise NazurinError('Dynamic not found') + if data['code'] != 0: + raise NazurinError('Failed to get dynamic: ' + + data['message']) + card = json.loads(data['data']['card']['card']) return card async def fetch(self, dynamic_id: int) -> Illust: @@ -26,6 +37,8 @@ async def fetch(self, dynamic_id: int) -> Illust: def getImages(self, card, dynamic_id: int) -> List[Image]: """Get all images in a dynamic card.""" + if 'item' not in card.keys() or 'pictures' not in card['item'].keys(): + raise NazurinError("No image found") pics = card['item']['pictures'] imgs = list() for index, pic in enumerate(pics): diff --git a/nazurin/sites/Danbooru/api.py b/nazurin/sites/Danbooru/api.py index 83f232a4..1fdd3c58 100644 --- a/nazurin/sites/Danbooru/api.py +++ b/nazurin/sites/Danbooru/api.py @@ -21,7 +21,7 @@ def __init__(self, site='danbooru'): async def getPost(self, post_id: Optional[int] = None, md5: Optional[str] = None): - """Fetch an post.""" + """Fetch a post.""" try: if post_id: post = await self.post_show(post_id) diff --git a/nazurin/sites/Gelbooru/api.py b/nazurin/sites/Gelbooru/api.py index 0bb50be1..cf87dfd8 100644 --- a/nazurin/sites/Gelbooru/api.py +++ b/nazurin/sites/Gelbooru/api.py @@ -12,6 +12,7 @@ async def getPost(self, post_id: int): post_id) async with Request() as request: async with request.get(api) as response: + response.raise_for_status() response = await response.json() if 'post' not in response.keys(): raise NazurinError('Post not found') diff --git a/nazurin/sites/Twitter/api.py b/nazurin/sites/Twitter/api.py index 5adcef05..a8d9a982 100644 --- a/nazurin/sites/Twitter/api.py +++ b/nazurin/sites/Twitter/api.py @@ -3,9 +3,11 @@ from nazurin.models import Caption, Illust, Image from nazurin.utils import Request +from nazurin.utils.decorators import network_retry from nazurin.utils.exceptions import NazurinError class Twitter(object): + @network_retry async def getTweet(self, status_id: int): """Get a tweet from API.""" # Old: 'https://syndication.twitter.com/tweets.json?ids='+ status_id +'&lang=en' @@ -15,6 +17,7 @@ async def getTweet(self, status_id: int): async with request.get(api) as response: if response.status == 404: raise NazurinError('Tweet not found or unavailable.') + response.raise_for_status() tweet = await response.json() return tweet diff --git a/nazurin/sites/Wallhaven/api.py b/nazurin/sites/Wallhaven/api.py index cc6729cb..ca2e752a 100644 --- a/nazurin/sites/Wallhaven/api.py +++ b/nazurin/sites/Wallhaven/api.py @@ -3,11 +3,13 @@ from nazurin.models import Caption, Illust, Image from nazurin.utils import Request +from nazurin.utils.decorators import network_retry from nazurin.utils.exceptions import NazurinError from .config import API_KEY class Wallhaven(object): + @network_retry async def getWallpaper(self, wallpaperId: str): """Get wallpaper information from API.""" api = 'https://wallhaven.cc/api/v1/w/' + wallpaperId @@ -21,9 +23,7 @@ async def getWallpaper(self, wallpaperId: str): raise NazurinError( 'You need to log in to view this wallpaper. ' + 'Please ensure that you have set a valid API key.') - if response.status == 429: - raise NazurinError( - 'Hit API rate limit, please try again later.') + response.raise_for_status() wallpaper = await response.json() if 'error' in wallpaper.keys(): raise NazurinError(wallpaper['error']) diff --git a/nazurin/sites/Weibo/api.py b/nazurin/sites/Weibo/api.py index 4d7ba0ca..1d8cd5ad 100644 --- a/nazurin/sites/Weibo/api.py +++ b/nazurin/sites/Weibo/api.py @@ -5,16 +5,17 @@ from nazurin.models import Caption, Illust, Image from nazurin.utils import Request +from nazurin.utils.decorators import network_retry from nazurin.utils.exceptions import NazurinError class Weibo(object): + @network_retry async def getPost(self, post_id: str): - """Fetch an post.""" + """Fetch a post.""" api = f"https://m.weibo.cn/detail/{post_id}" async with Request() as request: async with request.get(api) as response: - if not response.status == 200: - raise NazurinError('Post not found') + response.raise_for_status() html = await response.text() post = self.parseHtml(html) return post diff --git a/nazurin/sites/Zerochan/api.py b/nazurin/sites/Zerochan/api.py index 90fffd81..8df6d53a 100644 --- a/nazurin/sites/Zerochan/api.py +++ b/nazurin/sites/Zerochan/api.py @@ -3,22 +3,19 @@ from typing import List from urllib.parse import unquote -from aiohttp.client_exceptions import ClientResponseError from bs4 import BeautifulSoup from nazurin.models import Caption, Illust, Image from nazurin.utils import Request -from nazurin.utils.exceptions import NazurinError +from nazurin.utils.decorators import network_retry class Zerochan(object): + @network_retry async def getPost(self, post_id: int): async with Request() as request: async with request.get('https://www.zerochan.net/' + str(post_id)) as response: - try: - response.raise_for_status() - except ClientResponseError as err: - raise NazurinError(err) from None + response.raise_for_status() # Override post_id if there's a redirection TODO: Check if response.history: diff --git a/nazurin/utils/decorators.py b/nazurin/utils/decorators.py index d06de50c..77e4ee96 100644 --- a/nazurin/utils/decorators.py +++ b/nazurin/utils/decorators.py @@ -5,6 +5,8 @@ import tenacity from aiogram.types import ChatActions, Message from aiogram.utils.exceptions import RetryAfter +from aiohttp import ClientError, ClientResponseError +from tenacity import retry_if_exception, stop_after_attempt, wait_exponential from nazurin import config @@ -15,9 +17,20 @@ def after_log(retry_state): tenacity._utils.get_callback_name(retry_state.fn), retry_state.attempt_number, config.RETRIES) -retry = tenacity.retry(reraise=True, - stop=tenacity.stop_after_attempt(config.RETRIES), - after=after_log) +def exception_predicate(exception): + """Predicate to check if we should retry when an exception occurs.""" + if not isinstance(exception, + (ClientError, asyncio.exceptions.TimeoutError)): + return False + if isinstance(exception, ClientResponseError): + return exception.status in [408, 429, 500, 502, 503, 504] + return True + +network_retry = tenacity.retry(reraise=True, + stop=stop_after_attempt(config.RETRIES), + after=after_log, + retry=retry_if_exception(exception_predicate), + wait=wait_exponential(multiplier=1, max=8)) def chat_action(action: str): """Sends `action` while processing.""" @@ -33,6 +46,7 @@ async def wrapped_func(message: Message, *args, **kwargs): return decorator def async_wrap(func): + """Transform a synchronous function to an asynchronous one.""" @wraps(func) async def run(*args, loop=None, executor=None, **kwargs): if loop is None: diff --git a/nazurin/utils/network.py b/nazurin/utils/network.py index 104c907f..e4301833 100644 --- a/nazurin/utils/network.py +++ b/nazurin/utils/network.py @@ -1,19 +1,14 @@ -from aiohttp import ClientSession, TCPConnector +from aiohttp import ClientSession, ClientTimeout, TCPConnector -from nazurin.config import PROXY, UA - -from .decorators import retry +from nazurin.config import PROXY, TIMEOUT, UA class Request(ClientSession): - get = retry(ClientSession.get) - post = retry(ClientSession.post) - put = retry(ClientSession.put) - patch = retry(ClientSession.patch) - delete = retry(ClientSession.delete) - head = retry(ClientSession.head) - request = retry(ClientSession.request) - - def __init__(self, cookies=None, headers=None, **kwargs): + """Wrapped ClientSession with default user agent, timeout and proxy support.""" + def __init__(self, + cookies=None, + headers=None, + timeout=ClientTimeout(total=TIMEOUT), + **kwargs): if not headers: headers = dict() headers.update({'User-Agent': UA}) @@ -24,4 +19,5 @@ def __init__(self, cookies=None, headers=None, **kwargs): cookies=cookies, headers=headers, trust_env=True, + timeout=timeout, **kwargs)