rzr/rzr

#!/usr/bin/env python

from argparse import RawTextHelpFormatter
from colour import Color
from openrazer.client import DeviceManager
from pathlib import Path

import argparse
import os
import time
import toml
import sys


def error(msg, e=True):
    """Print error to stderr and exit eventually."""
    print("ERROR: {}".format(msg), file=sys.stderr)
    if e:
        exit(1)


def get_color_tuple(color_string):
    """Convert a color string to a RGB tuple in range [0, 255]."""
    # Convert color string to RGB tuple
    try:
        color = Color(color_string).rgb
    except:

        # Try again with leading #
        try:
            color = Color("#{}".format(color_string)).rgb
        except:
            raise Exception("'{}' is not a valid color".format(color_string))

    # Scale RGB tuple from [0, 1] to [0, 255]
    color_tuple = tuple(map(lambda x: int(x * 255), color))

    return color_tuple


def apply_device_profile(device_profile):
    """Apply a profile section on a device."""
    global device_manager, lightmap_directory

    # Get openrazer device
    device = next(
        (
            d
            for d in device_manager.devices
            if d.name.lower().strip() == device_profile["name"].lower().strip()
        ),
        None,
    )
    if not device:
        error("device '{}' not available".format(device_profile["name"]))
    elif not device.has("lighting_led_matrix"):
        error("device '{}' not supported".format(device_profile["name"]))

    # Open lightmap
    try:
        lightmap = toml.load(
            "{}/{}.toml".format(lightmap_directory, device_profile["lightmap"])
        )
    except FileNotFoundError:
        error(
            "the lightmap '{}' for device '{}' doesn't exist".format(
                device_profile["lightmap"], device_profile["name"]
            ),
            False,
        )
        list_lightmaps()
        exit(1)
    except Exception as e:
        error(
            "failed to load lightmap '{}' for device '{}': {}".format(
                device_profile["lightmap"], device_profile["name"], e
            )
        )

    # Set light colors
    try:

        # If lights are declared apply profile, else turn lights off
        if "lights" in device_profile:
            for light in device_profile["lights"]:
                color_tuple = get_color_tuple((device_profile["lights"][light]))
                device.fx.advanced.matrix[lightmap[light][0], lightmap[light][1]] = (
                    color_tuple
                )
        else:
            device.fx.none()

        # Apply light colors
        device.fx.advanced.draw()

    except KeyError:
        error(
            "light '{}' is not available in lightmap '{}' for device '{}'".format(
                light, device_profile["lightmap"], device_profile["name"]
            )
        )
    except Exception as e:
        error(
            "failed to set light '{}' for device '{}': {}".format(
                light, device_profile["name"], e
            )
        )


def apply_profile(name):
    """Apply a profile by name."""
    global profile_directory

    try:
        profile = toml.load("{}/{}.toml".format(profile_directory, name))
    except FileNotFoundError:
        error("the profile '{}' doesn't exist".format(name), False)
        list_profiles()
        exit(1)
    except Exception as e:
        error("couldn't load profile '{}': {}".format(name, e))
    for device in profile:

        # Check if mandatory attributes "name" and "lightmap" exist
        if "name" not in profile[device]:
            error("'name' attribute is missing for device '{}'".format(device))
        if "lightmap" not in profile[device]:
            error("'lightmap' attribute is missing for device '{}'".format(device))

        apply_device_profile(profile[device])

    print("Profile '{}' applied".format(name))


def iterate_lights():
    """Iterate through every matrix position of every device and light one key after another."""
    global device_manager

    try:

        # Turn all lights off
        for device in device_manager.devices:
            device.fx.none()
            device.fx.advanced.draw()

        # Iterate through all devices
        for device in device_manager.devices:

            # Wait five seconds
            for i in range(5, 0, -1):
                print("{} will be iterated in {} seconds".format(device.name, i))
                time.sleep(1)

            # Turn on one light every second
            for i in range(device.fx.advanced.rows):
                for j in range(device.fx.advanced.cols):
                    device.fx.advanced.matrix[i, j] = (255, 255, 255)
                    device.fx.advanced.draw()
                    print("{}: [{}, {}]".format(device.name, i, j))
                    time.sleep(1)

    except Exception as e:
        error("failed to iterate device '{}': {}".format(device.name, e))


def list_devices():
    """Print a list of all available and supported devices."""
    print("The following devices are available and supported:")
    for device in device_manager.devices:
        if device.has("lighting_led_matrix"):
            print("  - {}".format(device.name))


def list_lightmaps():
    """Print a list of all available lightmaps."""
    global lightmap_directory

    if len(os.listdir(lightmap_directory)) > 0:
        print("Available lightmaps:")
        for lightmap_file in sorted(os.listdir(lightmap_directory)):
            print("  - {}".format(lightmap_file[:-5]))
    else:
        print("No lightmaps available")


def list_profiles():
    """Print a list of all available profiles."""
    global profile_directory

    if len(os.listdir(profile_directory)) > 0:
        print("Available profiles:")
        for profile_file in sorted(os.listdir(profile_directory)):
            print("  - {}".format(profile_file[:-5]))
    else:
        print("No profiles available")


def main():
    global device_manager, lightmap_directory, profile_directory

    # Parse arguments
    parser = argparse.ArgumentParser(
        description="A simple command line frontend for OpenRazer.",
        formatter_class=RawTextHelpFormatter,
    )
    parser.add_argument(
        "command",
        metavar="COMMAND",
        nargs="?",
        default=None,
        help="one of the following:\n  list-devices   - list available devices\n  list-lightmaps - list available lightmaps\n  list-profiles  - list available profiles\n  iterate-lights - iterate though all lights of all devices\n  <PROFILE>      - apply the given profile",
    )
    parser.add_argument(
        "-ld",
        "--lightmap-directory",
        default="{}/.config/rzr/lightmaps".format(Path.home()),
        help="path to directory with lightmaps (default: ~/.config/rzr/lightmaps)",
    )
    parser.add_argument(
        "-pd",
        "--profile-directory",
        default="{}/.config/rzr/profiles".format(Path.home()),
        help="path to directory with profiles (default: ~/.config/rzr/profiles)",
    )
    args = parser.parse_args()

    # Create device manager
    try:
        device_manager = DeviceManager()
    except Exception as e:
        error("failed to load device manager: {}".format(e), False)
        print("Is the openrazer-daemon running?")
        exit(1)

    # Check if directories exist
    directories_available = True
    lightmap_directory = args.lightmap_directory
    profile_directory = args.profile_directory

    if not os.path.exists(lightmap_directory):
        error("lightmap directory '{}' doesn't exist".format(lightmap_directory), False)
        create = input("Create the directory? [y/N] ")
        if create in ["y", "Y"]:
            Path(lightmap_directory).mkdir(parents=True, exist_ok=True)
        else:
            directories_available = False
    if not os.path.exists(profile_directory):
        error("profile directory '{}' doesn't exist".format(profile_directory), False)
        create = input("Create the directory? [y/N] ")
        if create in ["y", "Y"]:
            Path(profile_directory).mkdir(parents=True, exist_ok=True)
        else:
            directories_available = False
    if not directories_available:
        exit(1)

    if not args.command:

        # Apply 'default' profile if existent, else print help
        if "default.toml" in os.listdir(profile_directory):
            apply_profile("default")
        else:
            parser.print_help()

    # Execute command
    elif args.command == "list-devices":
        list_devices()
    elif args.command == "list-lightmaps":
        list_lightmaps()
    elif args.command == "list-profiles":
        list_profiles()
    elif args.command == "iterate-lights":
        iterate_lights()
    elif args.command:
        apply_profile(args.command)


if __name__ == "__main__":
    main()