1
 2from PIL import Image, UnidentifiedImageError
 3from io import BytesIO
 4import requests, random, time, logging
 5
 6base_url = "https://danbooru.donmai.us/"
 7base_url_test = "https://testbooru.donmai.us/"
 8
 9page_suffix = "post/index.json"
10post_suffix = "posts/"
11
12ratings = [
13    'g', # general
14    's', # sensitive
15    'q', # questionable
16    'e', # explicit
17]
18
19rating_normal = "g,s"
20rating_lewd = "q"
21
22limit = 100
23max_pages = 1000
24sleep_seconds = 3
25max_tries = 5
26
27supported_file_types = [
28    ".jpg",
29    ".jpeg",
30    ".png"
31]
32
33def _valid_extension(fname: str):
34    for t in supported_file_types:
35        if fname.lower().endswith(t):
36            return True
37    return False
38
39def get_random_image(rating=rating_normal):
40    params = {
41        "limit": limit,
42        #"tags": "order:change_desc rating:" + rating,
43        "tags": "rating:" + rating,
44    }
45    count = 0
46    while count < max_tries:
47        params['page'] = random.randint(1, max_pages)
48        page = requests.get(base_url + page_suffix, params).json()
49        n = random.randint(0, limit - 1)
50        
51        #print("Page: " + str(params['page']))
52        #print("File: " + str(n))
53        try:
54            file_url = page[n]['file_url']
55            if not _valid_extension(file_url):
56                raise Exception
57            r = requests.get(file_url)
58            img = Image.open(BytesIO(r.content))
59            link = base_url + post_suffix + str(page[n]['id'])
60            return img, link
61            
62        except (KeyError, IndexError, Exception):
63            logging.warning("Can't display image.")
64        except UnidentifiedImageError:
65            logging.warning("Unidentified image: " + file_url)
66        
67        count += 1
68        logging.warning(f"Try #{count} failed.\n")
69        #time.sleep(sleep_seconds)
70    logging.error(f"Reached {count} tries. Giving up.")
71    return None, None