Add support for old Robot Brains with swapped pins (#187)

* Add support for old RB with swapped pins

* move wait time to set_en and set_g0

* Apply suggestion from @falkoschindler

* Apply suggestion from @falkoschindler

* add support for use of espresso in docker

* refactor GPIO controller implementation to support gpiod v1 and v2 APIs

* some more cleanup

---------

Co-authored-by: Falko Schindler <falko@zauberzeug.com>

Author: Johannes-Thiel

espresso.py

@@ -4,31 +4,91 @@
 import sys
 import time
 from contextlib import contextmanager
-from typing import Generator, Optional
+from typing import Generator
 from pathlib import Path
 import argparse
+
 try:
     import gpiod
 except ImportError:
-    gpiod = None  # type: ignore[assignment]
-    print('gpiod module not found. Please install it using "sudo apt install python3-gpiod". Ignore this error if you are not on a Jetson board.')
+    GPIOD_VERSION = None
+else:
+    GPIOD_VERSION = 2 if hasattr(gpiod, 'request_lines') else 1
+
+
+class GpioController:
+    def __init__(self, en: str, g0: str) -> None:
+        pass
+
+    def request_outputs(self, consumer: str) -> None:
+        """Request output lines."""
+
+    def set_value(self, key: str, value: int) -> None:
+        """Set the value of a line."""
+
+    def release(self) -> None:
+        """Release the lines."""
+
+
+class GpioControllerV1(GpioController):
+
+    def __init__(self, en: str, g0: str) -> None:
+        chip = gpiod.Chip('gpiochip0')
+        self._lines = {
+            'en': chip.find_line(en),
+            'g0': chip.find_line(g0),
+        }
+
+    def request_outputs(self, consumer: str) -> None:
+        for line in self._lines.values():
+            line.request(consumer=consumer, type=gpiod.LINE_REQ_DIR_OUT)
+
+    def set_value(self, key: str, value: int) -> None:
+        self._lines[key].set_value(value)
+
+    def release(self) -> None:
+        for line in self._lines.values():
+            line.release()
+
+
+class GpioControllerV2(GpioController):
+    CHIP_PATH = '/dev/gpiochip0'
+
+    def __init__(self, en: str, g0: str) -> None:
+        chip = gpiod.Chip(self.CHIP_PATH)
+        self._offsets = {
+            'en': chip.line_offset_from_id(en),
+            'g0': chip.line_offset_from_id(g0),
+        }
+        chip.close()
+        self._request = None
+
+    def request_outputs(self, consumer: str) -> None:
+        from gpiod.line import Direction  # pylint: disable=import-outside-toplevel
+        config = {offset: gpiod.LineSettings(direction=Direction.OUTPUT) for offset in self._offsets.values()}
+        self._request = gpiod.request_lines(self.CHIP_PATH, consumer=consumer, config=config)
+
+    def set_value(self, key: str, value: int) -> None:
+        from gpiod.line import Value  # pylint: disable=import-outside-toplevel
+        if self._request:
+            self._request.set_value(self._offsets[key], Value.ACTIVE if value else Value.INACTIVE)
 
+    def release(self) -> None:
+        if self._request:
+            self._request.release()
+            self._request = None
 
-JETPACK: Optional[int] = None
+
+DEFAULT_DEVICE = '/dev/tty.SLAB_USBtoUART'
 path = Path('/etc/nv_tegra_release')
 if path.exists() and (match := re.search(r'R(\d+)', path.read_text(encoding='utf-8'))):
     major = int(match.group(1))
     if major == 35:
-        JETPACK = 5
+        DEFAULT_DEVICE = '/dev/ttyTHS0'
     elif major == 36:
-        JETPACK = 6
+        DEFAULT_DEVICE = '/dev/ttyTHS1'
     else:
         raise RuntimeError(f'Unsupported L4T (Linux for Tegra) version: {major}')
-DEFAULT_DEVICE = {
-    5: '/dev/ttyTHS0',
-    6: '/dev/ttyTHS1',
-    None: '/dev/tty.SLAB_USBtoUART',
-}[JETPACK]
 
 parser = argparse.ArgumentParser(description='Flash and control an ESP32 microcontroller from a Jetson board')
 
@@ -38,6 +98,7 @@
 parser.add_argument('--bootloader', default='build/bootloader/bootloader.bin', help='Path to bootloader')
 parser.add_argument('--partition-table', default='build/partition_table/partition-table.bin',
                     help='Path to partition table')
+parser.add_argument('--swap', action='store_true', help='Swap En and G0 pins for piggyboard version lower than v0.5')
 parser.add_argument('--firmware', default='build/lizard.bin', help='Path to firmware binary')
 parser.add_argument('--chip', choices=['esp32', 'esp32s3'], default='esp32', help='ESP chip type')
 parser.add_argument('-d', '--dry-run', action='store_true', help='Dry run')
@@ -47,6 +108,7 @@
 ON = 1 if args.nand else 0
 OFF = 0 if args.nand else 1
 DRY_RUN = args.dry_run
+SWAP = args.swap
 CHIP = args.chip
 DEVICE = args.device
 FLASH_FREQ = {'esp32': '40m', 'esp32s3': '80m'}.get(CHIP, '40m')
@@ -55,54 +117,48 @@
 PARTITION_TABLE = args.partition_table
 FIRMWARE = args.firmware
 
-if JETPACK:
-    chip = gpiod.Chip('gpiochip0')
-    en = chip.find_line('PR.04')
-    g0 = chip.find_line('PAC.06')
+EN = 'PR.04'
+G0 = 'PAC.06'
+if SWAP:
+    EN, G0 = G0, EN
+gpio = {
+    None: GpioController,
+    1: GpioControllerV1,
+    2: GpioControllerV2,
+}[GPIOD_VERSION](EN, G0)
 
 
 @contextmanager
 def _pin_config() -> Generator[None, None, None]:
     """Configure the EN and G0 pins to control the microcontroller."""
-    if JETPACK:
-        print_bold('Configuring EN and G0 pins...')
-        en.request(consumer='espresso', type=gpiod.LINE_REQ_DIR_OUT)
-        g0.request(consumer='espresso', type=gpiod.LINE_REQ_DIR_OUT)
-        time.sleep(0.5)
+    print_bold('Configuring EN and G0 pins...')
+    gpio.request_outputs(consumer='espresso')
+    time.sleep(0.5)
     yield
-    if JETPACK:
-        _release_pins()
+    _release_pins()
 
 
 def _release_pins() -> None:
     """Release pins."""
-    if JETPACK:
-        en.release()
-        g0.release()
+    gpio.release()
 
 
 @contextmanager
 def _flash_mode() -> Generator[None, None, None]:
     """Bring the microcontroller into flash mode."""
-    if JETPACK:
-        print_bold('Bringing the microcontroller into flash mode...')
-        set_en(ON)
-        time.sleep(0.5)
-        set_g0(ON)
-        time.sleep(0.5)
-        set_en(OFF)
-        time.sleep(0.5)
+    print_bold('Bringing the microcontroller into flash mode...')
+    set_en(ON)
+    set_g0(ON)
+    set_en(OFF)
     yield
-    if JETPACK:
-        _reset()
+    _reset()
 
 
 def enable() -> None:
     """Enable the microcontroller."""
     print_bold('Enabling the microcontroller...')
     with _pin_config():
         set_g0(OFF)
-        time.sleep(0.5)
         set_en(OFF)
 
 
@@ -123,9 +179,7 @@ def reset() -> None:
 def _reset() -> None:
     """Set pins to reset the microcontroller."""
     set_g0(OFF)
-    time.sleep(0.5)
     set_en(ON)
-    time.sleep(0.5)
     set_en(OFF)
 
 
@@ -184,13 +238,15 @@ def run(*run_args: str) -> bool:
 def set_en(value: int) -> None:
     print(f'  Setting EN pin to {value}')
     if not DRY_RUN:
-        en.set_value(value)
+        gpio.set_value('en', value)
+        time.sleep(0.5)
 
 
 def set_g0(value: int) -> None:
     print(f'  Setting G0 pin to {value}')
     if not DRY_RUN:
-        g0.set_value(value)
+        gpio.set_value('g0', value)
+        time.sleep(0.5)
 
 
 def print_ok(message: str) -> None:
@@ -207,6 +263,8 @@ def print_fail(message: str) -> None:
 
 def main(command: str) -> None:
     print_ok('Espresso dry-running...' if DRY_RUN else 'Espresso running...')
+    if GPIOD_VERSION is None and not DRY_RUN:
+        print_fail('Module gpiod not found. Espresso is not able to control the EN and G0 pins.')
     if command == 'enable':
         enable()
     elif command == 'disable':