Address PR 22 review feedback

src/pyldraw3_tui/app.py

@@ -87,7 +87,7 @@ def compose(self) -> ComposeResult:
     def on_mount(self) -> None:
         """Classify the data source and start loading the catalog."""
         self.theme = DARK_THEME
-        model_view = self.query_one("#model-view", ModelView)
+        model_view = self.query_one("#model-view", expect_type=ModelView)
         model_view.set_source(self.source)
         if self._model_path is not None:
             model_view.load_model(self._model_path)
@@ -112,7 +112,7 @@ def _start_catalog_load(self, state: SourceState) -> None:
                 "Building the parts index — the first load can take a while…",
                 timeout=10,
             )
-        self.query_one("#catalog-view", CatalogView).loading = True
+        self.query_one("#catalog-view", expect_type=CatalogView).loading = True
         self._catalog_worker = self._load_catalog()
 
     @property
@@ -147,24 +147,24 @@ def _load_catalog(self) -> None:
         self.call_from_thread(self._catalog_ready, parts)
 
     def _catalog_failed(self, reason: str) -> None:
-        self.query_one("#catalog-view", CatalogView).loading = False
+        self.query_one("#catalog-view", expect_type=CatalogView).loading = False
         self.notify(f"Could not load the catalog: {reason}", severity="error")
 
     def _catalog_ready(self, parts: Parts) -> None:
         self.parts = parts
         self.search_index = SearchIndex.from_catalog(parts.catalog)
-        catalog_view = self.query_one("#catalog-view", CatalogView)
+        catalog_view = self.query_one("#catalog-view", expect_type=CatalogView)
         catalog_view.set_parts(parts, self.search_index)
         catalog_view.loading = False
-        self.query_one("#model-view", ModelView).set_parts(parts)
-        if self.query_one("#main-tabs", TabbedContent).active == "catalog":
-            self.query_one("#parts-list", PartsList).focus()
+        self.query_one("#model-view", expect_type=ModelView).set_parts(parts)
+        if self.query_one("#main-tabs", expect_type=TabbedContent).active == "catalog":
+            self.query_one("#parts-list", expect_type=PartsList).focus()
         self.notify(f"Catalog loaded: {len(parts.catalog.by_code)} parts")
 
     # ------------------------------------------------------------- helpers
 
     def _selected_entry(self) -> CatalogEntry | None:
-        entry = self.query_one("#catalog-view", CatalogView).selected_entry
+        entry = self.query_one("#catalog-view", expect_type=CatalogView).selected_entry
         if entry is None:
             self.notify("No part selected.", severity="warning")
         return entry
@@ -181,8 +181,8 @@ def _copy_choice(self, choice: str | None) -> None:
 
     def focus_part_in_catalog(self, code: str) -> None:
         """Switch to the catalog tab and select a part by code."""
-        self.query_one("#main-tabs", TabbedContent).active = "catalog"
-        self.query_one("#catalog-view", CatalogView).focus_part(code)
+        self.query_one("#main-tabs", expect_type=TabbedContent).active = "catalog"
+        self.query_one("#catalog-view", expect_type=CatalogView).focus_part(code)
 
     def help_sections(self) -> BindingSections:
         """Collect binding tables for the help screen, grouped by owner."""
@@ -200,7 +200,7 @@ def help_sections(self) -> BindingSections:
 
     def action_show_tab(self, tab: str) -> None:
         """Activate the Catalog or Model tab."""
-        self.query_one("#main-tabs", TabbedContent).active = tab
+        self.query_one("#main-tabs", expect_type=TabbedContent).active = tab
         if tab == "model":
             self.query_one("#piece-table", expect_type=PieceTable).focus()
 
@@ -256,8 +256,8 @@ def action_open_model_prompt(self) -> None:
     def _open_model(self, path: str | None) -> None:
         if not path:
             return
-        self.query_one("#main-tabs", TabbedContent).active = "model"
-        self.query_one("#model-view", ModelView).load_model(path)
+        self.query_one("#main-tabs", expect_type=TabbedContent).active = "model"
+        self.query_one("#model-view", expect_type=ModelView).load_model(path)
 
     async def action_regenerate_index(self) -> None:
         """Delete the persistent index and rebuild it from the library."""
@@ -279,7 +279,7 @@ def action_copy_bom_json(self) -> None:
         self._copy_bom(as_json=True)
 
     def _copy_bom(self, *, as_json: bool) -> None:
-        rows = self.query_one("#model-view", ModelView).bom_rows
+        rows = self.query_one("#model-view", expect_type=ModelView).bom_rows
         if not rows:
             self.notify("No BOM to copy — open a model first.", severity="warning")
             return

src/pyldraw3_tui/screens/catalog.py

@@ -70,32 +70,34 @@ def set_parts(self, parts: Parts, search: SearchIndex) -> None:
         """Populate every pane once the catalog is loaded."""
         self._parts = parts
         self._search = search
-        self.query_one("#category-tree", CategoryTree).set_catalog(parts.catalog)
-        self.query_one("#part-detail", PartDetail).set_parts(parts)
+        self.query_one("#category-tree", expect_type=CategoryTree).set_catalog(
+            parts.catalog
+        )
+        self.query_one("#part-detail", expect_type=PartDetail).set_parts(parts)
         self._refresh_list()
 
     @property
     def selected_entry(self) -> CatalogEntry | None:
         """The catalog entry under the parts-list cursor."""
-        return self.query_one("#parts-list", PartsList).highlighted_entry
+        return self.query_one("#parts-list", expect_type=PartsList).highlighted_entry
 
     def focus_part(self, code: str) -> None:
         """Reset scope/filter and move the cursor to a part code."""
         if self._parts is None:
             return
         self._scope = CategoryScope()
         self._filter = ""
-        filter_box = self.query_one("#filter-box", FilterBox)
+        filter_box = self.query_one("#filter-box", expect_type=FilterBox)
         with filter_box.prevent(FilterBox.Changed):
             filter_box.value = ""
         self._refresh_list()
-        parts_list = self.query_one("#parts-list", PartsList)
+        parts_list = self.query_one("#parts-list", expect_type=PartsList)
         parts_list.focus_code(code)
         parts_list.focus()
 
     def action_focus_filter(self) -> None:
         """Move focus into the live filter box."""
-        self.query_one("#filter-box", FilterBox).focus()
+        self.query_one("#filter-box", expect_type=FilterBox).focus()
 
     @on(CategorySelected)
     def _category_selected(self, event: CategorySelected) -> None:
@@ -112,14 +114,14 @@ def _filter_changed(self, event: FilterChanged) -> None:
     @on(PartHighlighted)
     def _part_highlighted(self, event: PartHighlighted) -> None:
         event.stop()
-        self.query_one("#part-detail", PartDetail).show_entry(event.entry)
+        self.query_one("#part-detail", expect_type=PartDetail).show_entry(event.entry)
 
     def _refresh_list(self) -> None:
         if self._parts is None:
             return
         entries = self._scope.entries(self._parts.catalog)
         if self._filter and self._search is not None:
             entries = tuple(self._search.filter(self._filter, within=entries))
-        parts_list = self.query_one("#parts-list", PartsList)
+        parts_list = self.query_one("#parts-list", expect_type=PartsList)
         parts_list.set_entries(entries)
         parts_list.border_title = f"{self._scope.label} ({len(entries)})"

src/pyldraw3_tui/screens/model.py

@@ -193,7 +193,7 @@ def set_parts(self, parts: Parts) -> None:
     @property
     def bom_rows(self) -> list[BomRow]:
         """The currently displayed bill-of-materials rows."""
-        return self.query_one("#bom-table", BomTable).rows_data
+        return self.query_one("#bom-table", expect_type=BomTable).rows_data
 
     def load_model(self, path: Path | str) -> None:
         """Open a model file and show its root model."""
@@ -214,21 +214,21 @@ def load_model(self, path: Path | str) -> None:
         title = Path(path).name
         if (description := model.description) is not None:
             title = f"{title} — {description}"
-        self.query_one("#model-title", Static).update(title)
-        select = self.query_one("#submodel-select", Select)
+        self.query_one("#model-title", expect_type=Static).update(title)
+        select = self.query_one("#submodel-select", expect_type=Select)
         root_label = model.name or Path(path).name
         options = [(f"(root) {root_label}", ROOT_KEY)]
         options += [(name, name) for name in model.submodels]
         with select.prevent(Select.Changed):
             select.set_options(options)
             select.value = ROOT_KEY
-        mode_select = self.query_one("#view-mode-select", Select)
+        mode_select = self.query_one("#view-mode-select", expect_type=Select)
         with mode_select.prevent(Select.Changed):
             mode_select.value = MODEL_MODE
         self._sync_mode_ui()
         self._render_model()
         self._render_issues()
-        self.query_one("#piece-table", PieceTable).focus()
+        self.query_one("#piece-table", expect_type=PieceTable).focus()
 
     @on(Select.Changed, "#view-mode-select")
     def _view_mode_changed(self, event: Select.Changed) -> None:
@@ -268,23 +268,33 @@ def _show_error(self, message: str) -> None:
         self._selected_instruction_step = 1
         self._step_occurrence_counts = None
         self._step_counts_section = None
-        submodel_select = self.query_one("#submodel-select", Select)
+        submodel_select = self.query_one("#submodel-select", expect_type=Select)
         with submodel_select.prevent(Select.Changed):
             submodel_select.set_options([])
         self._view_mode = MODEL_MODE
-        mode_select = self.query_one("#view-mode-select", Select)
+        mode_select = self.query_one("#view-mode-select", expect_type=Select)
         with mode_select.prevent(Select.Changed):
             mode_select.value = MODEL_MODE
         self._sync_mode_ui()
         self.add_class("errored")
-        self.query_one("#model-error", Static).update(f"[bold red]Error:[/] {message}")
-        self.query_one("#model-title", Static).update("No model open")
-        self.query_one("#piece-table", PieceTable).set_occurrences([], self._parts)
-        self.query_one("#stats-panel", StatsPanel).update("Model has no pieces.")
-        self.query_one("#instruction-details", InstructionDetails).update("")
-        self.query_one("#pli-table", BomTable).set_rows([], self._parts)
-        self.query_one("#bom-table", BomTable).set_rows([], self._parts)
-        self.query_one("#directives-table", DirectivesTable).set_directives([])
+        self.query_one("#model-error", expect_type=Static).update(
+            f"[bold red]Error:[/] {message}"
+        )
+        self.query_one("#model-title", expect_type=Static).update("No model open")
+        self.query_one("#piece-table", expect_type=PieceTable).set_occurrences(
+            [], self._parts
+        )
+        self.query_one("#stats-panel", expect_type=StatsPanel).update(
+            "Model has no pieces."
+        )
+        self.query_one("#instruction-details", expect_type=InstructionDetails).update(
+            ""
+        )
+        self.query_one("#pli-table", expect_type=BomTable).set_rows([], self._parts)
+        self.query_one("#bom-table", expect_type=BomTable).set_rows([], self._parts)
+        self.query_one("#directives-table", expect_type=DirectivesTable).set_directives(
+            []
+        )
         self._clear_instruction_selectors()
 
     def _selected_model(self) -> Model | None:
@@ -296,7 +306,7 @@ def _selected_model(self) -> Model | None:
             return self._model.submodel_view(self._selected_key)
         except UnknownSubmodelError:
             self._selected_key = ROOT_KEY
-            select = self.query_one("#submodel-select", Select)
+            select = self.query_one("#submodel-select", expect_type=Select)
             with select.prevent(Select.Changed):
                 select.value = ROOT_KEY
             return self._model
@@ -315,16 +325,16 @@ def _render_whole_model(self) -> None:
         occurrences = list(
             model.iter_occurrences(include_steps=len(steps) > 1),
         )
-        self.query_one("#piece-table", PieceTable).set_occurrences(
+        self.query_one("#piece-table", expect_type=PieceTable).set_occurrences(
             occurrences,
             self._parts,
         )
-        self.query_one("#stats-panel", StatsPanel).show_occurrences(
+        self.query_one("#stats-panel", expect_type=StatsPanel).show_occurrences(
             occurrences,
             self._parts,
             steps=len(steps),
         )
-        self.query_one("#bom-table", BomTable).set_rows(
+        self.query_one("#bom-table", expect_type=BomTable).set_rows(
             bill_of_materials(model, parts=self._parts),
             self._parts,
         )
@@ -347,25 +357,28 @@ def _render_instruction_step(self) -> None:
         )
         # pyldraw3 guarantees the cumulative expansion is the ordered
         # concatenation of each step's added occurrences; if that ever
-        # drifts, blank step labels beat crashing the render.
+        # drifts, falling back to each occurrence's own step number
+        # beats crashing the render.
         step_numbers = labels if len(labels) == len(occurrences) else None
-        self.query_one("#piece-table", PieceTable).set_occurrences(
+        self.query_one("#piece-table", expect_type=PieceTable).set_occurrences(
             occurrences,
             self._parts,
             step_numbers=step_numbers,
         )
-        self.query_one("#stats-panel", StatsPanel).show_occurrences(
+        self.query_one("#stats-panel", expect_type=StatsPanel).show_occurrences(
             occurrences,
             self._parts,
             steps=len(section.steps),
         )
-        self.query_one("#instruction-details", InstructionDetails).show_step(
+        self.query_one(
+            "#instruction-details", expect_type=InstructionDetails
+        ).show_step(
             step,
             total_steps=len(section.steps),
             added_count=counts[step.number - 1],
             cumulative_count=len(occurrences),
         )
-        self.query_one("#pli-table", BomTable).set_rows(
+        self.query_one("#pli-table", expect_type=BomTable).set_rows(
             step.added_bill_of_materials(
                 parts=self._parts,
                 expand_submodels=True,
@@ -374,15 +387,15 @@ def _render_instruction_step(self) -> None:
             self._parts,
             title="Parts list",
         )
-        self.query_one("#bom-table", BomTable).set_rows(
+        self.query_one("#bom-table", expect_type=BomTable).set_rows(
             step.cumulative_bill_of_materials(
                 parts=self._parts,
                 expand_submodels=True,
                 respect_lpub=True,
             ),
             self._parts,
         )
-        self.query_one("#directives-table", DirectivesTable).set_directives(
+        self.query_one("#directives-table", expect_type=DirectivesTable).set_directives(
             step.directives,
         )
 
@@ -414,7 +427,9 @@ def _build_instruction_document(self, *, reset: bool) -> None:
         if reset or self._selected_instruction_section not in section_names:
             self._selected_instruction_section = document.root.name
             self._selected_instruction_step = 1
-        section_select = self.query_one("#instruction-section-select", Select)
+        section_select = self.query_one(
+            "#instruction-section-select", expect_type=Select
+        )
         options = [
             (
                 f"(root) {section.name}" if section.is_root else section.name,
@@ -435,7 +450,7 @@ def _set_step_options(self, *, reset: bool) -> None:
         numbers = {step.number for step in section.steps}
         if reset or self._selected_instruction_step not in numbers:
             self._selected_instruction_step = section.steps[0].number
-        step_select = self.query_one("#instruction-step-select", Select)
+        step_select = self.query_one("#instruction-step-select", expect_type=Select)
         with step_select.prevent(Select.Changed):
             step_select.set_options(
                 [(f"Step {step.number}", step.number) for step in section.steps],
@@ -465,20 +480,22 @@ def _instruction_selection(
         return section, fallback
 
     def _clear_instruction_selectors(self) -> None:
-        section_select = self.query_one("#instruction-section-select", Select)
+        section_select = self.query_one(
+            "#instruction-section-select", expect_type=Select
+        )
         with section_select.prevent(Select.Changed):
             section_select.set_options([])
         self._clear_step_selector()
 
     def _clear_step_selector(self) -> None:
-        step_select = self.query_one("#instruction-step-select", Select)
+        step_select = self.query_one("#instruction-step-select", expect_type=Select)
         with step_select.prevent(Select.Changed):
             step_select.set_options([])
 
     def _sync_mode_ui(self) -> None:
         instructions = self._view_mode == INSTRUCTIONS_MODE
         self.set_class(instructions, "instructions")
-        tabs = self.query_one("#model-tabs", TabbedContent)
+        tabs = self.query_one("#model-tabs", expect_type=TabbedContent)
         if not instructions and tabs.active in _INSTRUCTION_ONLY_TABS:
             tabs.active = "tab-pieces"
         for pane_id in _INSTRUCTION_ONLY_TABS:
@@ -489,7 +506,7 @@ def action_toggle_instructions(self) -> None:
         next_mode = (
             MODEL_MODE if self._view_mode == INSTRUCTIONS_MODE else INSTRUCTIONS_MODE
         )
-        self.query_one("#view-mode-select", Select).value = next_mode
+        self.query_one("#view-mode-select", expect_type=Select).value = next_mode
 
     def action_previous_instruction_step(self) -> None:
         """Select the preceding step in the current instruction section."""
@@ -509,7 +526,7 @@ def _move_instruction_step(self, offset: int) -> None:
             max(self._selected_instruction_step + offset, section.steps[0].number),
             section.steps[-1].number,
         )
-        self.query_one("#instruction-step-select", Select).value = target
+        self.query_one("#instruction-step-select", expect_type=Select).value = target
 
     def _render_issues(self) -> None:
         """Validate the open file and show the issues, whole file at once.
@@ -535,6 +552,8 @@ def _render_issues(self) -> None:
             else []
         )
         combined_issues = [*issues, *instruction_issues]
-        self.query_one("#issues-table", IssuesTable).set_issues(combined_issues)
-        tabs = self.query_one("#model-tabs", TabbedContent)
+        self.query_one("#issues-table", expect_type=IssuesTable).set_issues(
+            combined_issues
+        )
+        tabs = self.query_one("#model-tabs", expect_type=TabbedContent)
         tabs.get_tab("tab-issues").label = f"Issues ({len(combined_issues)})"

src/pyldraw3_tui/screens/setup.py

@@ -84,7 +84,7 @@ def _start(self, event: Button.Pressed) -> None:
         self._download_and_generate()
 
     def _status(self, message: str) -> None:
-        self.query_one("#setup-status", Static).update(message)
+        self.query_one("#setup-status", expect_type=Static).update(message)
 
     @work(thread=True, exclusive=True, group="setup")
     def _download_and_generate(self) -> None:
@@ -110,7 +110,7 @@ def progress(event: ProgressEvent) -> None:
 
     def _failed(self, reason: str) -> None:
         self.remove_class("working")
-        self.query_one("#setup-download", Button).disabled = False
+        self.query_one("#setup-download", expect_type=Button).disabled = False
         self._status(f"[red]Setup failed:[/] {reason}")
 
     def _succeeded(self) -> None: