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