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 = "rating:g,s"
20rating_lewd = "rating:q"
21
22
23
24supported_file_types = [
25    ".jpg",
26    ".jpeg",
27    ".png"
28]
29
30def _valid_extension(fname: str):
31    for t in supported_file_types:
32        if fname.lower().endswith(t):
33            return True
34    return False
35
36def get_random_image(rating=rating_normal, tags=""):
37    limit = 100
38    max_pages = 1000
39    sleep_seconds = 3
40    max_tries = 5
41
42    params = {
43        "limit": limit,
44        #"tags": "order:change_desc rating:" + rating,
45        "tags": rating + tags,
46    }
47    
48    if tags != "":
49        max_pages = 50
50        
51    count = 0
52    while count < max_tries:
53        params['page'] = random.randint(1, max_pages)
54        page = requests.get(base_url + page_suffix, params).json()
55        n = random.randint(0, params['limit'] - 1)
56        
57        #print("Page: " + str(params['page']))
58        #print("File: " + str(n))
59        try:
60            file_url = page[n]['file_url']
61            if not _valid_extension(file_url):
62                raise Exception
63            r = requests.get(file_url)
64            img = Image.open(BytesIO(r.content))
65            link = base_url + post_suffix + str(page[n]['id'])
66            return img, link
67            
68        except (KeyError, IndexError, Exception):
69            logging.warning("Can't display image.")
70        except UnidentifiedImageError:
71            logging.warning("Unidentified image: " + file_url)
72        
73        count += 1
74        logging.warning(f"Try #{count} failed.\n")
75        #time.sleep(sleep_seconds)
76    logging.error(f"Reached {count} tries. Giving up.")
77    return None, None