CircuitPython ESP32 Wi-Fi Setup Helper

setup_wifi.py

# SPDX-FileCopyrightText: 2023 Tammy Makes Things
# SPDX-License-Identifier: MIT

import os
import board
import busio
from digitalio import DigitalInOut
import adafruit_requests as requests
import adafruit_esp32spi.adafruit_esp32spi_socket as socket
from adafruit_esp32spi import adafruit_esp32spi

try:
    from typing import Optional, Tuple

    SetupWiFiReturnType = Tuple[
        Optional[busio.SPI], Optional[adafruit_esp32spi.ESP_SPIcontrol]
    ]
except ImportError:
    pass


def setup_wifi(
    connect_via: Optional[str] = os.getenv("SETUP_WIFI_CONNECT_VIA", "esp32_onboard"),
    max_retries: Optional[int] = int(os.getenv("SETUP_WIFI_MAX_RETRIES", "5")),
    esp_debug: Optional[int] = bool(os.getenv("SETUP_WIFI_ESP_DEBUG", "")),
    esp32_cs_pin_name: Optional[str] = None,
    esp32_ready_pin_name: Optional[str] = None,
    esp32_reset_pin_name: Optional[str] = None,
    wifi_network_name: Optional[str] = None,
    wifi_network_password: Optional[str] = None,
    verbose_logging: Optional[bool] = False,
) -> SetupWiFiReturnType:
    """
    Automates and simplifies the setup of Wi-Fi on CircuitPython devices.

    Args:
        connect_via (str, optional): Specify how the ESP32 is connected. Valid
            values are 'esp32_onboard', 'airlift_shield', 'airlift_featherwing',
            or 'esp32_external'. These specify the default pin names that will
            be used to connect to the ESP32 device. If you provide any other
            value, you'll need to specify the pin names manually. Defaults to
            the value of the environment variable ``SETUP_WIFI_CONNECT_VIA``, or
            ``esp32_onboard`` if that variable isn't set in ``settings.toml``.
            
            The pin value defaults supplied by this parameter are taken from the
            Learn guide "Adafruit PyPortal - IoT for CircuitPython". See
            <https://learn.adafruit.com/adafruit-pyportal/internet-connect>.

            The default pin values changed by this setting are:

            ===============             ==========      ============    =============
            ``connect_via``             CS Pin          READY Pin       RESET Pin
            ===============             ==========      ============    =============
            ``esp32_onboard``           ``ESP_CS``      ``ESP_BUSY``    ``ESP_RESET``
            ``airlift_featherwing``     ``D13``         ``D11``         ``D12``
            ``airlift_shield``          ``D10``         ``D7``          ``D5``
            ``esp32_external``          ``D9``          ``D10``         ``D5``
            ===============             ==========      ============    =============

            You can also select one of these modes and then override one or more
            pin settings; any pins not overridden will be set as per this table.

        max_retries (int, optional): The maximum number of retries when attempting
            to connect to Wi-Fi. Defaults to the value of the environment variable
            ``SETUP_WIFI_MAX_RETRIES``, or ``5`` if that variable isn't set in
            ``settings.toml``.

        esp_debug (bool, optional): If ``True``, enables debugging on the
            :py:class:`ESP_SPIcontrol` object. Defaults to the value of the
            environment variable ``SETUP_WIFI_ESP_DEBUG``, or ``False`` if that
            variable isn't set in ``settings.toml``.

        esp32_cs_pin_name (str, optional): The name of the ESP32 CS pin. Defaults
            based on the value of `connect_via`, but can also be overridden by
            setting the environment variable ``SETUP_WIFI_PIN_ESP32_CS`` in
            ``settings.toml``.

        esp32_ready_pin_name (str, optional): The name of the ESP32 READY pin. Defaults
            based on the value of `connect_via`, but can also be overridden by
            setting the environment variable ``SETUP_WIFI_PIN_ESP32_READY`` in
            ``settings.toml``.

        esp32_reset_pin_name (str, optional): The name of the ESP32 RESET pin. Defaults
            based on the value of `connect_via`, but can also be overridden by
            setting the environment variable ``SETUP_WIFI_PIN_ESP32_RESET`` in
            ``settings.toml``.

        wifi_network_name (str, optional): The SSID of the Wi-Fi network to
            connect to. Defaults to the value of the environment variable
            ``CIRCUITPY_WIFI_SSID``, but can also be overridden here for added
            flexibility.

        wifi_network_password (str, optional): The password of the Wi-Fi network to
            connect to. Defaults to the value of the environment variable
            ``CIRCUITPY_WIFI_PASSWORD``, but can also be overridden here for added
            flexibility.

        verbose_logging (bool, optional): Enable more verbose debugging output.
            Defaults to ``False`` if not supplied.

    Returns:
        SetupWiFiReturnType - If the connection is successful, returns the
            SPI and SPIcontrol objects so they can be reused. If an error
            occurs anywhere in the process, returns ``None`` for both
            objects.
    """

    if verbose_logging:
        print("===> setup_wifi starting")
        print("")

    wifi_ssid = wifi_network_name or os.getenv("CIRCUITPY_WIFI_SSID", None)
    wifi_password = wifi_network_password or os.getenv("CIRCUITPY_WIFI_PASSWORD", None)

    try:
        if not all([wifi_ssid, wifi_password]):
            if verbose_logging:
                print("* Wi-Fi Credentials not configured - skipping Wi-Fi Setup")
            return (None, None)

        esp32_cs_pin, esp32_ready_pin, esp32_reset_pin = (
            os.getenv("SETUP_WIFI_PIN_ESP32_CS", esp32_cs_pin_name),
            os.getenv("SETUP_WIFI_PIN_ESP32_READY", esp32_ready_pin_name),
            os.getenv("SETUP_WIFI_PIN_ESP32_RESET", esp32_reset_pin_name),
        )
        esp32_cs, esp32_ready, esp32_reset = (None, None, None)

        if connect_via == "esp32_onboard":
            if verbose_logging:
                print("- Connection method = esp32_onboard")
            esp32_cs_pin = esp32_cs_pin or "ESP_CS"
            esp32_ready_pin = esp32_ready_pin or "ESP_BUSY"
            esp32_reset_pin = esp32_reset or "ESP_RESET"
        elif connect_via == "airlift_shield":
            if verbose_logging:
                print("- Connection method = airlift_shield")
            esp32_cs_pin = esp32_cs_pin or "D10"
            esp32_ready_pin = esp32_ready_pin or "D7"
            esp32_reset_pin = esp32_reset_pin or "D5"
        elif connect_via == "airlift_featherwing":
            if verbose_logging:
                print("- Connection method = airlift_featherwing")
            esp32_cs_pin = esp32_cs_pin or "D13"
            esp32_ready_pin = esp32_ready_pin or "D11"
            esp32_reset_pin = esp32_reset_pin or "D12"
        elif connect_via == "esp32_external":
            if verbose_logging:
                print("- Connection method = esp32_external")
            esp32_cs_pin = esp32_cs_pin or "D9"
            esp32_ready_pin = esp32_ready_pin or "D10"
            esp32_reset_pin = esp32_reset_pin or "D5"

        if not all([esp32_cs_pin, esp32_ready_pin, esp32_reset_pin]):
            if verbose_logging:
                print("* ESP32 pins not properly specified!")
            return (None, None)

        if verbose_logging:
            print(
                "- configured esp32 pins: cs=board.{}, ready=board.{}, reset=board.{}".format(
                    esp32_cs_pin, esp32_ready_pin, esp32_reset_pin
                )
            )

        esp32_cs = DigitalInOut(getattr(board, esp32_cs_pin))
        esp32_ready = DigitalInOut(getattr(board, esp32_ready_pin))
        esp32_reset = DigitalInOut(getattr(board, esp32_reset_pin))

        if not all([esp32_cs, esp32_ready, esp32_reset]):
            if verbose_logging:
                print("* ESP32 pins not configured correctly!")
            return (None, None)

        spi = busio.SPI(board.SCK, board.MOSI, board.MISO)
        if not spi:
            if verbose_logging:
                print("* Could not create busio.SPI object!")
            return (None, None)

        esp = adafruit_esp32spi.ESP_SPIcontrol(spi, esp32_cs, esp32_ready, esp32_reset)
        if not esp:
            if verbose_logging:
                print("* Could not create adafruit_esp32spi.ESP_SPIcontrol object!")
            return (None, None)

        retry_count = 0

        if esp.status == adafruit_esp32spi.WL_IDLE_STATUS:
            if verbose_logging:
                print(
                    "- Wi-Fi found in idle status [firmware version={}, MAC addr {}]".format(
                        str(esp.firmware_version), [hex(i) for i in esp.MAC_address]
                    )
                )
            while not esp.is_connected:
                try:
                    if verbose_logging:
                        print(
                            "- Trying to connect to Wi-Fi (attempt {} of {})".format(
                                retry_count + 1, max_retries
                            )
                        )
                    esp.connect_AP(wifi_ssid, wifi_password)
                except OSError as exc:
                    if verbose_logging:
                        print("* Could not connect to Wi-Fi: {}!".format(str(exc)))
                    retry_count = retry_count + 1
                    if retry_count > max_retries:
                        if verbose_logging:
                            print("* Giving up!")
                        return (None, None)
                    else:
                        if verbose_logging:
                            print("* Retrying...")
                print(
                    "= Wi-Fi connected to '{}' (RSSI={}, IP={})".format(
                        str(esp.ssid, "utf-8"), esp.rssi, esp.pretty_ip(esp.ip_address)
                    )
                )
                if esp_debug:
                    if verbose_logging:
                        print("- Enabling ESP32SPI Debug Mode")
                    esp._debug = True

        requests.set_socket(socket, esp)
        return (spi, esp)
    finally:
        if verbose_logging:
            print("")
            print("===> setup_wifi done")