Update SPI reassignment script for Bookworm

raspi-spi-reassign.py

@@ -18,12 +18,18 @@ shell.group="SPI Reassign"
 allowed_gpios = (4, 5, 6, 7, 8, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27)
 spi0_default_pins = (8, 7)
 
+DISABLE_GPIO_TEXT = "Disabled"
+
+boot_dir = "/boot/firmware"
+if not shell.exists(boot_dir) or not shell.isdir(boot_dir):
+    boot_dir = "/boot"
+
 """
 For now this will only ask about SPI0, but we can later add SPI1
 """
 
 def valid_pins(ce0_pin, ce1_pin):
-    if ce0_pin is None or ce1_pin is None:
+    if ce0_pin is None and ce1_pin is not None:
         return False
     if ce0_pin == ce1_pin:
         return False
@@ -41,19 +47,35 @@ def enable_spi():
     print("Enabling SPI")
     shell.run_command("sudo raspi-config nonint do_spi 0")
 
+def spi_disabled():
+    return shell.run_command("sudo raspi-config nonint get_spi", suppress_message=True, return_output=True).strip() == "1"
+
 def remove_custom():
-    shell.pattern_replace("/boot/config.txt", 'dtoverlay=spi0-2cs,cs.*?\n', multi_line=True)
+    shell.pattern_replace(f"{boot_dir}/config.txt", 'dtoverlay=spi0-[0-2]cs,cs.*?\n', multi_line=True)
+
+def format_gpio(gpio):
+    if gpio is None:
+        return DISABLE_GPIO_TEXT
+    return f"GPIO {gpio}"
+
+def gpio_options(pool):
+    options = []
+    for gpio in pool:
+        options.append(format_gpio(gpio))
+    return options
 
 def write_new_custom(ce0_pin, ce1_pin):
-    if (ce0_pin, ce1_pin) == spi0_default_pins:
-        enable_spi()
-    else:
+    if (ce0_pin, ce1_pin) != spi0_default_pins:
         overlay_command = "dtoverlay=spi0-2cs"
-        if ce0_pin != spi0_default_pins[0]:
+        if ce0_pin is None and ce1_pin is None:
+            overlay_command = "dtoverlay=spi0-0cs"
+        elif ce1_pin is None:
+            overlay_command = "dtoverlay=spi0-1cs"
+        if ce0_pin != spi0_default_pins[0] and ce0_pin is not None:
             overlay_command += ",cs0_pin={}".format(ce0_pin)
-        if ce1_pin != spi0_default_pins[1]:
+        if ce1_pin != spi0_default_pins[1] and ce1_pin is not None:
             overlay_command += ",cs1_pin={}".format(ce1_pin)
-        shell.write_text_file("/boot/config.txt", overlay_command + "\n")
+        shell.write_text_file(f"{boot_dir}/config.txt", overlay_command + "\n")
 
 @click.command()
 @click.option('--ce0', nargs=1, default=None, help="Specify a GPIO for CE0")
@@ -67,7 +89,6 @@ def main(ce0, ce1, reboot):
         auto_reboot = reboot.lower() == 'yes'
     if valid_pins(ce0, ce1):
         remove_custom()
-        disable_spi()
         write_new_custom(ce0, ce1)
         if auto_reboot:
             shell.reboot()
@@ -86,6 +107,16 @@ lines doesn't interfere with CircuitPython.
 Run time of < 1 minute. Reboot required for
 changes to take effect!
 """)
+    menu_options = [
+        "Reassign SPI Chip Enable Pins",
+        "Reset to Defaults Pins",
+        "Disable SPI",
+        "Exit",
+    ]
+
+    if spi_disabled():
+        menu_options[2] = "Enable SPI"
+
     shell.info("{} detected.\n".format(pi_model))
     if not shell.is_raspberry_pi():
         shell.bail("Non-Raspberry Pi board detected. This must be run on a Raspberry Pi")
@@ -93,37 +124,39 @@ changes to take effect!
     if os_identifier != "Raspbian":
         shell.bail("Sorry, the OS detected was {}. This script currently only runs on Raspberry Pi OS.".format(os_identifier))
     menu_selection = shell.select_n(
-        "Select an option:", (
-            "Reassign SPI Chip Enable Pins",
-            "Reset to Defaults Pins",
-            "Disable SPI",
-            "Exit",
-        )
+        "Select an option:", menu_options
     )
     if menu_selection == 1:
         while True:
             # Reassign
             gpio_pool = list(allowed_gpios)
+            gpio_pool.append(None)
             # Ask for pin for CE0
-            ce0_selection = shell.select_n("Select a new GPIO for CE0", ["GPIO {}".format(x) for x in gpio_pool])
+            ce0_selection = shell.select_n("Select a new GPIO for CE0", gpio_options(gpio_pool))
             ce0_pin = gpio_pool[ce0_selection - 1]
-            gpio_pool.remove(ce0_pin)
-            # Ask for pin for CE1
-            ce1_selection = shell.select_n("Select a new GPIO for CE1", ["GPIO {}".format(x) for x in gpio_pool])
-            ce1_pin = gpio_pool[ce1_selection - 1]
-            if shell.prompt("The new GPIO {} for CE0 and GPIO {} for CE1. Is this correct?".format(ce0_pin, ce1_pin)):
+            if ce0_pin is not None:
+                gpio_pool.remove(ce0_pin)
+                # Ask for pin for CE1
+                ce1_selection = shell.select_n("Select a new GPIO for CE1", gpio_options(gpio_pool))
+                ce1_pin = gpio_pool[ce1_selection - 1]
+            else:
+                ce1_pin = None
+            if shell.prompt(f"The new settings will be {format_gpio(ce0_pin)} for CE0 and {format_gpio(ce1_pin)} for CE1. Is this correct?"):
                 break
         remove_custom()
-        disable_spi()
         write_new_custom(ce0_pin, ce1_pin)
+        if spi_disabled():
+            enable_spi()
     elif menu_selection == 2:
         # Reset to Defaults
         remove_custom()
-        enable_spi()
     elif menu_selection == 3:
-        # Disable
-        remove_custom()
-        disable_spi()
+        # Enable/Disable SPI
+        if spi_disabled():
+            enable_spi()
+        else:
+            disable_spi()
+            remove_custom()
     elif menu_selection == 4:
         # Exit
         shell.exit(0)