pywal/wal

#!/usr/bin/env python
"""
wal - Generate and change colorschemes on the fly.
Created by Dylan Araps
"""
import argparse
import glob
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/"


class ColorFormats(object):  # pylint: disable=too-few-public-methods
    """Store colors in various formats."""
    x_colors = []
    sequences = []


# 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():
        images = []
        file_types = ('*.png', '*.jpg', '*.jpeg', '*.jpe', '*.gif')
        for files in file_types:
            images.extend(glob.glob(str(image) + "/" + files))

        rand_img = random.choice(images)
        rand_img = pathlib.Path(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."""
    ColorFormats.sequences.append("\\033]%s;%s\\007" % (str(index), color))

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

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

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


def set_color(index, color):
    """Build the escape sequence we need for each color."""
    ColorFormats.x_colors.append("*.color%s: %s\n" % (str(index), color))
    ColorFormats.x_colors.append("*color%s: %s\n" % (str(index), color))
    ColorFormats.sequences.append("\\033]4;%s;%s\\007" % (str(index), 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.
    for num, color in enumerate(colors):
        set_color(num, color)

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

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

    # Send the sequences to all open terminals.
    terminals = glob.glob("/dev/pts/[0-9]*")
    terminals.append(CACHE_DIR / "sequences")

    for term in terminals:
        with open(term, 'w') as file:
            file.write(sequences)

    print("colors: Set terminal colors")


# }}}


# WALLPAPER SETTING {{{


def set_wallpaper(img):
    """Set the wallpaper."""
    uname = os.uname

    if 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])

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

    else:
        subprocess.Popen(["gsettings", "set", "org.gnome.desktop.background",
                          "picture-uri", img])

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


# }}}


# EXPORT COLORS {{{


def export_generic(colors, col_format):
    """Export colors to var format."""
    # Loop over the colors and format them.
    colors = [col_format % (num, color) for num, color in enumerate(colors)]
    colors = ''.join(colors)

    return colors


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_shell(colors, export_file):
    """Export colors to shell format."""
    col_format = "color%s='%s'\n"
    colors = export_generic(colors, col_format)
    save_file(colors, export_file)

    print("export: Exported shell colors.")


def export_rofi(colors):
    """Append rofi colors to the x_colors list."""
    ColorFormats.x_colors.append("rofi.color-window: %s, %s, %s\n"
                                 % (colors[0], colors[0], colors[10]))
    ColorFormats.x_colors.append("rofi.color-normal: %s, %s, %s, %s, %s\n"
                                 % (colors[0], colors[15], colors[0],
                                    colors[10], colors[0]))
    ColorFormats.x_colors.append("rofi.color-active: %s, %s, %s, %s, %s\n"
                                 % (colors[0], colors[15], colors[0],
                                    colors[10], colors[0]))
    ColorFormats.x_colors.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."""
    ColorFormats.x_colors.append("emacs*background: %s\n" % (colors[0]))
    ColorFormats.x_colors.append("emacs*foreground: %s\n" % (colors[15]))


def export_xrdb(colors, export_file):
    """Export colors to xrdb."""
    colors = ''.join(colors)
    save_file(colors, export_file)

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

    print("export: Exported xrdb colors.")


def export_css(colors, export_file):
    """Export colors as css variables."""
    col_format = "\t--color%s: %s;\n"
    colors = ":root {\n%s}\n" % str(export_generic(colors, col_format))
    save_file(colors, export_file)

    print("export: Exported css colors.")


def export_scss(colors, export_file):
    """Export colors as scss variables."""
    col_format = "$color%s: %s;\n"
    colors = export_generic(colors, col_format)
    save_file(colors, export_file)

    print("export: Exported scss colors.")


def export_colors(colors):
    """Export colors in various formats."""
    export_plain(colors)
    export_shell(colors, CACHE_DIR / "colors.sh")

    # X based colors.
    export_rofi(colors)
    export_emacs(colors)
    export_xrdb(ColorFormats.x_colors, CACHE_DIR / "xcolors")

    # Web based colors.
    export_css(colors, CACHE_DIR / "colors.css")
    export_scss(colors, CACHE_DIR / "colors.scss")


# }}}


# 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 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)

    return 0


main()