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