feat(templates): Add interconnect selection

When using "zmk keyboard new" to create a shield, it will now ask you to
select an interconnect. The generated template will set this as the
"requires" value in the .zmk.yml file and will use the correct GPIO node
label in devicetree code.

Also fixed an issue where adding a feature to the .zmk.yml file that
wasn't explicitly listed in the Python types would cause an error.

zmk/commands/keyboard/add.py

@@ -13,8 +13,14 @@
 from ...build import BuildItem, BuildMatrix
 from ...config import get_config
 from ...exceptions import FatalError
-from ...hardware import Board, Keyboard, Shield, get_hardware, is_compatible
-from ...menu import show_menu
+from ...hardware import (
+    Board,
+    Keyboard,
+    Shield,
+    get_hardware,
+    is_compatible,
+    show_hardware_menu,
+)
 from ...repo import Repo
 from ...util import spinner
 
@@ -85,17 +91,15 @@ def keyboard_add(
 
     # Prompt the user for any necessary components they didn't specify
     if keyboard is None:
-        keyboard = show_menu(
-            "Select a keyboard:", hardware.keyboards, filter_func=_filter
-        )
+        keyboard = show_hardware_menu("Select a keyboard:", hardware.keyboards)
 
     if isinstance(keyboard, Shield):
         if controller is None:
             hardware.controllers = [
                 c for c in hardware.controllers if is_compatible(c, keyboard)
             ]
-            controller = show_menu(
-                "Select a controller:", hardware.controllers, filter_func=_filter
+            controller = show_hardware_menu(
+                "Select a controller:", hardware.controllers
             )
 
         # Sanity check that everything is compatible
@@ -116,11 +120,6 @@ def keyboard_add(
     console.print(f'Run "zmk code {keyboard.id}" to edit the keymap.')
 
 
-def _filter(item: Keyboard, text: str):
-    text = text.casefold().strip()
-    return text in item.id.casefold() or text in item.name.casefold()
-
-
 class KeyboardNotFound(FatalError):
     """Fatal error for an invalid keyboard ID"""
 

zmk/commands/keyboard/new.py

@@ -13,8 +13,11 @@
 from ...backports import StrEnum
 from ...config import get_config
 from ...exceptions import FatalError
+from ...hardware import Interconnect, get_hardware, show_hardware_menu
 from ...menu import detail_list, show_menu
+from ...repo import Repo
 from ...templates import get_template_files
+from ...util import spinner
 
 
 class KeyboardType(StrEnum):
@@ -49,6 +52,7 @@ class TemplateData:
 
 ID_PATTERN = re.compile(r"[a-z_]\w*")
 MAX_NAME_LENGTH = 16
+DEFAULT_INTERCONNECT = "pro_micro"
 
 
 def _validate_id(value: str):
@@ -133,6 +137,13 @@ def keyboard_new(
         KeyboardLayout | None,
         typer.Option("--layout", "-l", help="Keyboard hardware layout."),
     ] = None,
+    interconnect_id: Annotated[
+        str | None,
+        typer.Option(
+            "--interconnect",
+            help="If creating a shield, the interconnect ID for the controller board.",
+        ),
+    ] = None,
     force: Annotated[
         bool, typer.Option("--force", "-f", help="Overwrite existing files.")
     ] = False,
@@ -160,6 +171,11 @@ def keyboard_new(
     if not keyboard_type:
         keyboard_type = _prompt_keyboard_type()
 
+    if not interconnect_id and keyboard_type == KeyboardType.SHIELD:
+        interconnect = _prompt_interconnect(repo)
+    else:
+        interconnect = _get_interconnect(repo, interconnect_id)
+
     if not keyboard_platform:
         if keyboard_type == KeyboardType.BOARD:
             keyboard_platform = _prompt_keyboard_platform()
@@ -176,6 +192,7 @@ def keyboard_new(
         keyboard_name=keyboard_name,
         short_name=short_name,
         keyboard_id=keyboard_id,
+        interconnect=interconnect,
     )
 
     dest = board_root / template.dest
@@ -214,6 +231,42 @@ def _prompt_keyboard_type():
     return result.data
 
 
+def _prompt_interconnect(repo: Repo):
+    with spinner("Finding interconnects..."):
+        hardware = get_hardware(repo)
+
+    default_index = next(
+        (
+            i
+            for i, interconnect in enumerate(hardware.interconnects)
+            if interconnect.id == DEFAULT_INTERCONNECT
+        ),
+        0,
+    )
+
+    return show_hardware_menu(
+        "Select the interconnect for the controller board:",
+        hardware.interconnects,
+        default_index=default_index,
+    )
+
+
+def _get_interconnect(repo: Repo, interconnect_id: str | None):
+    if not interconnect_id:
+        return None
+
+    with spinner("Finding interconnects..."):
+        hardware = get_hardware(repo)
+
+    try:
+        return next(ic for ic in hardware.interconnects if ic.id == interconnect_id)
+    except StopIteration as ex:
+        raise FatalError(
+            f'"{interconnect_id}" is not a valid interconnect. '
+            'Run "zmk keyboard list --type interconnect" to list possible values.'
+        ) from ex
+
+
 def _prompt_keyboard_platform():
     items = detail_list(
         [
@@ -309,6 +362,11 @@ def ask(  # pyright: ignore[reportIncompatibleMethodOverride]
     KeyboardPlatform.NRF52840: "arm",
 }
 
+_DEFAULT_GPIO = "&gpio0"
+_PLATFORM_GPIO: dict[KeyboardPlatform, str] = {
+    KeyboardPlatform.NRF52840: "&gpio0",
+}
+
 
 def _get_template(
     keyboard_type: KeyboardType,
@@ -317,22 +375,33 @@ def _get_template(
     keyboard_name: str,
     short_name: str,
     keyboard_id: str,
+    interconnect: Interconnect | None = None,
 ):
     template = TemplateData()
     template.data["id"] = keyboard_id
     template.data["name"] = keyboard_name
     template.data["shortname"] = short_name
     template.data["keyboard_type"] = str(keyboard_type)
+    template.data["interconnect"] = ""
     template.data["arch"] = ""
+    template.data["gpio"] = _DEFAULT_GPIO
 
     match keyboard_type:
         case KeyboardType.SHIELD:
             template.folder = "shield/"
             template.dest = f"shields/{keyboard_id}"
 
+            if interconnect:
+                template.data["interconnect"] = interconnect.id
+                try:
+                    template.data["gpio"] = "&" + interconnect.node_labels["gpio"]
+                except KeyError:
+                    pass
+
         case _:
             arch = _PLATFORM_ARCH.get(keyboard_platform, _DEFAULT_ARCH)
             template.data["arch"] = arch
+            template.data["gpio"] = _PLATFORM_GPIO.get(keyboard_platform, _DEFAULT_GPIO)
 
             template.folder = f"board/{keyboard_platform}/"
             template.dest = f"{arch}/{keyboard_id}"

zmk/hardware.py

@@ -6,25 +6,36 @@
 from dataclasses import dataclass, field
 from functools import reduce
 from pathlib import Path
-from typing import Any, Literal, Type, TypeAlias, TypeGuard, TypeVar
+from typing import Any, Literal, Type, TypeAlias, TypedDict, TypeGuard, TypeVar
 
 import dacite
 
+from .menu import show_menu
 from .repo import Repo
 from .util import flatten
 from .yaml import read_yaml
 
-Feature: TypeAlias = Literal[
-    "keys", "display", "encoder", "underglow", "backlight", "pointer"
-]
+Feature: TypeAlias = (
+    Literal["keys", "display", "encoder", "underglow", "backlight", "pointer", "studio"]
+    | str
+)
 Output: TypeAlias = Literal["usb", "ble"]
 
-# TODO: dict should match { id: str, features: list[Feature] }
-Variant: TypeAlias = str | dict[str, str]
+
+class VariantDict(TypedDict):
+    """Keyboard variant with custom options"""
+
+    id: str
+    features: list[Feature]
+
+
+Variant: TypeAlias = str | VariantDict
 
 # TODO: replace with typing.Self once minimum Python version is >= 3.11
 _Self = TypeVar("_Self", bound="Hardware")
 
+_HW = TypeVar("_HW", bound="Hardware")
+
 
 @dataclass
 class Hardware:
@@ -241,3 +252,22 @@ def _find_hardware(path: Path) -> Generator[Hardware, None, None]:
 
             case "interconnect":
                 yield Interconnect.from_dict(meta)
+
+
+def _filter_hardware(item: Hardware, text: str):
+    text = text.casefold().strip()
+    return text in item.id.casefold() or text in item.name.casefold()
+
+
+def show_hardware_menu(
+    title: str,
+    items: Iterable[_HW],
+    **kwargs,
+) -> _HW:
+    """
+    Show a menu to select from a list of Hardware objects.
+
+    kwargs are passed through to zmk.menu.show_menu(), except for filter_func,
+    which is set to a function appropriate for filtering Hardware objects.
+    """
+    return show_menu(title=title, items=items, **kwargs, filter_func=_filter_hardware)

zmk/templates/__init__.py

@@ -7,7 +7,9 @@
     name: str -- The keyboard display name
     shortname: str -- A name abbreviated to <= 16 characters
     keyboard_type: str -- "board" or "shield"
-    arch: Optional[str] -- The board architecture, e.g "arm"
+    interconnect: str -- The interconnect ID for the controller board. May be empty.
+    arch: str -- The board architecture, e.g "arm". May be empty.
+    gpio: str -- The default node label for GPIO, e.g. "&gpio0".
 
 """
 

zmk/templates/board/nrf52840/common_inner.dtsi

@@ -5,7 +5,7 @@
         // If you have a GPIO routed to a status LED, set it here.
         // Otherwise, delete this block.
         blue_led: led_0 {
-            gpios = </* &gpio0 0 */ GPIO_ACTIVE_HIGH>;
+            gpios = <${gpio} 0 GPIO_ACTIVE_HIGH>;
         };
     };
 

zmk/templates/common/kscan.dtsi

@@ -1,4 +1,3 @@
-<%page args="node = '&gpio0'" />
     kscan: kscan {
         // If the hardware does not use a switch matrix, change this to the
         // appropriate driver and update the properties below to match.
@@ -13,12 +12,12 @@
         // Replace these comments with the GPIO pins in the matrix.
         // See https://zmk.dev/docs/development/hardware-integration/new-shield#shield-overlays
         col-gpios
-            = </* ${node} 0 */ GPIO_ACTIVE_HIGH>
-            , </* ${node} 0 */ GPIO_ACTIVE_HIGH>
+            = <${gpio} 0 GPIO_ACTIVE_HIGH>
+            , <${gpio} 0 GPIO_ACTIVE_HIGH>
             ;
 
         row-gpios
-            = </* ${node} 0 */ (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
-            , </* ${node} 0 */ (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
+            = <${gpio} 0 (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
+            , <${gpio} 0 (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
             ;
     };

zmk/templates/common/kscan_split_common.dtsi

@@ -1,4 +1,3 @@
-<%page args="node = '&gpio0'" />
     kscan: kscan {
         // If the hardware does not use a switch matrix, change this to the
         // appropriate driver and update the properties below to match.
@@ -14,7 +13,7 @@
         // Replace these comments with the GPIO pins in the matrix.
         // See https://zmk.dev/docs/development/hardware-integration/new-shield#shield-overlays
         row-gpios
-            = </* ${node} 0 */ (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
-            , </* ${node} 0 */ (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
+            = <${gpio} 0 (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
+            , <${gpio} 0 (GPIO_ACTIVE_HIGH | GPIO_PULL_DOWN)>
             ;
     };

zmk/templates/common/kscan_split_left.dtsi

@@ -1,9 +1,8 @@
-<%page args="node = '&gpio0'" />
 &kscan {
     // Replace these comments with the GPIO pins in the matrix for the left side.
     // See https://zmk.dev/docs/development/hardware-integration/new-shield#shield-overlays
     col-gpios
-        = </* ${node} 0 */ GPIO_ACTIVE_HIGH>
-        , </* ${node} 0 */ GPIO_ACTIVE_HIGH>
+        = <${gpio} 0 GPIO_ACTIVE_HIGH>
+        , <${gpio} 0 GPIO_ACTIVE_HIGH>
         ;
 };

zmk/templates/common/kscan_split_right.dtsi

@@ -1,4 +1,3 @@
-<%page args="node = '&gpio0'" />
 &default_transform {
     // Set this to the number of columns on the left side.
     col-offset = <2>;
@@ -8,7 +7,7 @@
     // Replace these comments with the GPIO pins in the matrix for the right side.
     // See https://zmk.dev/docs/development/hardware-integration/new-shield#shield-overlays
     col-gpios
-        = </* ${node} 0 */ GPIO_ACTIVE_HIGH>
-        , </* ${node} 0 */ GPIO_ACTIVE_HIGH>
+        = <${gpio} 0 GPIO_ACTIVE_HIGH>
+        , <${gpio} 0 GPIO_ACTIVE_HIGH>
         ;
 };

zmk/templates/common/shield_left.overlay

@@ -4,4 +4,4 @@
 
 #include "${id}.dtsi"
 
-<%include file="kscan_split_left.dtsi" args="node = '&pro_micro'" />
+<%include file="kscan_split_left.dtsi" />

zmk/templates/common/shield_right.overlay

@@ -4,4 +4,4 @@
 
 #include "${id}.dtsi"
 
-<%include file="kscan_split_right.dtsi" args="node = '&pro_micro'" />
+<%include file="kscan_split_right.dtsi" />

zmk/templates/shield/split/${id}.dtsi

@@ -1,6 +1,6 @@
 <%inherit file="/common/shield.overlay" />
 <%block name="kscan">
-<%include file="/common/kscan_split_common.dtsi" args="node = '&pro_micro'" />
+<%include file="/common/kscan_split_common.dtsi" />
 </%block>
 <%block name="matrix_transform">
 <%include file="/common/matrix_transform_split.dtsi" />

zmk/templates/shield/split/${id}.zmk.yml

@@ -1,5 +1,5 @@
 <%inherit file="/common/hardware.zmk.yml" />
-# Set this to the interconnect the shield uses.
-# Run "zmk keyboard list --type interconnect" to get the possible values.
+% if interconnect:
 requires:
-  - pro_micro
+  - ${interconnect}
+% endif

zmk/templates/shield/unibody/${id}.overlay

@@ -1,4 +1,4 @@
 <%inherit file="/common/shield.overlay" />
 <%block name="kscan">
-<%include file="/common/kscan.dtsi" args="node = '&pro_micro'" />
+<%include file="/common/kscan.dtsi" />
 </%block>

zmk/templates/shield/unibody/${id}.zmk.yml

@@ -1,5 +1,5 @@
 <%inherit file="/common/hardware.zmk.yml" />
-# Set this to the interconnect the shield uses.
-# Run "zmk keyboard list --type interconnect" to get the possible values.
+% if interconnect:
 requires:
-  - pro_micro
+  - ${interconnect}
+% endif