diff --git a/src/bilifm/season.py b/src/bilifm/season.py index cd2ad5b..5e45ab1 100644 --- a/src/bilifm/season.py +++ b/src/bilifm/season.py @@ -2,7 +2,7 @@ import typer -from .util import request +from .util import request, Retry headers: dict[str, str] = { "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36", @@ -32,16 +32,12 @@ def get_videos(self): "page_size": self.page_size, } + @Retry(self.__response_succeed, self.__handle_error_response) def wrapped_request(): """wrap request with retry""" - for _ in range(self.retry): - res = request( - method="get", url=self.season_url, params=params, headers=headers - ).json() - if self.__response_succeed(res): - return res - self.__handle_error_response(res) - return None + return request( + method="get", url=self.season_url, params=params, headers=headers + ).json() res = wrapped_request() if res is None: diff --git a/src/bilifm/series.py b/src/bilifm/series.py index 2b27cc0..e58e300 100644 --- a/src/bilifm/series.py +++ b/src/bilifm/series.py @@ -2,12 +2,7 @@ import typer -from .util import request - -headers: dict[str, str] = { - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36", - "Referer": "https://www.bilibili.com", -} +from .util import Retry, request class Series: @@ -30,16 +25,10 @@ def get_videos(self): "current_id": self.uid, } + @Retry(self.__response_succeed, self.__handle_error_response) def wrapped_request(): """wrap request with retry""" - for _ in range(self.retry): - res = request( - method="get", url=self.series_url, params=params, headers=headers - ).json() - if self.__response_succeed(res): - return res - self.__handle_error_response(res) - return None + return request(method="get", url=self.series_url, params=params).json() res = wrapped_request() if res is None: diff --git a/src/bilifm/util.py b/src/bilifm/util.py index d42b3fe..7f4d6dd 100644 --- a/src/bilifm/util.py +++ b/src/bilifm/util.py @@ -4,6 +4,7 @@ import urllib.parse from functools import reduce from hashlib import md5 +from typing import Callable import requests import typer @@ -192,3 +193,24 @@ def check_path(path: str): Directory = Annotated[str, typer.Option("-o", "--directory", callback=change_directory)] Path = typer.Argument(callback=check_path) + + +class Retry: + """Retry decorator""" + + def __init__(self, response_succeed, handle_error_response, total=3) -> None: + self.total = total + self.__response_succeed = response_succeed + self.__handle_error_response = handle_error_response + pass + + def __call__(self, request_func: Callable) -> Callable: + def wrapped_request(*args, **kwargs): + for _ in range(self.total): + res = request_func(*args, **kwargs) + if self.__response_succeed(res): + return res + self.__handle_error_response(res) + return None + + return wrapped_request