polychromatic-cli

#!/usr/bin/python3
#
# Polychromatic is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Polychromatic is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Polychromatic. If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2017-2024 Luke Horwell <code@horwell.me>
#

"""
Control Razer devices from the command line. Useful for commamd line users or bash scripting.
"""

import argparse
import os
import signal
import sys

import colorama

import polychromatic.base as base_class
import polychromatic.common as common
import polychromatic.effects as effects
import polychromatic.procpid as procpid
from polychromatic.backends._backend import Backend as Backend

VERSION = "0.9.0"

########################################
# Set up variables
########################################
# TODO: Refactor later
base = base_class.PolychromaticBase()
base.init_base(__file__, sys.argv)
dbg = base.dbg
_ = base._
signal.signal(signal.SIGINT, signal.SIG_DFL)
verbose = False

########################################
# Formatting
########################################
CHECKMARK = "✔"

def print_columns(data, print_header=False):
    """
    Prints a basic 'pretty' table ensuring each column has enough room.

    Params:
        data = [
            ["Row 1 Column 1", "Row 1 Column 2", "Row 1 Column 3"]
            ["Row 2 Column 1", "Row 2 Column 2", "Row 2 Column 3"]
        ]
    """
    # Basic mode
    if args.no_pretty_column:
        for row in data:
            print("      ".join(row))
        return

    # Pretty mode
    total_cols = 0
    col_widths = {}
    col_pos = {}

    def _back_to_start():
        # Move cursor to beginning of line
        for i in range(0, total_cols * 10):
            print(colorama.Cursor.BACK(), end="")

    def _jump_to_pos(x):
        # Move cursor to specified position
        _back_to_start()
        for i in range(0, x):
            print(colorama.Cursor.FORWARD(), end="")

    # Calculate the columns dimensions before printing
    for row in data:
        for index, col in enumerate(row):
            line = col

            # Strip formatting characters from calculation
            for char in [dbg.error, dbg.success, dbg.grey, dbg.normal]:
                line = line.replace(char, "")

            # Make note of largest width and total columns.
            try:
                if len(line) > col_widths[index]:
                    col_widths[index] = len(line)
            except KeyError:
                col_widths[index] = len(line)
            total_cols = total_cols + col_widths[index]

    for index, col in enumerate(row):
        # Precalculate the column positions
        col_pos[index] = 0
        for i in range(0, index):
            col_pos[index] += col_widths[i] + 2

    # Print the actual data
    for index, row in enumerate(data):
        # Print header background
        if print_header and index == 0:
            print(colorama.Back.WHITE + colorama.Fore.BLACK, end="")

        # Print aligned columns
        for index, col in enumerate(row):
            _jump_to_pos(col_pos[index])
            print(col,end="")
        print(colorama.Style.RESET_ALL)


########################################
# Parse arguments
########################################
parser = argparse.ArgumentParser(add_help=False)
parser._optionals.title = _("These arguments can be specified")

# Select a device
parser.add_argument("-d", "--device", action="store", choices=common.FORM_FACTORS)
parser.add_argument("-n", "--name", action="store")
parser.add_argument("-s", "--serial", action="store")

# Output details only
parser.add_argument("-l", "--list-devices", action="store_true")
parser.add_argument("-k", "--list-options", action="store_true")

# Device manipulation
parser.add_argument("-z", "--zone", action="store")
parser.add_argument("-o", "--option", action="store")
parser.add_argument("-p", "--parameter", action="store")
parser.add_argument("-c", "--colours", action="store")
parser.add_argument("-e", "--effect", action="store")

# Special handling
parser.add_argument("--dpi", action="store")

# Misc
parser.add_argument("--version", action="store_true")
parser.add_argument("-h", "--help", action="store_true")
parser.add_argument("--locale", action="store")
parser.add_argument("--no-pretty-column", action="store_true")
parser.add_argument("-v", "--verbose", action="store_true")

args = parser.parse_args()

if args.locale:
    base.reinit_locales(args.locale)
    _ = base._

if not len(sys.argv) > 1:
    dbg.stdout(_("No arguments passed."), dbg.error)
    dbg.stdout(_("Type polychromatic-cli --help to see possible combinations."))
    sys.exit(0)

if args.version:
    app_version, git_commit, py_version = common.get_versions(VERSION)
    print("Polychromatic", app_version)
    if git_commit:
        print("Commit:", git_commit)
    print("Python:", py_version)
    sys.exit(0)

if args.verbose:
    dbg.verbose_level = 1
    verbose = True

if args.help:
    rows = [
        [
            dbg.action + "  -d, --device" + dbg.warning + " {id}",
            dbg.normal + _("Select the device by form factor. If omitted, use all devices")
        ],
        [
            "",
            dbg.action + "[{0}]".format(", ".join(common.FORM_FACTORS))
        ],
        [
            dbg.action + "  -n, --name" + dbg.warning + " Razer BlackWidow",
            dbg.normal + _("Select a device by name")
        ],
        [
            dbg.action + "  -s, --serial" + dbg.warning + " XX1234567890123",
            dbg.normal + _("Select a device by serial number")
        ],
        ["", ""],
        [
            dbg.action + "  -l, --list-devices",
            dbg.normal + _("List connected device(s) and their zone(s)")
        ],
        [
            dbg.action + "  -k, --list-options",
            dbg.normal + _("List all possible combinations for selected device(s)")
        ],
        ["", ""],
        [
            dbg.action + "  -z, --zone" + dbg.warning + " main",
            dbg.normal + _("Make change to specified zone(s). Comma separated. If omitted, use all available zones")
        ],
        [
            dbg.action + "  -o, --option" + dbg.warning + " brightness",
            dbg.normal + _("Set option on the device")
        ],
        [
            dbg.action + "  -p, --parameter" + dbg.warning + " 50",
            dbg.normal + _("Set parameter for the option (if applicable)")
        ],
        [
            dbg.action + "  -c, --colours" + dbg.warning + " \"#RRGGBB,#RRGGBB\"",
            dbg.normal + _("Set colours in hex format (#RRGGBB). Comma separated")
        ],
        [
            dbg.action + "  -e, --effect" + dbg.warning + " <name/path>",
            dbg.normal + _("Set custom (software) effect by name or path")
        ],
        [
            "",
            dbg.warning + _("To set hardware effects, use --option (-o) instead")
        ],
        [
            dbg.action + "  --dpi" + dbg.warning + " <X>[,Y]",
            dbg.normal + _("Set dots per inch.")
        ],
        ["", ""],
        [
            dbg.action + "  --version",
            dbg.normal + _("Print program version and exit")
        ],
        [
            dbg.action + "  -h, --help",
            dbg.normal + _("Show this help message and exit")
        ],
        [
            dbg.action + "  --no-pretty-column",
            dbg.normal + _("Do not print with pretty columns")
        ],
        [
            dbg.action + "  --locale" + dbg.warning + " \"" + base.i18n.get_current_locale() + "\"",
            dbg.normal + _("Force a specific language, e.g. de_DE")
        ],
        [
            dbg.action + "  -v, --verbose",
            dbg.normal + _("Output additional data and debug messages")
        ],
        ["",""]
    ]

    dbg.stdout(_("Available Options"))
    dbg.stdout("---------------------------")
    print_columns(rows)
    sys.exit(0)


########################################
# Select devices
########################################
if verbose:
    dbg.stdout("Loading backends...", dbg.action)
base.middleman.init()

if verbose:
    dbg.stdout("Loaded {0} backend(s):".format(len(base.middleman.backends)), dbg.success)
    for backend in base.middleman.backends:
        dbg.stdout(" - {0}".format(backend.name), dbg.success)

    if len(base.middleman.not_installed) > 0:
        dbg.stdout("Not installed: " + ", ".join(base.middleman.not_installed), dbg.debug)

    if len(base.middleman.bad_init) > 0:
        dbg.stdout("Unsuccessful:", dbg.warning)
        for backend in base.middleman.bad_init:
            dbg.stdout(" - {0}".format(backend.backend_id), dbg.warning)

    if len(base.middleman.import_errors) > 0:
        dbg.stdout("Errors:", dbg.warning)
        for error in base.middleman.import_errors:
            dbg.stdout(" - {0}".format(error), dbg.warning)

if len(base.middleman.backends) == 0:
    dbg.stdout(_("No backends available."), dbg.error)
    sys.exit(1)

device_list = []
unsupported_list = []

if args.name:
    device = base.middleman.get_device_by_name(str(args.name))
    if device:
        device_list = [device]
    else:
        dbg.stdout("{0} {1}".format(_("Could not find a device with this name:"), args.name), dbg.error)
        sys.exit(1)

elif args.serial:
    device= base.middleman.get_device_by_serial(str(args.serial))
    if device:
        device_list = [device]
    else:
        dbg.stdout("{0} {1}".format(_("Could not find a device with this serial:"), args.serial), dbg.error)
        sys.exit(1)

elif args.device:
    device_list = base.middleman.get_devices_by_form_factor(args.device)
    if not device_list:
        dbg.stdout("{0} {1}".format(_("No devices found that match this form factor:"), args.device), dbg.error)
        dbg.stdout("{0} {1}".format(_("Avaliable form factors:"), ", ".join(common.FORM_FACTORS)), dbg.error)
        sys.exit(1)

else:
    device_list = base.middleman.get_devices()
    unsupported_list = base.middleman.get_unsupported_devices()

device_list = sorted(device_list, key=lambda device: device.name)

if len(device_list) == 0:
    dbg.stdout(_("No devices connected."), dbg.error)
    sys.exit(1)

########################################
# List devices, zones or current status.
########################################
if args.list_devices:
    dbg.stdout("Connected Devices:", dbg.success)
    table = [[
        _("Name") + " (-n)",
        _("Form Factor") + " (-d)",
        _("Serial") + " (-s)",
        _("Zones") + " (-z)"
    ]]

    for device in device_list:
        table.append([
            "{0}{1}".format(dbg.normal, device.name),
            "{0}{1}".format(dbg.action, device.form_factor["id"]),
            "{0}{1}".format(dbg.warning, device.serial),
            "{0}{1}".format(dbg.magenta, ", ".join(map(lambda zone: zone.zone_id, device.zones)))
        ])

    for device in unsupported_list:
        table.append([
            "{0}{1} {2}".format(dbg.error, _("Unrecognised:"), device.name),
            "",
            "",
            ""
        ])

    print_columns(table, True)

if args.list_options:
    for device in device_list:
        print("")
        rows = []

        dbg.stdout(device.name, dbg.action)
        dbg.stdout("".join("–" for c in list(device.name)), dbg.action)

        # Print device information
        def _add_row(label, value, colour):
            rows.append([colour + label, colour + value])

        if device.vid and device.pid:
            _add_row("VID:PID", "{0}:{1}".format(device.vid, device.pid), dbg.normal)

        if device.serial:
            _add_row(_("Serial"), device.serial, dbg.normal)

        if device.firmware_version:
            _add_row(_("Firmware Version"), device.firmware_version, dbg.normal)

        if device.keyboard_layout:
            _add_row(_("Keyboard Layout"), device.keyboard_layout, dbg.normal)

        if device.matrix:
            row_suffix = _("row") if str(device.matrix.rows)[-1] == 1 else _("rows")
            col_suffix = _("column") if str(device.matrix.cols)[-1] == 1 else _("columns")
            dimensions = "{0} {1}, {2} {3}".format(device.matrix.rows, row_suffix, device.matrix.cols, col_suffix)
        else:
            _add_row(_("Custom Effects"), _("Unsupported"), dbg.error)

        if device.matrix and not device.monochromatic:
            _add_row(_("Custom Effects"), "{0} ({1})".format(_("Yes"), dimensions), dbg.success)

        if device.matrix and device.monochromatic:
            _add_row(_("Custom Effects"), "{0} ({1})".format(_("Yes, Limited RGB"), dimensions), dbg.success)

        if device.battery:
            battery_label = f"{str(device.battery.percentage)}%"
            if device.battery.percentage <= 15:
                battery_colour = dbg.error
            elif device.battery.percentage >= 50:
                battery_colour = dbg.success
            else:
                battery_colour = dbg.warning
            if device.battery.is_charging:
                battery_colour = dbg.success
                battery_label += " (charging)"
            _add_row(_("Battery"), battery_label, battery_colour)

        # Space between details and next table
        _add_row("", "", dbg.normal)
        print_columns(rows)
        # Print DPI range
        if device.dpi:
            dpi = device.dpi
            dbg.stdout(_("Supports DPI (--dpi) between 1 and 2").replace("1", str(dpi.min)).replace("2", str(dpi.max)) + "\n", dbg.success)

        # Print options and their parameter/colour inputs
        rows = [[
            _("Zone") + " (-z)",
            _("Option") + " (-o)",
            _("Parameters") + " (-p)",
            _("Colours") + " (-c)",
        ]]

        def _add_combination(zone, option, params):
            colours = str(option.colours_required) if option.colours_required > 0 else ""
            rows.append([zone.zone_id, option.uid, params, colours])

        device.refresh()
        for zone in device.zones:
            for option in zone.options:
                column_params = ""
                column_colours = ""

                if option.parameters:
                    params = ", ".join(str(param.data) for param in option.parameters)
                    _add_combination(zone, option, params)
                    continue

                if isinstance(option, Backend.ToggleOption):
                    _add_combination(zone, option, "0, 1")

                elif isinstance(option, Backend.SliderOption):
                    _add_combination(zone, option, "{0} - {1} {2}".format(str(option.min), str(option.max), option.suffix_plural))

                else:
                    _add_combination(zone, option, "")

        print_columns(rows, True)

# After using a --list-* parameter, exit here as devices won't be manipulated.
if args.list_devices or args.list_options:
    sys.exit()


########################################
# Set DPI
########################################
if args.dpi:
    for device in device_list:
        if device.form_factor["id"] != "mouse":
            continue

        if not device.dpi:
            dbg.stdout("{0}: {1}".format(device.name, _("DPI is not supported. Or, it may need to be set using -o/--option.")), dbg.error)

        dpi_values = args.dpi.split(",")

        for number in dpi_values:
            if int(number) > device.dpi.max:
                dbg.stdout(f"{device.name}: {_('DPI too high:')} {number} ({_('Maximum')}: {str(device.dpi.max)})", dbg.error)
                sys.exit(1)

            if int(number) < device.dpi.min:
                dbg.stdout(f"{device.name}: {_('DPI too low:')} {number} ({_('Minimum')}: {str(device.dpi.min)})", dbg.error)
                sys.exit(1)

        if verbose:
            dbg.stdout(f"{device.name}: Setting DPI to {','.join(dpi_values)}", dbg.success)

        if len(dpi_values) == 1:
            x = y = dpi_values[0]
        else:
            x = dpi_values[0]
            y = dpi_values[1]

        try:
            device.dpi.set(x, y)
        except Exception as e:
            dbg.stdout(f"{device.name}: {_('Failed to set DPI due to an error:')}\n{common.get_exception_as_string(e)}", dbg.error)

    sys.exit(0)


########################################
# Set custom software effect
########################################
if args.effect:
    effectman = effects.EffectFileManagement()
    requested_path = None

    # User specifies absolute path
    if os.path.isfile(args.effect):
        requested_path = os.path.abspath(args.effect)

    # User specifies effect by name
    else:
        requested_name = args.effect
        file_list = effectman.get_item_list()
        for effect in file_list:
            if effect["name"] == requested_name:
                requested_path = effect["path"]
                break

    if not requested_path:
        dbg.stdout(_("Unable to locate the effect '[]'").replace("[]", args.effect), dbg.error)
        sys.exit(1)

    data = effectman.get_item(requested_path)
    if type(data) == int:
        # Error details already printed by the get_item() function
        sys.exit(1)

    procmgr = procpid.ProcessManager("helper")
    procmgr.start_component(["--run-fx", requested_path, "--device-name", data["map_device"]])
    if verbose:
        dbg.stdout("Successfully executed request.", dbg.success)
    sys.exit(0)


########################################
# Set an option
########################################
apply_zone = args.zone
apply_option = args.option
apply_param = args.parameter
apply_colours = []

# Validate colour input (hash at the beginning, 6 length)
if args.colours:
    valid = True
    for colour in args.colours.split(","):
        if colour[0] != "#":
            # Be lenient with missing hashes
            colour = "#" + colour
        if len(colour) == 4:
            valid = False
            dbg.stdout(_("3 byte hexadecimal values are unsupported.").replace("[]", colour), dbg.warning)
        elif len(colour) != 7:
            valid = False
            dbg.stdout(_("Colour [] is not a valid hex value.").replace("[]", colour), dbg.warning)
        apply_colours.append(colour)

    if not valid:
        sys.exit(1)

if not apply_option:
    dbg.stdout(_("Please specify --option (-o). Use --list-options (-k) to view available combinations."), dbg.error)
    sys.exit(1)


def print_device_error(device, zone, reason, item):
    zone_label = ""
    if len(device.zones) > 1:
        zone_label = " ({0})".format(zone.zone_id)
    dbg.stdout("{0}{1}: {2} '{3}'".format(device.name, zone_label, reason, str(item)), dbg.error)


def _get_zone_object(device, zone_id):
    for zone in device.zones:
        if zone.zone_id == zone_id:
            return zone
    return None


def _get_option_object(zone, option_id):
    for option in zone.options:
        if str(option.uid) == option_id:
            return option
    return None


def _get_parameter_object(option, param_id):
    for param in option.parameters:
        if str(param.data) == param_id:
            return param
    return None


def _apply_option(device, option, parameter, colours):
    base.middleman.stop_software_effect(device.serial)
    try:
        option.colours = colours
        option.apply(parameter)
    except Exception as e:
        print_device_error(device, zone, _("Failed to set option"), option.uid)
        dbg.stdout(common.get_exception_as_string(e), dbg.error)
        return False

    if verbose:
        dbg.stdout("{0}: OK".format(device.name), dbg.success)
    return True


def _apply_device_options(device, zone_id):
    """
    Returns a boolean indicating success or failure. This will be used later for
    the exit code.
    """
    zone = _get_zone_object(device, zone_id)

    if not zone:
        print_device_error(device, zone, _("No such zone"), zone_id)
        return False

    option = _get_option_object(zone, apply_option)

    if not option:
        print_device_error(device, zone, _("No such option"), apply_option)
        return False

    # Validate option has enough colours
    if option.colours_required:
        if not apply_colours:
            print_device_error(device, zone, _("Missing --colours (-c) for option"), apply_option)
            return False

        if len(apply_colours) > option.colours_required:
            dbg.stdout(_("Too many colours specified. Ignoring excess."), dbg.warning)

        elif option.colours_required < len(apply_colours):
            print_device_error(device, zone, _("Insufficient --colours (-c) for option"), apply_option)
            dbg.stdout(_("Required colours:") + " " + str(option.colours_required), dbg.warning)
            return False

    # Validate parameters for effects and multiple choice are present
    param = _get_parameter_object(option, apply_param)

    if option.parameters:
        param_ids = list(map(lambda param: str(param.data), option.parameters))

        if not apply_param:
            print_device_error(device, zone, _("Missing --parameter (-p) for option"), apply_option)
            dbg.stdout(_("Available parameters:") + " " + ", ".join(param_ids), dbg.warning)
            return False

        if not apply_param in param_ids:
            print_device_error(device, zone, _("Unsupported parameter"), apply_param)
            dbg.stdout(_("Available parameters:") + " " + ", ".join(param_ids), dbg.warning)
            return False

        if param.colours_required:
            if not apply_colours:
                print_device_error(device, zone, _("Missing --colours (-c) for parameter"), apply_param)
                return False

            if len(apply_colours) > param.colours_required:
                dbg.stdout(_("Too many colours specified. Ignoring excess."), dbg.warning)

            elif param.colours_required < len(apply_colours):
                print_device_error(device, zone, _("Insufficient --colours (-c) for parameter"), apply_param)
                dbg.stdout(_("Required colours:") + " " + str(param.colours_required), dbg.warning)
                return False

        return _apply_option(device, option, param.data, apply_colours)

    # Validate parameters for other types of options
    if isinstance(option, Backend.ToggleOption):
        if not apply_param:
            print_device_error(device, zone, _("Missing --parameter (-p) for option"), apply_option)
            dbg.stdout(_("Available parameters:") + " 0, off, 1, on", dbg.warning)
            return False

        if apply_param in ["on", "1"]:
            return _apply_option(device, option, True, apply_colours)
        elif apply_param in ["off", "0"]:
            return _apply_option(device, option, False, apply_colours)

        print_device_error(device, zone, _("Invalid parameter"), apply_param)
        dbg.stdout(_("Available parameters:") + " 0, off, 1, on", dbg.warning)
        return False

    elif isinstance(option, Backend.SliderOption):
        if not apply_param:
            print_device_error(device, zone, _("Missing --parameter (-p) for option"), apply_option)
            dbg.stdout(_("Available parameters:") + " {0} - {1}".format(str(option.min), str(option.max)), dbg.warning)
            return False

        return _apply_option(device, option, int(apply_param), apply_colours)

    # Option does not take parameters
    return _apply_option(device, option, None, apply_colours)


results = []
for device in device_list:
    # User specifies the zone(s)
    if apply_zone:
        results.append(_apply_device_options(device, apply_zone))

    # Or, loop through them all
    else:
        for zone in device.zones:
            results.append(_apply_device_options(device, zone.zone_id))

if False in results:
    sys.exit(1)

sys.exit(0)