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        try:
57            file_url = page[n]['file_url']
58            if not _valid_extension(file_url):
59                raise Exception
60            r = requests.get(file_url)
61            img = Image.open(BytesIO(r.content))
62            link = base_url + post_suffix + str(page[n]['id'])
63            return img, link
64            
65        except (KeyError, IndexError, Exception):
66            logging.warning("Can't display image.")
67        except UnidentifiedImageError:
68            logging.warning("Unidentified image: " + file_url)
69        
70        count += 1
71        logging.warning(f"Try #{count} failed.\n")
72        #time.sleep(sleep_seconds)
73    logging.error(f"Reached {count} tries. Giving up.")
74    return None, None