Files
kiauh/kiauh/components/klipper_firmware/menus/klipper_flash_menu.py
dw-0 01afe1fe77 chore: ruff organize imports
Signed-off-by: Dominik Willner <th33xitus@gmail.com>
2024-05-20 10:52:18 +02:00

443 lines
17 KiB
Python

# ======================================================================= #
# Copyright (C) 2020 - 2024 Dominik Willner <th33xitus@gmail.com> #
# #
# This file is part of KIAUH - Klipper Installation And Update Helper #
# https://github.com/dw-0/kiauh #
# #
# This file may be distributed under the terms of the GNU GPLv3 license #
# ======================================================================= #
import textwrap
import time
from typing import Optional, Type
from components.klipper_firmware.firmware_utils import (
find_firmware_file,
find_uart_device,
find_usb_device_by_id,
find_usb_dfu_device,
get_sd_flash_board_list,
start_flash_process,
)
from components.klipper_firmware.flash_options import (
ConnectionType,
FlashCommand,
FlashMethod,
FlashOptions,
)
from components.klipper_firmware.menus.klipper_flash_error_menu import (
KlipperNoBoardTypesErrorMenu,
KlipperNoFirmwareErrorMenu,
)
from components.klipper_firmware.menus.klipper_flash_help_menu import (
KlipperFlashCommandHelpMenu,
KlipperFlashMethodHelpMenu,
KlipperMcuConnectionHelpMenu,
)
from core.menus import FooterType, Option
from core.menus.base_menu import BaseMenu
from utils.constants import COLOR_CYAN, COLOR_RED, COLOR_YELLOW, RESET_FORMAT
from utils.input_utils import get_number_input
from utils.logger import Logger
# noinspection PyUnusedLocal
# noinspection PyMethodMayBeStatic
class KlipperFlashMethodMenu(BaseMenu):
def __init__(self, previous_menu: Optional[Type[BaseMenu]] = None):
super().__init__()
self.help_menu = KlipperFlashMethodHelpMenu
self.input_label_txt = "Select flash method"
self.footer_type = FooterType.BACK_HELP
self.flash_options = FlashOptions()
def set_previous_menu(self, previous_menu: Optional[Type[BaseMenu]]) -> None:
from core.menus.advanced_menu import AdvancedMenu
self.previous_menu: Type[BaseMenu] = (
previous_menu if previous_menu is not None else AdvancedMenu
)
def set_options(self) -> None:
self.options = {
"1": Option(self.select_regular, menu=False),
"2": Option(self.select_sdcard, menu=False),
}
def print_menu(self) -> None:
header = " [ MCU Flash Menu ] "
subheader = f"{COLOR_YELLOW}ATTENTION:{RESET_FORMAT}"
subline1 = f"{COLOR_YELLOW}Make sure to select the correct method for the MCU!{RESET_FORMAT}"
subline2 = f"{COLOR_YELLOW}Not all MCUs support both methods!{RESET_FORMAT}"
color = COLOR_CYAN
count = 62 - len(color) - len(RESET_FORMAT)
menu = textwrap.dedent(
f"""
/=======================================================\\
| {color}{header:~^{count}}{RESET_FORMAT} |
|-------------------------------------------------------|
| Select the flash method for flashing the MCU. |
| |
| {subheader:<62} |
| {subline1:<62} |
| {subline2:<62} |
|-------------------------------------------------------|
| |
| 1) Regular flashing method |
| 2) Updating via SD-Card Update |
| |
"""
)[1:]
print(menu, end="")
def select_regular(self, **kwargs):
self.flash_options.flash_method = FlashMethod.REGULAR
self.goto_next_menu()
def select_sdcard(self, **kwargs):
self.flash_options.flash_method = FlashMethod.SD_CARD
self.goto_next_menu()
def goto_next_menu(self, **kwargs):
if find_firmware_file():
KlipperFlashCommandMenu(previous_menu=self.__class__).run()
else:
KlipperNoFirmwareErrorMenu().run()
# noinspection PyUnusedLocal
# noinspection PyMethodMayBeStatic
class KlipperFlashCommandMenu(BaseMenu):
def __init__(self, previous_menu: Optional[Type[BaseMenu]] = None):
super().__init__()
self.help_menu = KlipperFlashCommandHelpMenu
self.input_label_txt = "Select flash command"
self.footer_type = FooterType.BACK_HELP
self.flash_options = FlashOptions()
def set_previous_menu(self, previous_menu: Optional[Type[BaseMenu]]) -> None:
self.previous_menu: Type[BaseMenu] = (
previous_menu if previous_menu is not None else KlipperFlashMethodMenu
)
def set_options(self) -> None:
self.options = {
"1": Option(self.select_flash, menu=False),
"2": Option(self.select_serialflash, menu=False),
}
self.default_option = Option(self.select_flash, menu=False)
def print_menu(self) -> None:
menu = textwrap.dedent(
"""
/=======================================================\\
| |
| Which flash command to use for flashing the MCU? |
| 1) make flash (default) |
| 2) make serialflash (stm32flash) |
| |
"""
)[1:]
print(menu, end="")
def select_flash(self, **kwargs):
self.flash_options.flash_command = FlashCommand.FLASH
self.goto_next_menu()
def select_serialflash(self, **kwargs):
self.flash_options.flash_command = FlashCommand.SERIAL_FLASH
self.goto_next_menu()
def goto_next_menu(self, **kwargs):
KlipperSelectMcuConnectionMenu(previous_menu=self.__class__).run()
# noinspection PyUnusedLocal
# noinspection PyMethodMayBeStatic
class KlipperSelectMcuConnectionMenu(BaseMenu):
def __init__(
self, previous_menu: Optional[Type[BaseMenu]] = None, standalone: bool = False
):
super().__init__()
self.previous_menu = previous_menu
self.__standalone = standalone
self.help_menu = KlipperMcuConnectionHelpMenu
self.input_label_txt = "Select connection type"
self.footer_type = FooterType.BACK_HELP
self.flash_options = FlashOptions()
def set_previous_menu(self, previous_menu: Optional[Type[BaseMenu]]) -> None:
self.previous_menu: Type[BaseMenu] = (
previous_menu if previous_menu is not None else KlipperFlashCommandMenu
)
def set_options(self) -> None:
self.options = {
"1": Option(method=self.select_usb, menu=False),
"2": Option(method=self.select_dfu, menu=False),
"3": Option(method=self.select_usb_dfu, menu=False),
}
def print_menu(self) -> None:
header = "Make sure that the controller board is connected now!"
color = COLOR_YELLOW
count = 62 - len(color) - len(RESET_FORMAT)
menu = textwrap.dedent(
f"""
/=======================================================\\
| {color}{header:^{count}}{RESET_FORMAT} |
|-------------------------------------------------------|
| |
| How is the controller board connected to the host? |
| 1) USB |
| 2) UART |
| 3) USB (DFU mode) |
| |
"""
)[1:]
print(menu, end="")
def select_usb(self, **kwargs):
self.flash_options.connection_type = ConnectionType.USB
self.get_mcu_list()
def select_dfu(self, **kwargs):
self.flash_options.connection_type = ConnectionType.UART
self.get_mcu_list()
def select_usb_dfu(self, **kwargs):
self.flash_options.connection_type = ConnectionType.USB_DFU
self.get_mcu_list()
def get_mcu_list(self, **kwargs):
conn_type = self.flash_options.connection_type
if conn_type is ConnectionType.USB:
Logger.print_status("Identifying MCU connected via USB ...")
self.flash_options.mcu_list = find_usb_device_by_id()
elif conn_type is ConnectionType.UART:
Logger.print_status("Identifying MCU possibly connected via UART ...")
self.flash_options.mcu_list = find_uart_device()
elif conn_type is ConnectionType.USB_DFU:
Logger.print_status("Identifying MCU connected via USB in DFU mode ...")
self.flash_options.mcu_list = find_usb_dfu_device()
if len(self.flash_options.mcu_list) < 1:
Logger.print_warn("No MCUs found!")
Logger.print_warn("Make sure they are connected and repeat this step.")
# if standalone is True, we only display the MCUs to the user and return
if self.__standalone and len(self.flash_options.mcu_list) > 0:
Logger.print_ok("The following MCUs were found:", prefix=False)
for i, mcu in enumerate(self.flash_options.mcu_list):
print(f" ● MCU #{i}: {COLOR_CYAN}{mcu}{RESET_FORMAT}")
time.sleep(3)
return
self.goto_next_menu()
def goto_next_menu(self, **kwargs):
KlipperSelectMcuIdMenu(previous_menu=self.__class__).run()
# noinspection PyUnusedLocal
# noinspection PyMethodMayBeStatic
class KlipperSelectMcuIdMenu(BaseMenu):
def __init__(self, previous_menu: Optional[Type[BaseMenu]] = None):
super().__init__()
self.flash_options = FlashOptions()
self.mcu_list = self.flash_options.mcu_list
self.input_label_txt = "Select MCU to flash"
self.footer_type = FooterType.BACK_HELP
def set_previous_menu(self, previous_menu: Optional[Type[BaseMenu]]) -> None:
self.previous_menu: Type[BaseMenu] = (
previous_menu
if previous_menu is not None
else KlipperSelectMcuConnectionMenu
)
def set_options(self) -> None:
self.options = {
f"{i}": Option(self.flash_mcu, False, f"{i}")
for i in range(len(self.mcu_list))
}
def print_menu(self) -> None:
header = "!!! ATTENTION !!!"
header2 = f"[{COLOR_CYAN}List of available MCUs{RESET_FORMAT}]"
color = COLOR_RED
count = 62 - len(color) - len(RESET_FORMAT)
menu = textwrap.dedent(
f"""
/=======================================================\\
| {color}{header:^{count}}{RESET_FORMAT} |
|-------------------------------------------------------|
| Make sure, to select the correct MCU! |
| ONLY flash a firmware created for the respective MCU! |
| |
|{header2:-^64}|
"""
)[1:]
for i, mcu in enumerate(self.mcu_list):
mcu = mcu.split("/")[-1]
menu += f" ● MCU #{i}: {COLOR_CYAN}{mcu}{RESET_FORMAT}\n"
print(menu, end="\n")
def flash_mcu(self, **kwargs):
index = int(kwargs.get("opt_index"))
selected_mcu = self.mcu_list[index]
self.flash_options.selected_mcu = selected_mcu
if self.flash_options.flash_method == FlashMethod.SD_CARD:
KlipperSelectSDFlashBoardMenu(previous_menu=self.__class__).run()
elif self.flash_options.flash_method == FlashMethod.REGULAR:
KlipperFlashOverviewMenu(previous_menu=self.__class__).run()
# noinspection PyUnusedLocal
# noinspection PyMethodMayBeStatic
class KlipperSelectSDFlashBoardMenu(BaseMenu):
def __init__(self, previous_menu: Optional[Type[BaseMenu]] = None):
super().__init__()
self.flash_options = FlashOptions()
self.available_boards = get_sd_flash_board_list()
self.input_label_txt = "Select board type"
def set_previous_menu(self, previous_menu: Optional[Type[BaseMenu]]) -> None:
self.previous_menu: Type[BaseMenu] = (
previous_menu if previous_menu is not None else KlipperSelectMcuIdMenu
)
def set_options(self) -> None:
self.options = {
f"{i}": Option(self.board_select, False, f"{i}")
for i in range(len(self.available_boards))
}
def print_menu(self) -> None:
if len(self.available_boards) < 1:
KlipperNoBoardTypesErrorMenu().run()
else:
menu = textwrap.dedent(
"""
/=======================================================\\
| Please select the type of board that corresponds to |
| the currently selected MCU ID you chose before. |
| |
| The following boards are currently supported: |
|-------------------------------------------------------|
"""
)[1:]
for i, board in enumerate(self.available_boards):
line = f" {i}) {board}"
menu += f"|{line:<55}|\n"
print(menu, end="")
def board_select(self, **kwargs):
board = int(kwargs.get("opt_index"))
self.flash_options.selected_board = self.available_boards[board]
self.baudrate_select()
def baudrate_select(self, **kwargs):
menu = textwrap.dedent(
"""
/=======================================================\\
| If your board is flashed with firmware that connects |
| at a custom baud rate, please change it now. |
| |
| If you are unsure, stick to the default 250000! |
\\=======================================================/
"""
)[1:]
print(menu, end="")
self.flash_options.selected_baudrate = get_number_input(
question="Please set the baud rate",
default=250000,
min_count=0,
allow_go_back=True,
)
KlipperFlashOverviewMenu(previous_menu=self.__class__).run()
# noinspection PyUnusedLocal
# noinspection PyMethodMayBeStatic
class KlipperFlashOverviewMenu(BaseMenu):
def __init__(self, previous_menu: Optional[Type[BaseMenu]] = None):
super().__init__()
self.flash_options = FlashOptions()
self.input_label_txt = "Perform action (default=Y)"
def set_previous_menu(self, previous_menu: Optional[Type[BaseMenu]]) -> None:
self.previous_menu: Type[BaseMenu] = previous_menu
def set_options(self) -> None:
self.options = {
"Y": Option(self.execute_flash, menu=False),
"N": Option(self.abort_process, menu=False),
}
self.default_option = Option(self.execute_flash, menu=False)
def print_menu(self) -> None:
header = "!!! ATTENTION !!!"
color = COLOR_RED
count = 62 - len(color) - len(RESET_FORMAT)
method = self.flash_options.flash_method.value
command = self.flash_options.flash_command.value
conn_type = self.flash_options.connection_type.value
mcu = self.flash_options.selected_mcu
board = self.flash_options.selected_board
baudrate = self.flash_options.selected_baudrate
subheader = f"[{COLOR_CYAN}Overview{RESET_FORMAT}]"
menu = textwrap.dedent(
f"""
/=======================================================\\
| {color}{header:^{count}}{RESET_FORMAT} |
|-------------------------------------------------------|
| Before contuining the flashing process, please check |
| if all parameters were set correctly! Once you made |
| sure everything is correct, start the process. If any |
| parameter needs to be changed, you can go back (B) |
| step by step or abort and start from the beginning. |
|{subheader:-^64}|
"""
)[1:]
menu += f" ● MCU: {COLOR_CYAN}{mcu}{RESET_FORMAT}\n"
menu += f" ● Connection: {COLOR_CYAN}{conn_type}{RESET_FORMAT}\n"
menu += f" ● Flash method: {COLOR_CYAN}{method}{RESET_FORMAT}\n"
menu += f" ● Flash command: {COLOR_CYAN}{command}{RESET_FORMAT}\n"
if self.flash_options.flash_method is FlashMethod.SD_CARD:
menu += f" ● Board type: {COLOR_CYAN}{board}{RESET_FORMAT}\n"
menu += f" ● Baudrate: {COLOR_CYAN}{baudrate}{RESET_FORMAT}\n"
menu += textwrap.dedent(
"""
|-------------------------------------------------------|
| Y) Start flash process |
| N) Abort - Return to Advanced Menu |
"""
)
print(menu, end="")
def execute_flash(self, **kwargs):
start_flash_process(self.flash_options)
Logger.print_info("Returning to MCU Flash Menu in 5 seconds ...")
time.sleep(5)
KlipperFlashMethodMenu().run()
def abort_process(self, **kwargs):
from core.menus.advanced_menu import AdvancedMenu
AdvancedMenu().run()