1import logging
 2import requests
 3from io import BytesIO
 4
 5from PIL import Image, UnidentifiedImageError
 6
 7max_tries = 5
 8supported_file_types = [".jpg", ".jpeg", ".png"]
 9endpoint = "https://api.waifu.pics"
10
11class APIException(Exception):
12    def __init__(self, status_code: int, message: str):
13        self.status_code = status_code
14        self.message = message
15        super().__init__(self.message)
16
17    def __str__(self):
18        return f"Error {self.status_code}: {self.message}"
19
20def _valid_extension(fname: str):
21    for t in supported_file_types:
22        if fname.lower().endswith(t):
23            return True
24    return False
25
26def get_random_image(nsfw=False):
27    count = 0
28    while count < max_tries:
29        try:
30            apiUrl = endpoint + f"/{'nsfw' if nsfw else 'sfw'}/waifu"
31            response = requests.get(apiUrl)    
32
33            content = response.json()
34            if response.status_code != 200:
35                raise APIException(response.status_code, content['message'])
36            
37            url = content["url"]
38            if not _valid_extension(url):
39                raise AssertionError
40            
41            r = requests.get(url)
42            b = BytesIO(r.content)
43            image = Image.open(b)
44            return image, url
45
46        except (KeyError, IndexError, AssertionError, APIException):
47            logging.warning("Can't display image.")
48        except UnidentifiedImageError:
49            logging.warning("Unidentified image: " + url)
50
51        count += 1
52        logging.warning(f"Try #{count} failed.\n")
53    logging.error(f"Reached {count} tries. Giving up.")
54    return None, None