reddwall

#!/usr/bin/env python3
import json
from urllib.parse import urlparse
import urllib.request
import os
import random
import argparse
import dbus

cache = []
ROOT = os.getenv('HOME')
SOURCE_DIR = ROOT + "/Pictures/walls/"
CACHE_DIR = SOURCE_DIR + "/cache/"
WALL_DIR = SOURCE_DIR + "/wallpaper/"


def is_image(http_message):
    """
    Check whether an http response is of an image
    """
    return "image" in http_message.getheader("Content-Type")


def get_charset(http_message):
    """
    Get charset from an http response
    """
    cset = http_message.get_content_charset()
    if len(cset) == 0:
        cset = 'utf-8'
    return cset


def get_image(url):
    """
    Get image from url and return a file object for reading.
    It downloads the image and stores it in the cache directory.
    The image is identified by its filename. If a file with the
    same name exists, It will not download it.
    """
    filename = urlparse(url).path.split('/')[-1]
    if filename in cache:
        return open(CACHE_DIR + filename, "rb"), True
    img_response = urllib.request.urlopen(url)
    if not is_image(img_response):
        return None, False
    if "content-disposition" in img_response.info().keys():
        filename = img_response.info().get("content-disposition")
    fp = open(CACHE_DIR + filename, "wb")
    fp.write(img_response.read())
    fp.close()
    fp = open(CACHE_DIR + filename, "rb")
    return fp, True


def init():
    """
    Populate the entire list of images in cache directory.
    """
    global cache
    os.makedirs(CACHE_DIR, exist_ok=True)
    os.makedirs(WALL_DIR, exist_ok=True)
    cache = os.listdir(CACHE_DIR)


def get_from_cache():
    """
    Fetch image from local cache. Returns a file object.
    """
    print("Getting wallpaper from local cache")
    filename = random.choice(cache)
    image_file = open(CACHE_DIR + filename, "rb")
    return image_file


def get_from_subreddit(subred, limit=20):
    """
    Get an image from a subreddit alongwith options.
    Fetches the list of posts, selects a post randomly and tries to
    download it as image, if not an image then selects another post in random
    Returns file pointer to the downloaded file
    """
    print("Getting wallpaper from: " + subred)
    response = urllib.request.urlopen("http://reddit.com/r/" + subred + ".json?limit=" + str(limit))
    body = response.read()
    headers = response.getheader("Content-Type").split(';')

    # TODO: Handle non json requests
    charset = headers[1].split("=")[1].lower()
    if len(charset) == 0:
        charset = 'utf-8'  # Default
    json_response = body.decode(charset)
    js = json.loads(json_response)
    children = js.get("data").get("children")
    # TODO: Warning Possibility of infinite loop
    image_file = None
    while True:
        child = random.choice(children)
        data = child.get("data")
        print("Trying: " + data.get("title"))
        image_url = data.get("url")
        print("`- " + image_url)
        image_file, isimage = get_image(image_url)
        if isimage:
            print("Getting: " + data.get("title"))
            break
    return image_file

def set_plasma_wallpaper(path):
    """Set image at path as the plasma wallpaper

    :path: Path of the image file
    :returns: nothing

    """
    js = """
    var containers = desktops();
    for(i=0; i<containers.length; i++) {
        c = containers[i];
        c.wallpaperPlugin = "org.kde.image";
        c.currentConfigGroup = Array("Wallpaper", "org.kde.image", "General");
        c.writeConfig("Image", "%s")
    }
    """
    dbus.SessionBus().get_object('org.kde.plasmashell', '/PlasmaShell').evaluateScript(js % path)

def set_plasma_video_wallpaper(path):
    js = """
    var containers = desktops();
    for(i=0; i<containers.length; i++) {
        c = containers[i];
        c.wallpaperPlugin = "org.kde.video";
        c.currentConfigGroup = Array("Wallpaper", "org.kde.video", "General");
        c.writeConfig("Video", "file://%s")
    }
    """
    dbus.SessionBus().get_object('org.kde.plasmashell', '/PlasmaShell').evaluateScript(js % path)

if __name__ == "__main__":
    init()

    parser = argparse.ArgumentParser(description='Set wallpaper from reddit/cache')
    parser.add_argument('subreddit', nargs='?', help="Subreddit to fetch wallpaper from")
    parser.add_argument('--limit', '-l', type=int, default=20,
                        help="Number of posts to scan from a subreddit (default=20)")
    parser.add_argument('--file', '-f', help="Specify a particular file to put as wallpaper")
    parser.add_argument('--video', '-v', help="Specify a particular video to put as wallpaper")
    parser.add_argument('--url', '-u',
                        help="Download a specific file to put as wallpaper")  # TODO: Group -u/-f mutually excl.
    Args = parser.parse_args()
    try:
        # Ideally, there should be at most one file in the `walls` directory,
        oldwall = WALL_DIR + os.listdir(WALL_DIR)[0]
    except IndexError:
        oldwall = ""
    if Args.file is not None:
        img_file = open(Args.file, "rb")
    elif Args.video is not None:
        vd_file = open(Args.video, "rb")
        set_plasma_video_wallpaper(vd_file.name)
        exit()
    elif Args.url is not None:
        img_file, _ = get_image(Args.url)
    elif Args.subreddit is not None:
        subreddit = Args.subreddit
        img_file = get_from_subreddit(subreddit, Args.limit)
    else:
        img_file = get_from_cache()

    set_plasma_wallpaper(img_file.name)