qmk

#!/usr/bin/env python3


import argparse
import os
import shutil
import subprocess
import sys
import time

import visualizer

user = "issmirnov"

# override for non-standard keyboard locations. For example, keebio
# boards have one extra layer in the source tree.
keyboard_loc_override = {
    "levinson": "keebio/levinson",
    "boardsource_3x4": "boardsource/3x4"
}

def abs_path(*paths):
    dirname = os.path.abspath(os.path.dirname(__file__))
    return os.path.join(dirname, *paths)


class QmkBuilder(object):
    def __init__(self, keyboard, keep):
        self.qmk_root = abs_path("qmk_firmware")

        self.keep_files = keep
        self.local_config_dir = abs_path(keyboard)
        keyboard_target = keyboard_loc_override[keyboard] if keyboard in keyboard_loc_override else keyboard
        self.qmk_config_dir = os.path.join(self.qmk_root, "keyboards",
                                           keyboard_target, "keymaps", user)

        self.local_user_dir = abs_path("user")
        self.qmk_user_dir = os.path.join(self.qmk_root, "users", user)

    def _cleanup(self):
        if self.keep_files:
            return
        if os.path.exists(self.qmk_config_dir):
            shutil.rmtree(self.qmk_config_dir)
        if os.path.exists(self.qmk_user_dir):
            shutil.rmtree(self.qmk_user_dir)

    def __enter__(self):
        self._cleanup()
        shutil.copytree(self.local_config_dir, self.qmk_config_dir)
        shutil.copytree(self.local_user_dir, self.qmk_user_dir)
        return self

    def execute(self, script):
        cmd = os.path.join(self.local_config_dir, script)
        subprocess.check_call(cmd, cwd=self.qmk_root)

    def __exit__(self, *_):
        self._cleanup()

def build(keyboard, keep, push=False):
    with QmkBuilder(keyboard, keep) as builder:
        print("Building {}...".format(keyboard))
        builder.execute("build.sh")


def push(keyboard):
    print("Put your {} into bootlader mode".format(keyboard), end="")
    for _ in range(0, 3):
        print(".", end="")
        sys.stdout.flush()
        time.sleep(1)
    print()

    with QmkBuilder(keyboard, keep=False) as builder:
        print("Updating {}...".format(keyboard))
        builder.execute("push.sh")


def show(keyboard):
    visualizer.visualize(abs_path("user"), abs_path(keyboard))


def main(argv):
    parser = argparse.ArgumentParser(description="Build util for QMK")
    parser.add_argument("command", nargs=1, choices=["build", "push", "show"])
    parser.add_argument("keyboard", nargs=1)
    parser.add_argument("--keep", action='store_true',
                     help='Keep the files in the QMK repo', default=False)
    args = parser.parse_args()

    keyboard = args.keyboard[0].strip("/")
    if not os.path.isdir(abs_path(keyboard)):
        print("Invalid keyboard: {}".format(keyboard))
        return 1

    keep = args.keep
    if keep:
        print("--keep detected, will not delete QMK files")
    try:
        command = args.command[0]
        if command == "build":
            build(keyboard, keep)
        elif command == "push":
            push(keyboard)
        elif command == "show":
            show(keyboard)
    except KeyboardInterrupt:
        print()
        print("Aborting.")


if __name__ == "__main__":
    sys.exit(main(sys.argv))