generate_layout.py

import json
import argparse
import os
import sys
from pathlib import Path
from typing import List, Tuple, Optional

import keycodes

CPP_SKELETON = """#pragma once
#include "../common/key_properties.h"

// This file is generated by generate_layout.py. Do not edit manually.

namespace common::constants
{{

const int TOTAL_NUM_KEYS = {total_num_keys};
const int TOTAL_NUM_LEDS = {total_num_leds};
const int NUM_ROWS = {num_rows};
const int NUM_COLS = {num_cols};
constexpr float MAX_X = {max_x};
constexpr float MAX_Y = {max_y};

constexpr KeyDescription KEY_PROPERTIES[constants::TOTAL_NUM_KEYS] =
{{
    {rows}
}};

constexpr KeyDescription const* KEY_PROPERTIES_BY_ROW_COL[constants::NUM_ROWS][constants::NUM_COLS] =
{{
    {rows_and_cols}
}};

constexpr LEDDescription LED_PROPERTIES[constants::TOTAL_NUM_LEDS] =
{{
    {led_rows}
}};

}}
"""

ELM_SKELETON = """-- This file is generated by generate_layout.py. Do not edit manually.
module Generated.Layout exposing (..)

import Keyboard exposing (Key, Layout, KeyActions(..), Action(..))

initialLayout : Layout
initialLayout =
  [ {rows}
  ]

"""


class LEDDescription:
    def __init__(self, strip_index: int, x: float, y: float, key_description_index: Optional[int]):
        self.strip_index: int = strip_index
        self.x: float = x
        self.y: float = y
        self.key_description_index: Optional[int] = key_description_index

    def to_cpp(self) -> str:
        return f"LEDDescription {{ {self.strip_index}, {self.x}, {self.y}, {'nullptr' if self.key_description_index is None else '&KEY_PROPERTIES[' + str(self.key_description_index) + ']'} }}"


class KeyDescription:
    def __init__(self, row: int, col: int, x: float, y: float, height: float, width: float, default_key: int, default_key_str: Optional[str]):
        self.row: int = row
        self.col: int = col
        self.x: float = x
        self.y: float = y
        self.height: float = height
        self.width: float = width
        self.default_key: int = default_key
        self.default_key_str: Optional[str] = default_key_str
        self.led_strip_index: Optional[int] = None

    def __str__(self) -> str:
        return f"KeyDescription(row={self.row}, col={self.col}, x={self.x}, y={self.y}, height={self.height}, width={self.width}, default_key={self.default_key}, default_key_str={self.default_key_str})"

    def __repr__(self) -> str:
        return self.__str__()

    def to_cpp(self) -> str:
        return f"KeyDescription {{ {self.row}, {self.col}, {self.x}, {self.y}, {self.height}, {self.width}, {hex(self.default_key)} {'/* ' + self.default_key_str + ' */' if self.default_key_str is not None else ''}, {self.led_strip_index} }}"

    def to_elm(self, key_id: int) -> str:
        default_key = "Nothing"
        if self.default_key_str is not None:
            default_key = f"(Just <| Single {hex(self.default_key)} )"
        return f"Key {key_id} {self.row} {self.col} {self.x} {self.y} {self.height} {self.width} (LayerAction {{ actionLayer1 = {default_key}, actionLayer2 = Nothing, actionLayer3 = Nothing }} )"


def code_is_modifier(code: int) -> bool:
    return (code & 0xFF00) == 0xE000


def code_is_media(code: int) -> bool:
    return (code & 0xFF00) == 0xE400


def parse_args():
    parser = argparse.ArgumentParser(description="Generate C++ or Elm layout description from keyboard-layout-editor.com")
    parser.add_argument("-i", "--input",
                        help="Path to the input file, generated by keyboard-layout-editor.com",
                        type=Path, required=True)
    parser.add_argument("-o", "--output", help="Path to the generated C++ or Elm file.",
                        type=str, default="output.cpp")
    parser.add_argument("-l", "--language", type=str, default="cpp", choices=["cpp", "elm"], help="The language of the output file, either 'cpp' or 'elm'")
    parser.add_argument("-e", "--extra-led-descriptions", type=str, default=None, help="Path to a JSON file information on the LED:s which are not associated with any keys, such as the logo, as well as the starting direction for the snake pattern.")
    return parser.parse_args()


def extract_col(text: str) -> int:
    """Extracts the column number from the text, e.g. "col1" -> 1, "col10" -> 10"""
    col_half = text.lower().split("col")[1]
    if len(col_half) == 1 or not col_half[1].isdigit():
        return int(col_half[0])
    else:
        return int(col_half[0]) * 10 + int(col_half[1])


def extract_default_key(text: str) -> Optional[str]:
    col_num = extract_col(text)
    text_col_split = text.upper().split(f"COL{col_num}")
    key_first_half = keycodes.KEY_CODE_NAMES.get(text_col_split[0].strip().upper(), None)
    if key_first_half is not None:
        return key_first_half
    return keycodes.KEY_CODE_NAMES.get(text_col_split[1].strip().upper(), None)


def parse_json_layout(json_layout) -> Tuple[List[List[KeyDescription]], int]:
    max_col = 0
    curr_y = 0
    rows = []
    for row_num, row in enumerate(json_layout):
        curr_y += 1
        cols = []
        curr_x = 0
        width = 1
        set_width = False
        height = 1
        set_height = False
        for item in row:
            if isinstance(item, dict):
                if "y" in item:
                    curr_y += item["y"]
                if "x" in item:
                    curr_x += item["x"]
                if "w" in item:
                    width = item["w"]
                    set_width = True
                if "h" in item:
                    height = item["h"]
                    set_height = True
            else:
                if not set_width:
                    width = 1
                if not set_height:
                    height = 1

                col_num = extract_col(item)
                default_key_str = extract_default_key(item)
                default_key = keycodes.KEY_CODES.get(default_key_str, 0)
                key = KeyDescription(
                    row=row_num,
                    col=col_num,
                    x=curr_x,
                    y=curr_y,
                    height=height,
                    width=width,
                    default_key=default_key,
                    default_key_str=default_key_str)
                cols.append(key)

                curr_x += width
                set_width = False
                set_height = False

                max_col = max(max_col, col_num)

        rows.append(cols)
    return rows, max_col


def get_row_col_grid(rows: List[List[KeyDescription]], max_col: int) -> List[List[Optional[int]]]:
    grid = [[None] * (max_col + 1) for _ in range(len(rows))]
    index = 0
    for row in rows:
        for key in row:
            grid[key.row][key.col] = index
            index += 1
    return grid


def get_led_descriptions(rows: List[List[KeyDescription]], max_col: int, extra_led_descriptions: dict) -> List[LEDDescription]:
    direction = extra_led_descriptions["direction"]
    extra_leds = extra_led_descriptions["extra_leds"]

    extra_leds_by_index = {el["index"]: el for el in extra_leds}

    grid = get_row_col_grid(rows, max_col)
    curr_row = 0
    curr_col = 0 if direction == 1 else len(rows[0]) - 1
    index = 0
    leds = []
    while True:
        if curr_col == -1 or curr_col == len(rows[curr_row]):
            direction *= -1
            curr_row += 1
            if curr_row == len(rows):
                break

            curr_col = 0 if direction == 1 else len(rows[curr_row]) - 1

        if index in extra_leds_by_index:
            leds.append(LEDDescription(
                strip_index=index,
                x=extra_leds_by_index[index]["x"],
                y=extra_leds_by_index[index]["y"],
                key_description_index=None))
        else:
            key_description = rows[curr_row][curr_col]
            key_description.led_strip_index = index
            key_description_index = grid[key_description.row][key_description.col]
            leds.append(LEDDescription(
                strip_index=index,
                x=key_description.x,
                y=key_description.y,
                key_description_index=key_description_index))

            curr_col += direction

        index += 1
    return leds


def write_cpp(rows: List[List[KeyDescription]], max_col: int, output_path: str, extra_led_descriptions: dict):
    leds = get_led_descriptions(rows, max_col, extra_led_descriptions)

    num_rows = len(rows)
    num_cols = max_col + 1
    total_num_keys = sum([len(row) for row in rows])
    max_x = max(leds, key=lambda l: l.x).x
    max_y = max(leds, key=lambda l: l.y).y

    led_rows_str = ",\n    ".join([led.to_cpp() for led in leds])

    rows_str = ",\n    ".join([",\n    ".join(
        [key.to_cpp() for key in row]) for row in rows])

    grid = get_row_col_grid(rows, max_col)
    grid_rows = []
    for row_index, row in enumerate(grid):
        grid_cols = []
        for index in row:
            if index is None:
                grid_cols.append("nullptr")
            else:
                grid_cols.append(f"&KEY_PROPERTIES[{index}]")
        row = ",\n        ".join(grid_cols)
        grid_rows.append(f"{{\n        // Row {row_index}\n        {row}\n    }}")
    rows_and_cols = ",\n    ".join(grid_rows)

    cpp = CPP_SKELETON.format(
        total_num_keys=total_num_keys,
        rows=rows_str,
        num_rows=num_rows,
        rows_and_cols=rows_and_cols,
        num_cols=num_cols,
        led_rows=led_rows_str,
        total_num_leds=len(leds),
        max_x=max_x,
        max_y=max_y)
    with open(output_path, "w") as f:
        f.write(cpp)


def write_elm(rows: List[List[KeyDescription]], max_col: int, output_path: str):
    rows_list = []
    index = 0
    for row in rows:
        for key in row:
            rows_list.append(key.to_elm(index))
            index += 1

    rows_str = "\n  , ".join(rows_list)
    elm = ELM_SKELETON.format(rows=rows_str)
    with open(output_path, "w") as f:
        f.write(elm)


def main():
    args = parse_args()
    if not os.path.isfile(args.input):
        print(f"Input file {args.input} does not exist.")
        sys.exit(1)

    json_layout = None
    with open(args.input) as f:
        json_layout = json.load(f)

    extra_led_descriptions = {"direction": 1, "extra_leds": []}
    if args.extra_led_descriptions is not None:
        with open(args.extra_led_descriptions) as f:
            extra_led_descriptions = json.load(f)

    rows, max_col = parse_json_layout(json_layout)
    if args.language == "cpp":
        write_cpp(rows, max_col, args.output, extra_led_descriptions)
    elif args.language == "elm":
        write_elm(rows, max_col, args.output)


if __name__ == "__main__":
    main()