Updating Textual widgets in place

Question

I have built the following app that will show a text description. The user then selects the correct category (by pressing one of the SelectButtons) for this description and presses the "next" button (ActionButton). In doing so, the selected cateogry should be stored against the description, the stats should update to include the most recent selection, and the description should be updated to the next one to be categorised.

from abc import ABC, abstractmethod
from enum import Enum

from pyfiglet import Figlet
from rich import box
from rich.align import Align
from rich.panel import Panel
from rich.text import Text
from textual import events
from textual.app import App
from textual.reactive import Reactive
from textual.views import GridView
from textual.widget import Widget
from textual.widgets import Button, ButtonPressed, Header

LOREM = (
    "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore "
    "magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo "
    "consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. "
    "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."
)


class Category(str, Enum):
    CATEGORY_1 = "CATEGORY_1"
    CATEGORY_2 = "CATEGORY_2"
    CATEGORY_3 = "CATEGORY_3"
    CATEGORY_4 = "CATEGORY_4"
    CATEGORY_5 = "CATEGORY_5"
    CATEGORY_6 = "CATEGORY_6"
    CATEGORY_7 = "CATEGORY_7"
    CATEGORY_8 = "CATEGORY_8"
    CATEGORY_9 = "CATEGORY_9"
    CATEGORY_10 = "CATEGORY_10"
    CATEGORY_11 = "CATEGORY_11"
    CATEGORY_12 = "CATEGORY_12"
    CATEGORY_13 = "CATEGORY_13"
    CATEGORY_14 = "CATEGORY_14"
    CATEGORY_15 = "CATEGORY_15"
    CATEGORY_16 = "CATEGORY_16"


class TextBox(Widget):
    def __init__(self, name: str, text: str, is_titled: bool, align: str) -> None:
        super().__init__(name=name)
        self.text = text
        self.is_titled = is_titled
        self.align = align

    def render(self) -> Panel:
        if self.align == "left":
            align = Align.left
        elif self.align == "center":
            align = Align.center
        else:
            align = Align.right
        return Panel(
            align(self.text, vertical="middle"),
            title=self.name if self.is_titled else None,
            border_style="white",
            box=box.ROUNDED,
        )


class ButtonBase(Button, ABC):
    mouse_over: Reactive[bool] = Reactive(False)
    has_focus: Reactive[bool] = Reactive(False)

    def __init__(self, name: str, label: str, is_titled: bool, has_focus: bool) -> None:
        super().__init__(name=name, label=label)
        self.is_titled = is_titled
        self.has_focus = has_focus

    @abstractmethod
    def render(self) -> Panel:
        ...

    def on_enter(self) -> None:
        self.mouse_over = True

    def on_leave(self) -> None:
        self.mouse_over = False

    async def on_blur(self, event: events.Blur) -> None:
        self.has_focus = False

    async def on_click(self, event: events.Click) -> None:
        self.has_focus = not self.has_focus


class SelectButton(ButtonBase):
    def render(self) -> Panel:
        if self.has_focus:
            style = "green"
            border_style = "green"
            box_style = box.HEAVY
        elif self.mouse_over:
            style = "white"
            border_style = "green"
            box_style = box.HEAVY
        else:
            style = "white"
            border_style = "white" if self.is_titled else "blue"
            box_style = box.ROUNDED

        return Panel(
            Align.center(self.label, vertical="middle"),
            style=style,
            border_style=border_style,
            box=box_style,
            title=self.name if self.is_titled else None,
        )


class ActionButton(ButtonBase):
    clicked: bool = False

    def render(self) -> Panel:
        if self.mouse_over:
            style = "green"
            border_style = "green"
            box_style = box.HEAVY
        else:
            style = "white"
            border_style = "white" if self.is_titled is not None else "blue"
            box_style = box.ROUNDED

        return Panel(
            Align.center(self.label, vertical="middle"),
            style=style,
            border_style=border_style,
            box=box_style,
            title=self.name if self.is_titled else None,
        )

    async def on_click(self, event: events.Click) -> None:
        self.clicked = True


class FigletText:
    def __init__(self, text: str) -> None:
        self.text = text

    def __str__(self) -> str:
        font = Figlet(font="doh")
        return str(Text(font.renderText(self.text).rstrip("\n"), style="bold"))


class DescriptionGrid(GridView):
    def __init__(self, description_text: str, description_cleansed_text: str, current_category_text: Category) -> None:
        super().__init__()
        self.description_text = description_text
        self.description_cleansed_text = description_cleansed_text
        self.current_category_text = current_category_text

    def on_mount(self) -> None:
        self.grid.add_column("col-0", fraction=3)
        self.grid.add_column("col-1", fraction=1)
        self.grid.add_row("row-0", size=10)
        self.grid.add_row("row-1", size=10)
        self.grid.add_areas(
            description="col-0-start|col-0-end,row-0-start|row-0-end",
            description_cleansed="col-0-start|col-0-end,row-1-start|row-1-end",
            current_category="col-1-start|col-1-end,row-0-start|row-1-end",
        )
        self.grid.place(
            description=TextBox(name="Description", text=self.description_text, is_titled=True, align="left"),
            description_cleansed=TextBox(
                name="Description Cleansed", text=self.description_cleansed_text, is_titled=True, align="left"
            ),
            current_category=SelectButton(
                name="Current Category", label=self.current_category_text, is_titled=True, has_focus=True
            ),
        )


class CategoryGrid(GridView):
    def on_mount(self) -> None:
        self.grid.add_column("col", repeat=4)
        self.grid.add_row("row", repeat=4, size=5)
        self.grid.place(
            *[
                SelectButton(name=f"Category {i}", label=category, is_titled=False, has_focus=False)
                for i, category in enumerate(Category)
            ]
        )


class ControlGrid(GridView):
    def on_mount(self) -> None:
        self.grid.add_column("col-0")
        self.grid.add_column("col-1")
        self.grid.add_column("col-2")
        self.grid.add_column("col-3")
        self.grid.add_row("row-0", size=5)
        self.grid.add_row("row-1", size=5)
        self.grid.add_row("row-2", size=5)
        self.grid.add_row("row-3", size=5)
        self.grid.add_areas(
            back="col-0-start|col-0-end,row-0-start|row-3-end",
            next="col-3-start|col-3-end,row-0-start|row-3-end",
            stat_0="col-1-start|col-1-end,row-0-start|row-0-end",
            stat_1="col-2-start|col-2-end,row-0-start|row-0-end",
            stat_2="col-1-start|col-1-end,row-1-start|row-1-end",
            stat_3="col-2-start|col-2-end,row-1-start|row-1-end",
            stat_4="col-1-start|col-1-end,row-2-start|row-2-end",
            stat_5="col-2-start|col-2-end,row-2-start|row-2-end",
            stat_6="col-1-start|col-1-end,row-3-start|row-3-end",
            stat_7="col-2-start|col-2-end,row-3-start|row-3-end",
        )
        self.grid.place(
            back=ActionButton(name="Back", label=str(FigletText("<")), is_titled=True, has_focus=False),
            next=ActionButton(name="Next", label=str(FigletText(">")), is_titled=True, has_focus=False),
            stat_0=TextBox(name="Stat 0", text="0", is_titled=True, align="center"),
            stat_1=TextBox(name="Stat 1", text="0", is_titled=True, align="center"),
            stat_2=TextBox(name="Stat 2", text="0", is_titled=True, align="center"),
            stat_3=TextBox(name="Stat 3", text="0", is_titled=True, align="center"),
            stat_4=TextBox(name="Stat 4", text="0", is_titled=True, align="center"),
            stat_5=TextBox(name="Stat 5", text="0", is_titled=True, align="center"),
            stat_6=TextBox(name="Stat 6", text="0", is_titled=True, align="center"),
            stat_7=TextBox(name="Stat 7", text="0", is_titled=True, align="center"),
        )


class LabelCorrector(App):
    description_text: str
    description_cleansed_text: str
    current_category_text: Category
    selected_category: str

    async def on_mount(self, event: events.Mount) -> None:
        self.description_text = LOREM
        self.description_cleansed_text = LOREM[: len(LOREM) // 4]
        self.current_category_text = Category.CATEGORY_1
        await self.build_grid()

    async def handle_button_pressed(self, message: ButtonPressed) -> None:
        if isinstance(message.sender, ActionButton) and message.sender.name == "Next":
            self.log(f"Storing {self.selected_category} -> {self.description_text}")
            self.description_text = "New description text."
            self.description_cleansed_text = "New description cleansed text."
            self.current_category_text = Category.CATEGORY_2
            self.clear_grid()
            await self.build_grid()

        elif isinstance(message.sender, SelectButton) and message.sender.name.startswith("Category"):
            self.selected_category = Category(message.sender.label)

    def clear_grid(self) -> None:
        self.view.layout.docks.clear()
        self.view.widgets.clear()

    async def build_grid(self) -> None:
        header = Header(style="white")
        header.layout_size = 3
        description_grid = DescriptionGrid(
            description_text=self.description_text,
            description_cleansed_text=self.description_cleansed_text,
            current_category_text=self.current_category_text,
        )
        category_grid = CategoryGrid()
        control_grid = ControlGrid()
        self.selected_category = self.current_category_text
        await self.view.dock(header, description_grid, category_grid, control_grid)


LabelCorrector.run(log="textual.log")

However, I am having the following 3 problems:

1. When the app starts, the "Current Category" button does, and should have focus (green). However, when I then press any of the other category buttons, the focus should switch from the "Current Category" button to the selected button. Instead, both buttons seem to gain focus as both are now green. If I manually deselect and reselect the "Current Category" button by pressing it twice, then select another category, the focus switches fine as expected.
2. When the "Next" button is pressed, it should not take focus away from the selected category.
3. When the "Next button is pressed I need a way to store the selected category against the description in the LabelCorrector.handle_button_pressed() method. Then, the description text and current category text should be updated, and then current category button should be in focus by default again.

---

Edit:

I have managed to solve the second 2 problems and I've edited the code above to reflect this. However, I cannot figure out why the current category button doesn't lose focus when clicking on another button.