Added search filter toolbar support for size, color, and type, as well as safesearch-setting functionality.

Author: rabbitflyer5

src/cli.py

@@ -1,4 +1,5 @@
 import argparse, os, sys
+from enum import Enum
 
 # Should revisit this to look for xdg_downloads in env
 def get_download_path():
@@ -29,7 +30,6 @@ def check_pos_int(val: int):
     else:
         raise ValueError
 
-
 def get_arguments(argv=sys.argv):
     """ 
     The cli front end for the scraper.
@@ -40,7 +40,7 @@ def get_arguments(argv=sys.argv):
     Returns:
     parser.parse_args() -- A struct with all required info to run the scraper
     """
-    parser = argparse.ArgumentParser(description="Scrape google for images")
+    parser = argparse.ArgumentParser(description="Scrape Google for images")
     parser.add_argument("keyword", 
                         help="the phrase used to find images",
                         type=str,
@@ -59,6 +59,31 @@ def get_arguments(argv=sys.argv):
                         type=check_pos_int,
                         nargs="?",
                         default=1)
+    parser.add_argument("-s", "--size",
+                        help="Restrict your search to a certain size of image. Can be 'large', 'medium', or 'icon'.",
+                        type=str,
+                        nargs="?",
+                        choices=['large','medium','icon'],
+                        default='')
+    parser.add_argument("--color",
+                        help="Search for a certain color of image. Can be 'red', 'orange', 'yellow', 'green', 'teal', 'blue', 'purple', 'pink', 'white', 'gray', 'black', 'brown', 'grayscale', or 'transparent'.",
+                        type=str,
+                        nargs="?",
+                        choices=['red', 'orange', 'yellow', 'green', 'teal', 'blue', 'purple', 'pink', 'white', 'gray', 'black', 'brown', 'grayscale', 'transparent'],
+                        default='')
+    parser.add_argument("-k", "--type",
+                        help="The type of image to search for. Can be 'clipart', 'lineart', or 'animated'.",
+                        type=str,
+                        nargs="?",
+                        choices=['clipart', 'lineart', 'animated'],
+                        dest="imgtype",
+                        default='')
+    parser.add_argument("-p", "--safesearch",
+                        help="Force the use of a specific safesearch setting. Can be 'on' or 'off'.",
+                        type=str,
+                        nargs="?",
+                        choices=['on', 'off'],
+                        default='')
     args = parser.parse_args(argv[1:])
     # Set default directory
     if args.directory is None:

src/main.py

@@ -5,7 +5,7 @@
 
 def main():
     args = get_arguments(sys.argv)
-    scrape_images(args.keyword[0], args.count, args.directory, args.threads)
+    scrape_images(args.keyword[0], args.count, args.directory, args.threads, args.size, args.color, args.imgtype, args.safesearch)
 
 if __name__ == "__main__":
     main()

src/scraper.py

@@ -41,6 +41,91 @@ def add_filetype(file_path: str):
         eprint(err)
         return 1
 
+
+def process_image_size(val: str):
+    if (val == 'large'):
+        return "isz:l"
+    elif (val == 'medium'):
+        return "isz:m"
+    elif (val == 'icon'):
+        return "isz:i"
+    else:
+        return ""
+
+def process_image_color(val: str):
+    if (val == "grayscale"):
+        return "ic:gray"
+    elif (val == "transparent"):
+        return "ic:trans"
+    elif (val == "red"):
+        return "ic:specific%2Cisc:red"
+    elif (val == "orange"):
+        return "ic:specific%2Cisc:orange"
+    elif (val == "yellow"):
+        return "ic:specific%2Cisc:yellow"
+    elif (val == "green"):
+        return "ic:specific%2Cisc:green"
+    elif (val == "teal"):
+        return "ic:specific%2Cisc:teal"
+    elif (val == "blue"):
+        return "ic:specific%2Cisc:blue"
+    elif (val == "purple"):
+        return "ic:specific%2Cisc:purple"
+    elif (val == "pink"):
+        return "ic:specific%2Cisc:pink"
+    elif (val == "white"):
+        return "ic:specific%2Cisc:white"
+    elif (val == "gray"):
+        return "ic:specific%2Cisc:gray"
+    elif (val == "black"):
+        return "ic:specific%2Cisc:black"
+    elif (val == "brown"):
+        return "ic:specific%2Cisc:brown"
+    else:
+        return ""
+
+def process_image_type(val: str):
+    if (val == "clipart"):
+        return "itp:clipart"
+    elif (val == "lineart"):
+        return "itp:lineart"
+    elif (val == "animated"):
+        return "itp:animated"
+    else:
+        return ""
+
+def process_safesearch(val: str):
+    if (val == "on"):
+        return "on"
+    elif (val == "off"):
+        return "off"
+    else:
+        return ""
+
+
+def setup_url(searchurl: str, imgsize: str, imgcolor: str, imgtype: str, safesearch: str):
+    features = [searchurl]
+    subfeatures = [[],[]]
+    if (imgsize != None):
+        subfeatures[0] += [process_image_size(imgsize)]
+    if (imgcolor != None):
+        subfeatures[0] += [process_image_color(imgcolor)]
+    if (imgtype != None):
+        subfeatures[0] += [process_image_type(imgtype)]
+    if (safesearch != None):
+        subfeatures[1] += [process_safesearch(safesearch)]
+    
+    delim1 = "&"
+    delim2 = "%2C"
+    
+    if (subfeatures[0] != []):
+        features += ["tbs=" + delim2.join(subfeatures[0])]
+    if (subfeatures[1] != []):
+        features += ["safe=" + delim2.join(subfeatures[1])]
+    
+    return delim1.join(features)
+
+
 ############################# scraping helpers ################################
 
 def get_image_urls(query: str, page: int):
@@ -159,7 +244,7 @@ def get_manifest(search_key: str, image_cnt: int):
 
 ################################# main api ####################################
 
-def scrape_images(search_key, image_cnt, directory, threads):
+def scrape_images(search_key, image_cnt, directory, threads, size, color, imgtype, safesearch):
     """ 
     Request manifest, generate paths, save files, get filetype. 
     This is the only function that should be called externally. 
@@ -174,7 +259,8 @@ def scrape_images(search_key, image_cnt, directory, threads):
         print("savedir: {}".format(directory))
     if not os.path.exists(directory):
         os.makedirs(directory)
-
+    global search_url
+    search_url = setup_url(search_url, size, color, imgtype, safesearch)
     id_url_manifest = get_manifest(search_key, image_cnt)
     with ThreadPoolExecutor(max_workers=threads) as pool:
         with tqdm(total=len(id_url_manifest)) as progress: