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 <mwozniski@bloomberg.net>

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/_textual_hacks.py

@@ -0,0 +1,26 @@
+import dataclasses
+from typing import Dict
+from typing import Tuple
+from typing import Union
+
+from textual import binding
+from textual.binding import Binding
+from textual.dom import DOMNode
+
+# 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 `Bindings` type alias shows the two possible structures.
+# The `update_key_description` implementation works for both,
+# since we still support Textual versions below 0.61.
+
+Bindings = Union[Dict[str, "binding.ActiveBinding"], Dict[str, Tuple[DOMNode, Binding]]]
+
+
+def update_key_description(bindings: Bindings, key: str, description: str) -> None:
+    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

src/memray/reporters/tree.py

@@ -4,7 +4,6 @@
 import sys
 from dataclasses import dataclass
 from dataclasses import field
-from dataclasses import replace
 from typing import IO
 from typing import Any
 from typing import Callable
@@ -16,6 +15,7 @@
 
 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
@@ -37,6 +37,8 @@
 
 from memray import AllocationRecord
 from memray._memray import size_fmt
+from memray.reporters._textual_hacks import Bindings
+from memray.reporters._textual_hacks import update_key_description
 from memray.reporters.frame_tools import is_cpython_internal
 from memray.reporters.frame_tools import is_frame_from_import_system
 from memray.reporters.frame_tools import is_frame_interesting
@@ -382,6 +384,18 @@ def redraw_footer(self) -> None:
         self.app.query_one(Footer).highlight_key = "q"
         self.app.query_one(Footer).highlight_key = None
 
+    def rewrite_bindings(self, bindings: Bindings) -> None:
+        if self.import_system_filter is not None:
+            update_key_description(bindings, "i", "Show import system")
+        if self.uninteresting_filter is not None:
+            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 +409,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

@@ -1,5 +1,4 @@
 import contextlib
-import dataclasses
 import os
 import pathlib
 import sys
@@ -48,6 +47,8 @@
 from memray import AllocationRecord
 from memray import SocketReader
 from memray._memray import size_fmt
+from memray.reporters._textual_hacks import Bindings
+from memray.reporters._textual_hacks import update_key_description
 
 MAX_MEMORY_RATIO = 0.95
 
@@ -668,6 +669,19 @@ def redraw_footer(self) -> None:
         self.app.query_one(Footer).highlight_key = "q"
         self.app.query_one(Footer).highlight_key = None
 
+    def rewrite_bindings(self, bindings: Bindings) -> None:
+        if "space" in bindings and bindings["space"][1].description == "Pause":
+            if self.paused:
+                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 +784,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]