qmk

#!/usr/bin/env python3


import argparse
import os
import shutil
import subprocess
import getpass
import sys
import time
from keymaps import keyboard_keymap_paths
from keyboards import keyboard_paths


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


QMK_ROOT = abs_path('qmk_firmware')


class QmkBuilder(object):
    def __init__(self, keyboard, user):
        self.user = user

        self.local_user_dir = abs_path('user')
        self.local_keymap_dir = abs_path('keymaps', keyboard)

        self.keyboard_target = keyboard_keymap_paths[keyboard] if keyboard in keyboard_keymap_paths else keyboard

        self.qmk_keyboards_dir = os.path.join(QMK_ROOT, 'keyboards')
        self.qmk_keymap_dir = os.path.join(self.qmk_keyboards_dir, self.keyboard_target, 'keymaps', user)
        self.qmk_user_dir = os.path.join(QMK_ROOT, 'users', user)

    def __enter__(self):
        self._cleanup()

        for k, v in keyboard_paths.items():
          shutil.copytree(abs_path('keyboards', k), os.path.join(self.qmk_keyboards_dir, v))

        shutil.copytree(self.local_keymap_dir, self.qmk_keymap_dir)
        shutil.copytree(self.local_user_dir, self.qmk_user_dir)
        return self

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

    def _cleanup(self):
        for k, v in keyboard_paths.items():
            if os.path.exists(self.qmk_keymap_dir):
              shutil.rmtree(os.path.join(self.qmk_keyboards_dir, v))

        if os.path.exists(self.qmk_keymap_dir):
            shutil.rmtree(self.qmk_keymap_dir)

        if os.path.exists(self.qmk_user_dir):
            shutil.rmtree(self.qmk_user_dir)

    def execute(self, script):
        subprocess.run(script, cwd=QMK_ROOT, stderr=subprocess.STDOUT, shell=True)

    def build(self):
        self.execute('make {}:{}'.format(self.keyboard_target, self.user))

    def flash(self, programmer = 'flash'):
        self.execute('make {}:{}:{}'.format(self.keyboard_target, self.user, programmer))


def build(keyboard, user, programmer, flash=False):
    with QmkBuilder(keyboard, user) as builder:
        builder.build()

        if flash:
          builder.flash(programmer)


def flash(keyboard, user, programmer):
    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, user) as builder:
        builder.flash(programmer)


def main(argv):
    user = getpass.getuser()

    parser = argparse.ArgumentParser(description='Build util for QMK')
    parser.add_argument('keyboard')
    parser.add_argument('command', choices=['build', 'flash'])
    parser.add_argument('-p', '--programmer', default='flash', help='eg, flash, dfu, avrdude, teensy')
    parser.add_argument('-u', '--user', default=user, help='user (default: {})'.format(user))
    parser.add_argument('-f', '--flash', action='store_true', help='continue to flash after build')
    args = parser.parse_args()

    keyboard = args.keyboard.strip('/')
    if not os.path.isdir(abs_path('keymaps', keyboard)):
        print('Invalid keyboard: {}'.format(keyboard))
        return 1

    try:
        if args.command == 'build':
            build(keyboard, args.user, args.programmer, args.flash)
        elif args.command == 'flash':
            flash(keyboard, args.user, args.programmer)
    except KeyboardInterrupt:
        print()
        print('Aborting.')


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