Automate the Boring Stuff Chapter 12 - Image Site Downloader

Question

The project outline:

Write a program that goes to a photo-sharing site like Flickr or Imgur, searches for a category of photos, and then downloads all the resulting images. You could write a program that works with any photo site that has a search feature.

My solution:

import sys, requests, os, bs4, time
from selenium import webdriver
from selenium.webdriver.firefox.options import Options
from selenium.webdriver.firefox.service import Service
from selenium.webdriver.common.by import By

def get_image_posts(url):
    # Returns a list of imgur post urls
    res = requests.get(url)
    res.raise_for_status()
    soup = bs4.BeautifulSoup(res.text, "html.parser")
    image_posts = ["https://imgur.com" + a["href"] for a in soup.select(".post a", href=True)]
    if not image_posts:
        print("Could not find any image posts.")
    else:
        return image_posts

def download_images(image_posts, driver, download_folder):
    # Finds the image url in the post and downloads it.
    for image_post in image_posts:
        driver.get(image_post)
        time.sleep(1)
        content = driver.page_source.encode("utf-8").strip()
        soup = bs4.BeautifulSoup(content, "html.parser")
        image_elem = soup.select_one("img.image-placeholder")
        if not image_elem:
            print("Could not find an image.")
        else:
            image_url = image_elem["src"]
            res = requests.get(image_url)
            res.raise_for_status()
            with open(os.path.join(download_folder, os.path.basename(image_url)), "wb") as image_file:
                for chunk in res.iter_content(100000):
                    image_file.write(chunk)

def main():
    search_term = sys.argv[1]
    download_folder = sys.argv[2]
    os.makedirs(download_folder, exist_ok=True)
    service = Service(r"C:\Program Files\WebDrivers\geckdriver\geckodriver.exe")
    options = Options()
    options.binary_location = r"C:\Program Files\Mozilla Firefox\firefox.exe"
    driver = webdriver.Firefox(service=service, options=options)
    driver.get("https://imgur.com/")
    try:
        cookie_elem = driver.find_element(By. XPATH, "/html/body/div[1]/div/div/div/div[2]/div/button[2]")
        search_elem = driver.find_element(By.CLASS_NAME, "Searchbar-textInput")
    except:
        print('Was not able to find an element with that name.')
    cookie_elem.click()
    search_elem.send_keys(search_term)
    search_elem.submit()
    time.sleep(1)
    url = driver.current_url
    image_posts = get_image_posts(url)
    download_images(image_posts, driver, download_folder)
    
if __name__ == '__main__':
    main()

I chose the site Imgur. My soluton works, though it is slow and clunky so I decided to limit it to one image download per page.

Answer 1

String Concatenation

Use str.format or f-strings instead of +:

# go from this
"https://imgur.com" + a["href"]

# to either str.format:
"https://imgur.com{}".format(a["href"])

# or using an f-string for python 3.5+
f"https://imgur.com{a['href']}"

Image Posts Iteration

In the case of an empty list, I would still return it:

def get_image_posts(url):
    ~snip~
    if not image_posts:
       print("No image posts")
    
    # this is still a list, it just could be empty
    return image_posts

Then, I'd move the iteration out of the download_images function and into main:

def main():
    ~snip~
    for post in get_image_posts(url):
        download_image(post, driver, download_folder)

Then I'd probably handle looking up the image via browser separately from downloading it:

def get_image_url(post):
    driver.get(post)
    content = driver.page_source.encode("utf-8").strip()
    soup = bs4.BeautifulSoup(content, "html.parser")
    image_elem = soup.select_one("img.image-placeholder")
    return image_elem

Then, it's easy to filter for downloading:

def main():
    ~snip~
    for post in get_image_posts(url):
        src = get_image_url(post, driver)
        
        # skip empty sources
        if not src:
            continue

        url = src['url']

        download_image(url, download_folder)

        # move the sleep here so it's more explicit
        time.sleep(1)

And now download_image looks like:

def download_images(url, driver, download_folder):
    res = requests.get(url)
    res.raise_for_status()
    with open(os.path.join(download_folder, os.path.basename(url)), "wb") as image_file:
        for chunk in res.iter_content(100000):
            image_file.write(chunk)

Line Length

The with open line is a bit long, so it might be worth defining a path separately like:

def download_images(url, driver, download_folder):
    res = requests.get(url)
    res.raise_for_status()

    filepath = os.path.join(
        download_folder,
        os.path.basename(url)
    )

    with open(filepath, 'wb') as fh:
        # write to file

sys.args

Since you are unpacking sys.argv[1] and sys.argv[2], I'd do that outside of main:

def main(search_term, download_folder):
    ~snip~



if __name__ == "__main__":
    search_term, download_folder, *_ = sys.argv[1:]

    # pass them as args to main
    main(search_term, download_folder)

Passing them as args to main makes it easier to ascertain what arguments are to be expected, and it makes this more flexible. If you want to reuse this in another module, you have to know to pass these command line arguments for them to be accessible. Now that they are parameters, this isn't the case.

It also might be useful to define a usage string to tell the user what arguments are to be expected:

USAGE = """Downloads files from imgur.com using
a search_term to a specified download_folder.

Usage:
    python myscript.py "my search term" "/path/to/folder"
"""

~snip~
if __name__ == "__main__":
    try:
        search_term, download_folder, *_ = sys.argv[1:]
    except ValueError:
        print(USAGE)
    else:    
        # pass them as args to main
        main(search_term, download_folder)

I'd also look into the argparse library

try/except => print

This pattern can result in NameErrors, which for a user will be baffling. Raise any unhandled exceptions:

    try:
        cookie_elem = driver.find_element(By. XPATH, "/html/body/div[1]/div/div/div/div[2]/div/button[2]")
        search_elem = driver.find_element(By.CLASS_NAME, "Searchbar-textInput")
    # for a bare except, always raise
    except:
        print('Was not able to find an element with that name.')
        raise

The best thing to do is to figure out what exceptions will be raised and if you want to handle them. Catch them explicitly. Otherwise just a print won't do anything and you'll get search_elem not defined, which is true, but it masks the real error.