1
 2import requests, logging
 3from PIL import Image, UnidentifiedImageError
 4from io import BytesIO
 5
 6from anime_api.apis import WaifuPicsAPI
 7from anime_api.apis.waifu_pics.types import ImageCategory
 8
 9max_tries = 5
10supported_file_types = [ ".jpg", ".jpeg", ".png" ]
11api = WaifuPicsAPI()
12
13def _valid_extension(fname: str):
14    for t in supported_file_types:
15        if fname.lower().endswith(t):
16            return True
17    return False
18
19def get_random_image(nsfw=False):
20    
21    cat = ImageCategory.NSFW.WAIFU if nsfw else ImageCategory.SFW.WAIFU
22    
23    count = 0
24    while count < max_tries:
25        try:
26            img = api.get_random_image(category=cat)
27            if not _valid_extension(img.url):
28                raise Exception
29            r = requests.get(img.url)
30            image = Image.open(BytesIO(r.content))
31            
32            return image, img.url
33            
34        except (KeyError, IndexError, Exception):
35            logging.warning("Can't display image.")
36        except UnidentifiedImageError:
37            logging.warning("Unidentified image: " + img.url)
38        
39        count += 1
40        logging.warning(f"Try #{count} failed.\n")
41    logging.error(f"Reached {count} tries. Giving up.")
42    return None, None