diff --git a/utils.py b/utils.py index bdf13a0..5521122 100644 --- a/utils.py +++ b/utils.py @@ -3,7 +3,7 @@ import hashlib from pathlib import Path from functools import wraps -def _disk_cache_decorator(subdir_name: str, ext: str, read_func, write_func): +def _disk_cache_decorator(subdir_name: str, ext: str, read_func, write_func, should_cache = None): """ Generic internal caching decorator handling filename hashing and io abstraction. """ @@ -30,19 +30,20 @@ def _disk_cache_decorator(subdir_name: str, ext: str, read_func, write_func): result = func(*args, **kwargs) - cache_p.parent.mkdir(parents=True, exist_ok=True) - write_func(cache_p, result) - - # Write arguments to a .txt file for easy lookup - txt_p = cache_p.with_suffix('.txt') - txt_p.write_text(f"Args: {args}\nKwargs: {kwargs}", encoding="utf-8") + if should_cache is None or should_cache(result): + cache_p.parent.mkdir(parents=True, exist_ok=True) + write_func(cache_p, result) + + # Write arguments to a .txt file for easy lookup + txt_p = cache_p.with_suffix('.txt') + txt_p.write_text(f"Args: {args}\nKwargs: {kwargs}", encoding="utf-8") return result return wrapper return decorator -def with_disk_cache(subdir_name: str): +def with_disk_cache(subdir_name: str, should_cache = None): """ A decorator to cache function results to a local JSON file. The cache file is stored in `data//.json`. @@ -51,11 +52,12 @@ def with_disk_cache(subdir_name: str): subdir_name, ".json", read_func=lambda p: json.loads(p.read_text(encoding="utf-8")), - write_func=lambda p, res: p.write_text(json.dumps(res, ensure_ascii=False, indent=2), encoding="utf-8") + write_func=lambda p, res: p.write_text(json.dumps(res, ensure_ascii=False, indent=2), encoding="utf-8"), + should_cache=should_cache ) -def with_binary_disk_cache(subdir_name: str, ext: str = ".bin"): +def with_binary_disk_cache(subdir_name: str, ext: str = ".bin", should_cache = None): """ A decorator to cache binary function results to a local file. The cache file is stored in `data//`. @@ -64,5 +66,6 @@ def with_binary_disk_cache(subdir_name: str, ext: str = ".bin"): subdir_name, ext, read_func=lambda p: p.read_bytes(), - write_func=lambda p, res: p.write_bytes(res) + write_func=lambda p, res: p.write_bytes(res), + should_cache=should_cache ) diff --git a/utils_mteam.py b/utils_mteam.py index 1385a20..9369d1d 100644 --- a/utils_mteam.py +++ b/utils_mteam.py @@ -22,7 +22,14 @@ def _get_mteam_headers() -> dict: } -@with_disk_cache('search_mteam_torrents') +def is_mteam_success(res) -> bool: + """Predicate to check if an M-Team API response is successful.""" + if isinstance(res, dict): + code = str(res.get("code", "")) + return code == "0" or code == "200" + return False + +@with_disk_cache('search_mteam_torrents', should_cache=is_mteam_success) def search_mteam_torrents(imdb_url: str, page_number: int = 1, page_size: int = 100) -> dict: """ Search M-Team for torrents using IMDb URL. @@ -51,7 +58,7 @@ def search_mteam_torrents(imdb_url: str, page_number: int = 1, page_size: int = return response.json() -@with_disk_cache('mteam_imdb_info') +@with_disk_cache('mteam_imdb_info', should_cache=is_mteam_success) def mteam_imdb_info(id: str) -> dict: """ Search M-Team for torrents using IMDb URL. @@ -145,4 +152,4 @@ def format_mteam_torrent(t: dict) -> str: if __name__ == "__main__": - print(mteam_imdb_info('tt38872297')) + print(mteam_imdb_info('tt0371724'))