diff --git a/bootstrap/tui/widgets/locals_panel.py b/bootstrap/tui/widgets/locals_panel.py index fa86e14..531a7af 100644 --- a/bootstrap/tui/widgets/locals_panel.py +++ b/bootstrap/tui/widgets/locals_panel.py @@ -1,5 +1,7 @@ -from typing import Any, Dict +from typing import Any, Dict, List +import numpy as np +import torch from textual.app import ComposeResult from textual.containers import VerticalScroll from textual.widgets import Collapsible, Pretty, Static @@ -11,10 +13,17 @@ def compose(self) -> ComposeResult: async def add_locals(self, locals: Dict[str, Any]) -> None: for name, val in locals.items(): + title = f"{name}: {type(val).__name__}" + if isinstance(val, torch.Tensor) or isinstance(val, np.ndarray): + title += f", {val.shape}" + elif isinstance(val, List): + title += f", {len(val)} elements" + elif isinstance(val, Dict): + title += f", {len(val)} keys" await self.query_one("#locals").mount( Collapsible( Pretty(val), - title=name, + title=title, ) )