tui: Work around changes in Textual's bindings API

In Textual 0.61, `App.namespace_bindings` was removed in favor of
`Screen.active_bindings`. Update our implementation to work for both,
since we still support Textual versions below 0.61.

Signed-off-by: Matt Wozniski <[email protected]>

news/597.bugfix.rst

@@ -0,0 +1 @@
+Fix dynamic toggling between descriptions like "Pause" vs "Unpause" or "Show" vs "Hide" in the footer of the live-mode TUI and tree reporter. This was broken by changes introduced in Textual 0.61.

src/memray/reporters/tree.py

@@ -13,9 +13,11 @@
 from typing import Iterator
 from typing import Optional
 from typing import Tuple
+from typing import Union
 
 from rich.style import Style
 from rich.text import Text
+from textual import binding
 from textual import work
 from textual.app import App
 from textual.app import ComposeResult
@@ -42,6 +44,10 @@
 from memray.reporters.frame_tools import is_frame_interesting
 from memray.reporters.tui import _filename_to_module_name
 
+BindingsType = Union[
+    Dict[str, "binding.ActiveBinding"], Dict[str, Tuple[DOMNode, Binding]]
+]
+
 MAX_STACKS = int(sys.getrecursionlimit() // 2.5)
 
 StackElement = Tuple[str, str, int]
@@ -382,6 +388,34 @@ def redraw_footer(self) -> None:
         self.app.query_one(Footer).highlight_key = "q"
         self.app.query_one(Footer).highlight_key = None
 
+    def update_key_description(
+        self, bindings: BindingsType, key: str, description: str
+    ) -> None:
+        # In Textual 0.61, `App.namespace_bindings` was removed in favor of
+        # `Screen.active_bindings`. The two have a slightly different
+        # interface: a 2 item `tuple` was updated to a 3 item `namedtuple`.
+        # The `BindingsType` type alias shows the two possible structures.
+        # This implementation works for both, since we still support Textual
+        # versions below 0.61.
+        val = bindings[key]
+        binding = replace(val[1], description=description)
+        if type(val) is tuple:
+            bindings[key] = val[:1] + (binding,) + val[2:]  # type: ignore
+        else:
+            bindings[key] = val._replace(binding=binding)  # type: ignore
+
+    def rewrite_bindings(self, bindings: BindingsType) -> None:
+        if self.import_system_filter is not None:
+            self.update_key_description(bindings, "i", "Show import system")
+        if self.uninteresting_filter is not None:
+            self.update_key_description(bindings, "u", "Show uninteresting")
+
+    @property
+    def active_bindings(self) -> Dict[str, "binding.ActiveBinding"]:
+        bindings = super().active_bindings.copy()
+        self.rewrite_bindings(bindings)
+        return bindings
+
 
 class TreeApp(App[None]):
     def __init__(
@@ -395,23 +429,13 @@ def __init__(
     def on_mount(self) -> None:
         self.push_screen(self.tree_screen)
 
-    @property
-    def namespace_bindings(self) -> Dict[str, Tuple[DOMNode, Binding]]:
-        bindings = super().namespace_bindings.copy()
-        if self.import_system_filter is not None:
-            node, binding = bindings["i"]
-            bindings["i"] = (
-                node,
-                replace(binding, description="Show import system"),
-            )
-        if self.uninteresting_filter is not None:
-            node, binding = bindings["u"]
-            bindings["u"] = (
-                node,
-                replace(binding, description="Show uninteresting"),
-            )
-
-        return bindings
+    if hasattr(App, "namespace_bindings"):
+        # Removed in Textual 0.61
+        @property
+        def namespace_bindings(self) -> Dict[str, Tuple[DOMNode, Binding]]:
+            bindings = super().namespace_bindings.copy()  # type: ignore[misc]
+            self.tree_screen.rewrite_bindings(bindings)
+            return bindings  # type: ignore[no-any-return]
 
 
 @functools.lru_cache(maxsize=None)

src/memray/reporters/tui.py

@@ -18,12 +18,14 @@
 from typing import Optional
 from typing import Set
 from typing import Tuple
+from typing import Union
 from typing import cast
 
 from rich.markup import escape
 from rich.segment import Segment
 from rich.style import Style
 from rich.text import Text
+from textual import binding
 from textual import events
 from textual import log
 from textual.app import App
@@ -49,6 +51,10 @@
 from memray import SocketReader
 from memray._memray import size_fmt
 
+BindingsType = Union[
+    Dict[str, "binding.ActiveBinding"], Dict[str, Tuple[DOMNode, Binding]]
+]
+
 MAX_MEMORY_RATIO = 0.95
 
 
@@ -668,6 +674,35 @@ def redraw_footer(self) -> None:
         self.app.query_one(Footer).highlight_key = "q"
         self.app.query_one(Footer).highlight_key = None
 
+    def update_key_description(
+        self, bindings: BindingsType, key: str, description: str
+    ) -> None:
+        # In Textual 0.61, `App.namespace_bindings` was removed in favor of
+        # `Screen.active_bindings`. The two have a slightly different
+        # interface: a 2 item `tuple` was updated to a 3 item `namedtuple`.
+        # The `BindingsType` type alias shows the two possible structures.
+        # This implementation works for both, since we still support Textual
+        # versions below 0.61.
+        val = bindings[key]
+        binding = dataclasses.replace(val[1], description=description)
+        if type(val) is tuple:
+            bindings[key] = val[:1] + (binding,) + val[2:]  # type: ignore
+        else:
+            bindings[key] = val._replace(binding=binding)  # type: ignore
+
+    def rewrite_bindings(self, bindings: BindingsType) -> None:
+        if "space" in bindings and bindings["space"][1].description == "Pause":
+            if self.paused:
+                self.update_key_description(bindings, "space", "Unpause")
+            elif self.disconnected:
+                del bindings["space"]
+
+    @property
+    def active_bindings(self) -> Dict[str, Any]:
+        bindings = super().active_bindings.copy()
+        self.rewrite_bindings(bindings)
+        return bindings
+
 
 class UpdateThread(threading.Thread):
     def __init__(self, app: App[None], reader: SocketReader) -> None:
@@ -770,22 +805,11 @@ def on_snapshot_fetched(self, message: SnapshotFetched) -> None:
     def on_resize(self, event: events.Resize) -> None:
         self.set_class(0 <= event.size.width < 81, "narrow")
 
-    @property
-    def namespace_bindings(self) -> Dict[str, Tuple[DOMNode, Binding]]:
-        bindings = super().namespace_bindings.copy()
-
-        if (
-            "space" in bindings
-            and bindings["space"][1].description == "Pause"
-            and self.tui
-        ):
-            if self.tui.paused:
-                node, binding = bindings["space"]
-                bindings["space"] = (
-                    node,
-                    dataclasses.replace(binding, description="Unpause"),
-                )
-            elif self.tui.disconnected:
-                del bindings["space"]
-
-        return bindings
+    if hasattr(App, "namespace_bindings"):
+        # Removed in Textual 0.61
+        @property
+        def namespace_bindings(self) -> Dict[str, Tuple[DOMNode, Binding]]:
+            bindings = super().namespace_bindings.copy()  # type: ignore[misc]
+            if self.tui:
+                self.tui.rewrite_bindings(bindings)
+            return bindings  # type: ignore[no-any-return]