pywal/wal

#!/usr/bin/env python
"""
wal - Generate and change colorschemes on the fly.
Created by Dylan Araps
"""
import argparse
import os
import pathlib
import random
import re
import shutil
import subprocess
import sys


__version__ = "0.1"


# Internal variables.
COLOR_COUNT = 16
CACHE_DIR = pathlib.Path.home() / ".cache/wal/"


# pylint: disable=too-few-public-methods
class ColorType(object):
    """Store colors in various formats."""
    xrdb = []
    sequences = []
    shell = []
    scss = []
    css = [":root {\n"]
    putty = [
        "Windows Registry Editor Version 5.00\n",
        "[HKEY_CURRENT_USER\\Software\\SimonTatham\\PuTTY\\Sessions\\Wal]\n",
    ]


# ARGS {{{


def get_args():
    """Get the script arguments."""
    description = "wal - Generate colorschemes on the fly"
    arg = argparse.ArgumentParser(description=description)

    # Add the args.
    arg.add_argument("-c", action="store_true",
                     help="Delete all cached colorschemes.")

    arg.add_argument("-i", metavar="\"/path/to/img.jpg\"",
                     help="Which image or directory to use.")

    arg.add_argument("-n", action="store_true",
                     help="Skip setting the wallpaper.")

    arg.add_argument("-o", metavar="\"script_name\"",
                     help="External script to run after \"wal\".")

    arg.add_argument("-q", action="store_true",
                     help="Quiet mode, don\"t print anything.")

    arg.add_argument("-r", action="store_true",
                     help="Reload current colorscheme.")

    arg.add_argument("-t", action="store_true",
                     help="Fix artifacts in VTE Terminals. \
                           (Termite, xfce4-terminal)")

    return arg.parse_args()


def process_args(args):
    """Process args"""
    # If no args were passed.
    if not len(sys.argv) > 1:
        print("error: wal needs to be given arguments to run.")
        print("       Refer to \"wal -h\" for more info.")
        exit(1)

    # -q
    if args.q:
        sys.stdout = open("/dev/null", "w")
        sys.stderr = open("/dev/null", "w")

    # -c
    if args.c:
        shutil.rmtree(CACHE_DIR / "schemes")
        create_cache_dir()

    # -r
    if args.r:
        reload_colors(args.t)

    # -i
    if args.i:
        image = str(get_image(args.i))
        colors = get_colors(image)

        # Set the wallpaper.
        if not args.n:
            set_wallpaper(image)

        return colors


# }}}


# COLORSCHEME GENERATION {{{


def get_image(img):
    """Validate image input."""
    image = pathlib.Path(img)

    # Check if the user has Imagemagick installed.
    if not shutil.which("convert"):
        print("error: imagemagick not found, exiting...")
        print("error: wal requires imagemagick to function.")
        exit(1)

    if image.is_file():
        wal_img = image

    # Pick a random image from the directory.
    elif image.is_dir():
        file_types = (".png", ".jpg", ".jpeg", ".jpe", ".gif")
        images = [img for img in os.listdir(image)
                  if img.endswith(file_types)]

        rand_img = random.choice(images)
        rand_img = pathlib.Path(image / rand_img)

        if rand_img.is_file():
            wal_img = rand_img

    else:
        print("error: No valid image file found.")
        exit(1)

    print("image: Using image", wal_img)
    return wal_img


def imagemagick(color_count, img):
    """Call Imagemagick to generate a scheme."""
    colors = subprocess.Popen(["convert", img, "+dither", "-colors",
                               str(color_count), "-unique-colors", "txt:-"],
                              stdout=subprocess.PIPE)

    return colors.stdout.readlines()


def gen_colors(img):
    """Generate a color palette using imagemagick."""
    # Generate initial scheme.
    raw_colors = imagemagick(COLOR_COUNT, img)

    # If imagemagick finds less than 16 colors, use a larger source number
    # of colors.
    index = 0
    while len(raw_colors) - 1 < COLOR_COUNT:
        index += 1
        raw_colors = imagemagick(COLOR_COUNT + index, img)

        print("colors: Imagemagick couldn't generate a", COLOR_COUNT,
              "color palette, trying a larger palette size",
              COLOR_COUNT + index)

    # Remove the first element, which isn"t a color.
    del raw_colors[0]

    # Create a list of hex colors.
    colors = [re.search("#.{6}", str(col)).group(0) for col in raw_colors]
    return colors


def get_colors(img):
    """Generate a colorscheme using imagemagick."""
    # Cache file.
    cache_file = CACHE_DIR / "schemes" / img.replace("/", "_")
    cache_file = pathlib.Path(cache_file)

    # Cache the wallpaper name.
    with open(CACHE_DIR / "wal", "w") as file:
        file.write("%s\n" % (img))

    if cache_file.is_file():
        colors = read_colors(cache_file)

    else:
        print("colors: Generating a colorscheme...")

        # Generate the colors.
        colors = gen_colors(img)
        colors = sort_colors(colors)

        # Cache the colorscheme.
        with open(cache_file, "w") as file:
            file.write("\n".join(colors))

    print("colors: Generated colorscheme")
    return colors


def sort_colors(colors):
    """Sort the generated colors."""
    sorted_colors = []
    sorted_colors.append(colors[0])
    sorted_colors.append(colors[9])
    sorted_colors.append(colors[10])
    sorted_colors.append(colors[11])
    sorted_colors.append(colors[12])
    sorted_colors.append(colors[13])
    sorted_colors.append(colors[14])
    sorted_colors.append(colors[15])
    sorted_colors.append(set_grey(colors))
    sorted_colors.append(colors[9])
    sorted_colors.append(colors[10])
    sorted_colors.append(colors[11])
    sorted_colors.append(colors[12])
    sorted_colors.append(colors[13])
    sorted_colors.append(colors[14])
    sorted_colors.append(colors[15])
    return sorted_colors


# }}}


# SEND SEQUENCES {{{


def set_special(index, color):
    """Build the escape sequence for special colors."""
    ColorType.sequences.append("\\033]%s;%s\\007" % (str(index), color))

    if index == 10:
        ColorType.xrdb.append("URxvt*foreground: %s\n" % (color))
        ColorType.xrdb.append("XTerm*foreground: %s\n" % (color))

    elif index == 11:
        ColorType.xrdb.append("URxvt*background: %s\n" % (color))
        ColorType.xrdb.append("XTerm*background: %s\n" % (color))

    elif index == 12:
        ColorType.xrdb.append("URxvt*cursorColor: %s\n" % (color))
        ColorType.xrdb.append("XTerm*cursorColor: %s\n" % (color))


def set_color(index, color):
    """Build the escape sequence we need for each color."""
    index = str(index)
    ColorType.xrdb.append("*.color%s: %s\n" % (index, color))
    ColorType.xrdb.append("*color%s: %s\n" % (index, color))
    ColorType.sequences.append("\\033]4;%s;%s\\007" % (index, color))
    ColorType.shell.append("color%s='%s'\n" % (index, color))
    ColorType.css.append("\t--color%s: %s;\n" % (index, color))
    ColorType.scss.append("$color%s: %s;\n" % (index, color))
    ColorType.putty.append("\"Colour%s\"=\"%s\"\n"
                           % (index, hex_to_rgb(color)))


def set_grey(colors):
    """Set a grey color based on brightness of color0"""
    return {
        0: "#666666",
        1: "#666666",
        2: "#757575",
        3: "#999999",
        4: "#999999",
        5: "#8a8a8a",
        6: "#a1a1a1",
        7: "#a1a1a1",
        8: "#a1a1a1",
        9: "#a1a1a1",
    }.get(int(colors[0][1]), colors[7])


def send_sequences(colors, vte):
    """Send colors to all open terminals."""
    set_special(10, colors[15])
    set_special(11, colors[0])
    set_special(12, colors[15])
    set_special(13, colors[15])
    set_special(14, colors[0])

    # This escape sequence doesn"t work in VTE terminals.
    if not vte:
        set_special(708, colors[0])

    # Create the sequences.
    # pylint: disable=W0106
    [set_color(num, color) for num, color in enumerate(colors)]

    # Set a blank color that isn"t affected by bold highlighting.
    set_color(66, colors[0])

    # Decode the string.
    sequences = "".join(ColorType.sequences)
    sequences = bytes(sequences, "utf-8").decode("unicode_escape")

    # Get a list of terminals.
    terminals = ["%s%s" % ("/dev/pts/", term)
                 for term in os.listdir("/dev/pts/") if len(term) < 4]
    terminals.append(CACHE_DIR / "sequences")

    # Send the sequences to all open terminals.
    # pylint: disable=W0106
    [save_file(sequences, term) for term in terminals]

    print("colors: Set terminal colors")


# }}}


# WALLPAPER SETTING {{{


def get_desktop_env():
    """Identify the current running desktop environment."""
    desktop = os.getenv("XDG_CURRENT_DESKTOP")
    if desktop:
        return desktop

    desktop = os.getenv("DESKTOP_SESSION")
    if desktop:
        return desktop

    desktop = os.getenv("GNOME_DESKTOP_SESSION_ID")
    if desktop:
        return "GNOME"

    desktop = os.getenv("MATE_DESKTOP_SESSION_ID")
    if desktop:
        return "MATE"


def set_desktop_wallpaper(desktop, img):
    """Set the wallpaper for the desktop environment."""
    if "xfce" in desktop or "xubuntu" in desktop:
        subprocess.Popen(["xfconf-query", "--channel", "xfce4-desktop",
                          "--property",
                          "/backdrop/screen0/monitor0/image-path",
                          "--set", img])

        # XFCE requires two commands since they differ between versions.
        subprocess.Popen(["xfconf-query", "--channel", "xfce4-desktop",
                          "--property",
                          "/backdrop/screen0/monitor0/workspace0/last-image",
                          "--set", img])

    elif "muffin" in desktop or "cinnamon" in desktop:
        subprocess.Popen(["gsettings", "set",
                          "org.cinnamon.desktop.background",
                          "picture-uri", "file:///" + img])

    elif "mate" in desktop:
        subprocess.Popen(["gsettings", "set", "org.mate.background",
                          "picture-filename", img])

    elif "gnome" in desktop:
        subprocess.Popen(["gsettings", "set",
                          "org.gnome.desktop.background",
                          "picture-uri", "file:///" + img])


def set_wallpaper(img):
    """Set the wallpaper."""
    desktop = str(get_desktop_env()).lower()

    if desktop:
        set_desktop_wallpaper(desktop, img)

    else:
        if os.uname == "Darwin":
            subprocess.Popen(["osascript", "-e", "\"tell application \"Finder\" to set \
                               desktop picture to POSIX file\"" + img + "\""])

        elif shutil.which("feh"):
            subprocess.Popen(["feh", "--bg-fill", img])

        elif shutil.which("nitrogen"):
            subprocess.Popen(["nitrogen", "--set-zoom-fill", img])

        elif shutil.which("bgs"):
            subprocess.Popen(["bgs", img])

        elif shutil.which("hsetroot"):
            subprocess.Popen(["hsetroot", "-fill", img])

        elif shutil.which("habak"):
            subprocess.Popen(["habak", "-mS", img])

        else:
            print("error: No wallpaper setter found.")
            return

    print("wallpaper: Set the new wallpaper")
    return 0


# }}}


# EXPORT COLORS {{{


def save_colors(colors, export_file, message):
    """Export colors to var format."""
    colors = "".join(colors)
    save_file(colors, export_file)
    print("export: exported %s." % (message))


def export_plain(colors):
    """Export colors to a plain text file."""
    with open(CACHE_DIR / "colors", "w") as file:
        file.write("\n".join(colors))
    print("export: Exported plain colors")


def export_rofi(colors):
    """Append rofi colors to the x_colors list."""
    ColorType.xrdb.append("rofi.color-window: %s, %s, %s\n"
                          % (colors[0], colors[0], colors[10]))
    ColorType.xrdb.append("rofi.color-normal: %s, %s, %s, %s, %s\n"
                          % (colors[0], colors[15], colors[0],
                             colors[10], colors[0]))
    ColorType.xrdb.append("rofi.color-active: %s, %s, %s, %s, %s\n"
                          % (colors[0], colors[15], colors[0],
                             colors[10], colors[0]))
    ColorType.xrdb.append("rofi.color-urgent: %s, %s, %s, %s, %s\n"
                          % (colors[0], colors[9], colors[0],
                             colors[9], colors[15]))


def export_emacs(colors):
    """Set emacs colors."""
    ColorType.xrdb.append("emacs*background: %s\n" % (colors[0]))
    ColorType.xrdb.append("emacs*foreground: %s\n" % (colors[15]))


def export_xrdb(colors, export_file):
    """Export colors to xrdb."""
    save_colors(colors, export_file, "xrdb colors")

    # Merge the colors into the X db so new terminals use them.
    subprocess.Popen(["xrdb", "-merge", export_file])


def export_colors(colors):
    """Export colors in various formats."""
    export_plain(colors)
    save_colors(ColorType.shell, CACHE_DIR / "colors.sh", "shell variables")

    # X based colors.
    export_rofi(colors)
    export_emacs(colors)
    export_xrdb(ColorType.xrdb, CACHE_DIR / "xcolors")

    # Web based colors.
    ColorType.css.append("}\n")
    save_colors(ColorType.css, CACHE_DIR / "colors.css", "css variables")
    save_colors(ColorType.scss, CACHE_DIR / "colors.scss", "scss variables")

    # Text editor based colors.
    save_colors(ColorType.putty, CACHE_DIR / "colors-putty.reg", "putty theme")


# }}}


# OTHER FUNCTIONS {{{


def read_colors(color_file):
    """Read colors from a file"""
    with open(color_file) as file:
        colors = file.readlines()

    # Strip newlines from each list element.
    colors = [x.strip() for x in colors]
    return colors


def reload_colors(vte):
    """Reload colors."""
    with open(CACHE_DIR / "sequences") as file:
        sequences = file.read()

    # If vte mode was used, remove the problem sequence.
    if vte:
        sequences = re.sub(r"\]708;\#.{6}", "", sequences)

    # Decode the string.
    sequences = bytes(sequences, "utf-8").decode("unicode_escape")

    print(sequences, end="")
    quit()


def save_file(colors, export_file):
    """Write the colors to the file."""
    with open(export_file, "w") as file:
        file.write(colors)


def create_cache_dir():
    """Alias to create the cache dir."""
    pathlib.Path(CACHE_DIR / "schemes").mkdir(parents=True, exist_ok=True)


def hex_to_rgb(color):
    """Convert a hex color to rgb."""
    red, green, blue = list(bytes.fromhex(color.strip("#")))
    return "%s,%s,%s" % (red, green, blue)


# }}}


def main():
    """Main script function."""
    create_cache_dir()

    # Get the args.
    args = get_args()
    colors = process_args(args)

    # Set the colors.
    send_sequences(colors, args.t)
    export_colors(colors)

    # -o
    if args.o:
        subprocess.Popen(["nohup", args.o],
                         stdout=open("/dev/null", "w"),
                         stderr=open("/dev/null", "w"),
                         preexec_fn=os.setpgrp)

    # This saves 10ms.
    # pylint: disable=W0212
    os._exit(0)


main()