generate_keymap.py

# Reads the keymap tables from `keymap.md` and generates QMK and ZMK keymaps.
import sys
import os.path
from collections import namedtuple

HEADER = "// This file is automatically generated from README.md\n"

hyp = lambda z, q: (f"LA(LS(LC(LG({z})", "G(S(C(A({q}))))")

# Maps laels from the markdown table to key press codes in ZMK and QMK.
# "LABEL": ("ZMK CODE", "QMK CODE")
KEY_PRESS_CODES = {
    # Special Keys
    "ESC": ("ESC", "KC_ESC"),
    "RET": ("RET", "KC_ENT"),
    "TAB": ("TAB", "KC_TAB"),
    "SPC": ("SPC", "KC_SPC"),
    "BKSP": ("BKSP", "KC_BSPC"),
    "CMD_RET": ("LG(RET)", "S(KC_ENT)"),
    "ALT_BKSP": ("LA(BKSP)", "A(KC_BSPC)"),
    # Text Navigatioh Keys
    "LEFT": ("LEFT", "KC_LEFT"),
    "RIGHT": ("RIGHT", "KC_RIGHT"),
    "UP": ("UP", "KC_UP"),
    "DOWN": ("DOWN", "KC_DOWN"),
    "WORD_L": ("LA(LEFT)", "A(KC_LEFT)"),
    "WORD_R": ("LA(RIGHT)", "A(KC_RIGHT)"),
    "HOME": ("HOME", "KC_HOME"),
    "END": ("END", "KC_END"),
    # UI Navigation Keys
    "SPC_L": ("LC(LEFT)", "C(KC_LEFT)"),
    "SPC_R": ("LC(RIGHT)", "C(KC_RIGHT)"),
    "NXT_WIN": ("LC(F4)", "C(KC_F4)"),
    # Forward / Backwards (GUI + Bracket)
    "FWD": ("LG(LBKT)", "G(KC_LBRC)"),
    "BCK": ("LG(RBKT)", "G(KC_RBRC)"),
    # Prev / Next Tab (GUI + Shift Bracket)
    "TAB_L": ("LG(LS(LBKT))", "G(S(KC_LBRC))"),
    "TAB_R": ("LG(LS(RBKT))", "G(S(KC_RBRC))"),
    # Umlaut key
    "UML": ("LA(U)", "A(KC_U)"),
    # Symbols
    "`": ("GRAVE", "KC_GRAVE"),
    "~": ("LS(GRAVE)", "KC_TILDE"),
    "!": ("LS(N1)", "KC_EXCLAIM"),
    "@": ("LS(N2)", "KC_AT"),
    "#": ("LS(N3)", "KC_HASH"),
    "$": ("LS(N4)", "KC_DOLLAR"),
    "%": ("LS(N5)", "KC_PERCENT"),
    "^": ("LS(N6)", "KC_CIRCUMFLEX"),
    "&": ("LS(N7)", "KC_AMPERSAND"),
    "\*": ("LS(N8)", "KC_ASTERISK"),
    "(": ("LS(N9)", "KC_LEFT_PAREN"),
    ")": ("LS(N0)", "KC_RIGHT_PAREN"),
    "-": ("KP_MINUS", "KC_MINUS"),
    "\_": ("LS(MINUS)", "KC_UNDERSCORE"),
    "=": ("EQUAL", "KC_EQUAL"),
    "+": ("LS(EQUAL)", "KC_PLUS"),
    "[": ("LBKT", "KC_LBRC"),
    "{": ("LS(LBKT)", "KC_LCBR"),
    "]": ("RBKT", "KC_RBRC"),
    "}": ("LS(RBKT)", "KC_RCBR"),
    "\\": ("BSLH", "KC_BSLASH"),
    "PIPE": ("LS(BSLH)", "KC_PIPE"),
    ";": ("SEMI", "KC_SCOLON"),
    ":": ("LS(SEMI)", "KC_COLON"),
    "'": ("SQT", "KC_QUOTE"),
    '"': ("LS(SQT)", "KC_DOUBLE_QUOTE"),
    ",": ("COMMA", "KC_COMMA"),
    "<": ("LS(COMMA)", "KC_LEFT_ANGLE_BRACKET"),
    ".": ("DOT", "KC_DOT"),
    ">": ("LS(DOT)", "KC_RIGHT_ANGLE_BRACKET"),
    "/": ("FSLH", "KC_SLASH"),
    "?": ("LS(FSLH)", "KC_QUESTION"),
    # Modifiers
    "CMD": ("LGUI", "KC_LGUI"),
    "SHFT": ("LSHFT", "KC_LSFT"),
    "CTRL": ("LCTRL", "KC_LCTL"),
    "ALT": ("LALT", "KC_LALT"),
}

# Maps labels from the markdown table to special instructions for ZMK and QMK.
# "LABEL": ("ZMK", "QMK")
SPECIAL_LABELS = {
    # Bluetooth
    "BT_CLR": ("&bt BT_CLR", "KC_NO"),
    "BT_0": ("&bt BT_SEL 0", "KC_NO"),
    "BT_1": ("&bt BT_SEL 1", "KC_NO"),
    "BT_2": ("&bt BT_SEL 2", "KC_NO"),
    "BT_3": ("&bt BT_SEL 3", "KC_NO"),
}

# Maps labels from the markdown table to layer numbers.
LAYER_LABELS = {
    "RSE": 1,
    "LWR": 2,
    "HYP": 3,
    "ADJ": 4,
    "MOU": 5,
}


################################################################################
# Mapping to QMK


def get_qmk_key_press_code(label):
    if label.startswith('CMD_'):
        return f"G({get_qmk_key_press_code(label[4:])})"
    elif label.startswith('HYP\\_'):
        return f"G(S(A(C({get_qmk_key_press_code(label[5:])}))))"
    elif label.startswith('HYP_'):
        return f"G(S(A(C({get_qmk_key_press_code(label[4:])}))))"
    elif len(label) == 1 and (label.isalpha() or label.isdigit()):
        return f"KC_{label}"
    elif label in KEY_PRESS_CODES:
        return KEY_PRESS_CODES[label][1]
    return None


def map_key_label_to_qmk(label):
    code = get_qmk_key_press_code(label)
    if code:
        return code
    elif label in LAYER_LABELS:
        return f"MO({LAYER_LABELS[label]})"
    elif label in SPECIAL_LABELS:
        return SPECIAL_LABELS[label][1]
    elif not label:
        return "KC_NO"
    raise KeyError(f"Cannot map label {label} to qmk.")


def map_tap_hold_key_to_qmk(tap, hold):
    if not hold:
        return map_key_label_to_qmk(tap)
    if not tap:
        return map_key_label_to_qmk(hold)
    if not hold and not tap:
        return "&trans"

    tap_code = get_qmk_key_press_code(tap)
    hold_code = get_qmk_key_press_code(hold)
    if tap_code:
        if hold_code:
            hold_code = hold_code.replace("KC_", "MOD_")
            return f"MT({hold_code},{tap_code})"
        if hold in LAYER_LABELS:
            return f"LT({LAYER_LABELS[hold]},{tap_code})"
    raise KeyError(f"Cannot map thumb key ({tap}, {hold}) to qmk.")


def generate_qmk_combo(combo):
    if not combo:
        return ('KC_NO', 'KC_NO')
    return (", ".join((map_key_label_to_qmk(combo.a), map_key_label_to_qmk(combo.b))), map_key_label_to_qmk(combo.result))


def generate_qmk_layer(markdown_layer):
    rows = []
    for markdown_row in markdown_layer.rows:
        rows.append(", ".join(map_tap_hold_key_to_qmk(tap, hold)
                    for tap, hold in markdown_row))
    return ",\n".join(rows)


################################################################################
# Mapping to ZMK


def get_zmk_key_press_code(label):
    if label.startswith('CMD_'):
        return f"LG({get_zmk_key_press_code(label[4:])})"
    elif label.startswith('HYP_'):
        return f"LA(LS(LC(LG({get_zmk_key_press_code(label[4:])}))))"
    elif label.startswith('HYP\\_'):
        return f"LA(LS(LC(LG({get_zmk_key_press_code(label[5:])}))))"
    elif label.isalpha() and len(label) == 1:
        return label
    elif label.isdigit() and len(label) == 1:
        return f"N{label}"
    elif label in KEY_PRESS_CODES:
        return KEY_PRESS_CODES[label][0]
    return None


def map_key_label_to_zmk(label):
    code = get_zmk_key_press_code(label)
    if code:
        return f"&kp {code}"
    elif label in LAYER_LABELS:
        return f"&mo {LAYER_LABELS[label]}"
    elif label in SPECIAL_LABELS:
        return SPECIAL_LABELS[label][0]
    elif not label:
        return "&trans"
    raise KeyError(f"Cannot map label {label} to zmk.")


def map_tap_hold_key_to_zmk(tap, hold):
    if not hold:
        return map_key_label_to_zmk(tap)
    if not tap:
        return map_key_label_to_zmk(hold)
    if not hold and not tap:
        return "&trans"

    tap_code = get_zmk_key_press_code(tap)
    hold_code = get_zmk_key_press_code(hold)
    if tap_code:
        if hold_code:
            return f"&mt_ {hold_code} {tap_code}"
        if hold in LAYER_LABELS:
            return f"&lt_ {LAYER_LABELS[hold]} {tap_code}"
    raise KeyError(f"Cannot map thumb key ({tap}, {hold}) to zmk.")


def generate_zmk_combo(combo):
    # TODO: Implement combos for ZMK
    return ('0 0', '&trans')


def generate_zmk_layer(markdown_layer):
    rows = []
    for markdown_row in markdown_layer.rows[:3]:
        # Add padding since this is a 5 row layout with a 6 row corne firmware
        rows.append(
            "&trans "
            + " ".join(map_tap_hold_key_to_zmk(tap, hold)
                       for tap, hold in markdown_row)
            + " &trans"
        )
    rows.append(
        " ".join(
            map_tap_hold_key_to_zmk(tap, hold) for tap, hold in markdown_layer.rows[3]
        )
    )
    return "\n".join(rows)


################################################################################
# Markdown Table Parsing

TapHold = namedtuple("TapHold", ("tap", "hold"))
Combo = namedtuple("Combo", ("a", "b", "result"))


class MarkdownLayer(object):
    "Parses markdown tables into a list of labels for each keyboard row."

    @classmethod
    def parse(cls, source_lines):
        layers = []
        table = None
        title = None
        for line in source_lines:
            line = line.strip()
            if line.startswith("#"):
                title = line
            elif table is not None:
                if line.startswith("|"):
                    table.append(
                        [
                            col.replace("☐", "").strip().upper()
                            for col in line.split("|")
                        ]
                    )
                else:
                    if "hold" in title.lower():
                        layers[-1].add_hold_table(table)
                    elif "combos" in title.lower():
                        layers[-1].add_combos_table(table)
                    else:
                        layers.append(cls(table))
                    table = None
            elif line.startswith("| -"):
                table = []
        if table is not None:
            layers.append(cls(table))
        return layers

    def __init__(self, markdown_table):
        self.rows = [
            [TapHold(k, None) for k in row]
            for row in self.extract_labels(markdown_table)
        ]
        self.combos = []

    def extract_labels(self, markdown_table):
        for row in markdown_table[:3]:
            yield (row[1:6] + row[7:12])
        thumb_row = markdown_table[3]
        yield thumb_row[3:6] + thumb_row[7:10]

    def add_hold_table(self, markdown_table):
        for i, row in enumerate(self.extract_labels(markdown_table)):
            for j, key in enumerate(row):
                if key:
                    self.rows[i][j] = TapHold(self.rows[i][j].tap, key)

    def add_combos_table(self, markdown_table):
        self.combos = [Combo(r[1], r[2], r[4]) for r in markdown_table]

    def __repr__(self) -> str:
        return repr(self.rows)

################################################################################
# Generate Keymaps


def generate_keymap(markdown_layers, template_path, generate_layer_fn, generate_combo_fn):
    with open(template_path, "r") as template_file:
        template = template_file.read()
        combos = []
        for i, markdown_layer in enumerate(markdown_layers):
            template = template.replace(
                f"#LAYER_{i}#", generate_layer_fn(markdown_layer)
            )
            combos.extend(markdown_layer.combos)
        for combo_id in range(5):
            (trigger, result) = generate_combo_fn(
                combos[combo_id] if combo_id < len(combos) else None)
            template = template.replace(f"#COMBO_TRIGGER_{combo_id}#", trigger)
            template = template.replace(f"#COMBO_RESULT_{combo_id}#", result)
        return HEADER + template


def usage():
    print(f"Usage: {sys.argv[0]} qmk|zmk")


def main(type: str):
    root_dir = os.path.dirname(__file__)
    with open(os.path.join(root_dir, "README.md"), "r") as source:
        markdown_layers = MarkdownLayer.parse(source.readlines())
        if type == "zmk":
            print(
                generate_keymap(                    markdown_layers,
                    os.path.join(root_dir, "zmk_template.dtsi"),
                    generate_zmk_layer,
                    generate_zmk_combo,
                )
            )
        elif type == "qmk":
            print(
                generate_keymap(
                    markdown_layers,
                    os.path.join(root_dir, "qmk_template.c"),
                    generate_qmk_layer,
                    generate_qmk_combo,
                )
            )
        else:
            usage()


if len(sys.argv) > 1:
    main(sys.argv[1])
else:
    usage()