diff --git a/src/spectr/spectr.py b/src/spectr/spectr.py index 7f4021c..dab6df4 100644 --- a/src/spectr/spectr.py +++ b/src/spectr/spectr.py @@ -160,7 +160,7 @@ def on_key(self, event) -> None: elif event.key.lower() == "n": event.stop() self.confirm_quit = False - self.query_one("#overlay-text", TopOverlay).clear_prompt() + SpectrApp._get_overlay(self).clear_prompt() self.update_status_bar() def _is_splash_active(self) -> bool: @@ -169,6 +169,16 @@ def _is_splash_active(self) -> bool: self.screen_stack and isinstance(self.screen_stack[-1], SplashScreen) ) + def _get_overlay(self) -> TopOverlay: + """Return the overlay widget for the active screen.""" + stack = getattr(self, "screen_stack", None) + if stack: + try: + return stack[-1].query_one("#overlay-text", TopOverlay) + except Exception: # noqa: BLE001 + pass + return self.query_one("#overlay-text", TopOverlay) + def _prepend_open_positions(self) -> None: """Ensure open position symbols are at the start of ``ticker_symbols``.""" try: @@ -372,7 +382,7 @@ async def on_mount(self, event: events.Mount) -> None: await self.push_screen(SplashScreen(id="splash"), wait_for_dismiss=False) self.refresh() - overlay = self.query_one("#overlay-text", TopOverlay) + overlay = SpectrApp._get_overlay(self) self.voice_agent._on_speech_start = lambda: self.call_from_thread( overlay.start_voice_animation ) @@ -400,7 +410,7 @@ async def on_mount(self, event: events.Mount) -> None: # self.update_status_bar() if self.args.broker == "robinhood" and self.args.real_trades: - self.query_one("#overlay-text", TopOverlay).flash_message( + SpectrApp._get_overlay(self).flash_message( "Robinhood does NOT support PAPER TRADING!", style="bold red" ) log.debug("starting consumer task") @@ -472,7 +482,7 @@ def _poll_one_symbol(self, symbol: str, quote: dict | None = None, position=None log.error("[poll] signal error: %s", traceback.format_exc()) def _flash_error() -> None: - self.query_one("#overlay-text", TopOverlay).flash_message( + SpectrApp._get_overlay(self).flash_message( f"Strategy error: {exc}", style="bold red", ) @@ -788,14 +798,12 @@ async def action_quit(self): """Prompt the user for confirmation before quitting.""" if not self.confirm_quit: self.confirm_quit = True - self.query_one("#overlay-text", TopOverlay).show_prompt( - "Quit Y/N?", style="bold red" - ) + SpectrApp._get_overlay(self).show_prompt("Quit Y/N?", style="bold red") return # Second ESC cancels the prompt self.confirm_quit = False - self.query_one("#overlay-text", TopOverlay).clear_prompt() + SpectrApp._get_overlay(self).clear_prompt() self.update_status_bar() # ------------ Action Functions ------------- @@ -972,14 +980,14 @@ def action_ask_agent(self) -> None: if self._voice_worker and self._voice_worker.is_running: self.voice_agent.stop() self._voice_worker.cancel() - self.query_one("#overlay-text", TopOverlay).clear_prompt() + SpectrApp._get_overlay(self).clear_prompt() self.update_status_bar() return self._voice_worker = self.run_worker(self._ask_agent, thread=True) def _ask_agent(self) -> None: """Run the voice assistant and display errors in the overlay.""" - overlay = self.query_one("#overlay-text", TopOverlay) + overlay = SpectrApp._get_overlay(self) self.call_from_thread(overlay.show_prompt, "Listening...") try: self.voice_agent.listen_and_answer() @@ -1097,7 +1105,7 @@ def action_toggle_portfolio(self) -> None: # -------------- def update_view(self, symbol: str): - self.query_one("#overlay-text", TopOverlay).symbol = symbol + SpectrApp._get_overlay(self).symbol = symbol df = self.df_cache.get(symbol) if df is not None and not self.is_backtest: @@ -1122,7 +1130,7 @@ def update_status_bar(self): else: auto_trade_state = f"Auto-Trades: [BOLD RED]DISABLED[/BOLD RED] {live_icon}" - overlay = self.query_one("#overlay-text", TopOverlay) + overlay = SpectrApp._get_overlay(self) overlay.symbol = self.ticker_symbols[self.active_symbol_index] overlay.live_icon = live_icon if self.strategy_name: @@ -1135,7 +1143,7 @@ def update_status_bar(self): ) def flash_message(self, msg: str): - overlay = self.query_one("#overlay-text", TopOverlay) + overlay = SpectrApp._get_overlay(self) overlay.flash_message(f"ORDER FAILED: {msg}", 10) def set_real_trades(self, enabled: bool) -> None: @@ -1231,7 +1239,7 @@ def add_symbol(self, symbol: str) -> list[str]: self.args.symbols = self.ticker_symbols self.df_cache.setdefault(sym, pd.DataFrame()) cache.save_symbols_cache(self.ticker_symbols) - self.query_one("#overlay-text", TopOverlay).flash_message( + SpectrApp._get_overlay(self).flash_message( f"Added {sym}", duration=5.0, style="bold green" ) return self.ticker_symbols @@ -1246,7 +1254,7 @@ def remove_symbol(self, symbol: str) -> list[str]: "If we remove it from the watchlist we could miss a sell signal." ) self.voice_agent.say(msg) - self.query_one("#overlay-text", TopOverlay).flash_message( + SpectrApp._get_overlay(self).flash_message( f"Failed to remove {sym}, has open position!", duration=6.0, style="bold red", @@ -1261,7 +1269,7 @@ def remove_symbol(self, symbol: str) -> list[str]: cache.save_symbols_cache(self.ticker_symbols) if self.ticker_symbols: self.update_view(self.ticker_symbols[self.active_symbol_index]) - self.query_one("#overlay-text", TopOverlay).flash_message( + SpectrApp._get_overlay(self).flash_message( f"Removed {sym} from watchlist.", duration=5.0, style="bold yellow", @@ -1288,13 +1296,13 @@ async def on_backtest_submit(self, form: dict) -> None: Runs in a thread to keep the UI responsive. """ if self.strategy_class is None: - self.query_one("#overlay-text", TopOverlay).flash_message( + SpectrApp._get_overlay(self).flash_message( "No active strategy", style="bold red" ) return try: log.debug("Backtest starting...") - overlay = self.query_one("#overlay-text", TopOverlay) + overlay = SpectrApp._get_overlay(self) overlay.update_status("Running backtest...") symbol = form["symbol"] starting_cash = float(form["cash"]) @@ -1383,7 +1391,7 @@ async def on_backtest_submit(self, form: dict) -> None: ) except Exception as exc: - self.query_one("#overlay-text", TopOverlay).flash_message( + SpectrApp._get_overlay(self).flash_message( f"Back-test error: {exc}", style="bold red" ) log.error("Back-test error: %s", traceback.format_exc())