From 02ed90c633b7a0f53d26a64d8adb41591320d521 Mon Sep 17 00:00:00 2001
From: Dave Pearson <davep@davep.org>
Date: Fri, 28 Apr 2023 10:22:11 +0100
Subject: [PATCH] Add TreeNode.tree

Currently, in the various TreeNode messages, and the handlers you'd write to
handle them, there's no way to easily know *which* tree sent the message and
so which tree the node belongs to.

This commit adds public access to the tree reference to the nodes, so that
in an event handler the developer can check the tree involved in the event.

See #2413.
---
 CHANGELOG.md                     |  6 ++++++
 src/textual/widgets/_tree.py     |  5 +++++
 tests/tree/test_tree_messages.py | 36 ++++++++++++++++++++++++--------
 3 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ead1b47ff..96df25999 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## Unreleased
+
+### Added
+
+- Added `TreeNode.tree` as a read-only public attribute https://github.com/Textualize/textual/issues/2413
+
 ## [0.22.0] - 2023-04-27
 
 ### Fixed
diff --git a/src/textual/widgets/_tree.py b/src/textual/widgets/_tree.py
index 6fb5d5b2e..91ffe794e 100644
--- a/src/textual/widgets/_tree.py
+++ b/src/textual/widgets/_tree.py
@@ -125,6 +125,11 @@ class TreeNode(Generic[TreeDataType]):
         self._selected_ = False
         self._updates += 1
 
+    @property
+    def tree(self) -> Tree[TreeDataType]:
+        """The tree that this node is attached to."""
+        return self._tree
+
     @property
     def children(self) -> TreeNodes[TreeDataType]:
         """The child nodes of a TreeNode."""
diff --git a/tests/tree/test_tree_messages.py b/tests/tree/test_tree_messages.py
index 359dd1aa0..ef742319a 100644
--- a/tests/tree/test_tree_messages.py
+++ b/tests/tree/test_tree_messages.py
@@ -16,18 +16,26 @@ class TreeApp(App[None]):
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        self.messages: list[str] = []
+        self.messages: list[tuple[str, str]] = []
 
     def compose(self) -> ComposeResult:
         """Compose the child widgets."""
-        yield MyTree("Root")
+        yield MyTree("Root", id="test-tree")
 
     def on_mount(self) -> None:
         self.query_one(MyTree).root.add("Child")
         self.query_one(MyTree).focus()
 
-    def record(self, event: Message) -> None:
-        self.messages.append(event.__class__.__name__)
+    def record(
+        self,
+        event: Tree.NodeSelected[None]
+        | Tree.NodeExpanded[None]
+        | Tree.NodeCollapsed[None]
+        | Tree.NodeHighlighted[None],
+    ) -> None:
+        self.messages.append(
+            (event.__class__.__name__, event.node.tree.id or "Unknown")
+        )
 
     def on_tree_node_selected(self, event: Tree.NodeSelected[None]) -> None:
         self.record(event)
@@ -47,7 +55,10 @@ async def test_tree_node_selected_message() -> None:
     async with TreeApp().run_test() as pilot:
         await pilot.press("enter")
         await pilot.pause()
-        assert pilot.app.messages == ["NodeExpanded", "NodeSelected"]
+        assert pilot.app.messages == [
+            ("NodeExpanded", "test-tree"),
+            ("NodeSelected", "test-tree"),
+        ]
 
 
 async def test_tree_node_selected_message_no_auto() -> None:
@@ -56,7 +67,7 @@ async def test_tree_node_selected_message_no_auto() -> None:
         pilot.app.query_one(MyTree).auto_expand = False
         await pilot.press("enter")
         await pilot.pause()
-        assert pilot.app.messages == ["NodeSelected"]
+        assert pilot.app.messages == [("NodeSelected", "test-tree")]
 
 
 async def test_tree_node_expanded_message() -> None:
@@ -64,7 +75,7 @@ async def test_tree_node_expanded_message() -> None:
     async with TreeApp().run_test() as pilot:
         await pilot.press("space")
         await pilot.pause()
-        assert pilot.app.messages == ["NodeExpanded"]
+        assert pilot.app.messages == [("NodeExpanded", "test-tree")]
 
 
 async def test_tree_node_collapsed_message() -> None:
@@ -72,7 +83,10 @@ async def test_tree_node_collapsed_message() -> None:
     async with TreeApp().run_test() as pilot:
         await pilot.press("space", "space")
         await pilot.pause()
-        assert pilot.app.messages == ["NodeExpanded", "NodeCollapsed"]
+        assert pilot.app.messages == [
+            ("NodeExpanded", "test-tree"),
+            ("NodeCollapsed", "test-tree"),
+        ]
 
 
 async def test_tree_node_highlighted_message() -> None:
@@ -80,4 +94,8 @@ async def test_tree_node_highlighted_message() -> None:
     async with TreeApp().run_test() as pilot:
         await pilot.press("enter", "down")
         await pilot.pause()
-        assert pilot.app.messages == ["NodeExpanded", "NodeSelected", "NodeHighlighted"]
+        assert pilot.app.messages == [
+            ("NodeExpanded", "test-tree"),
+            ("NodeSelected", "test-tree"),
+            ("NodeHighlighted", "test-tree"),
+        ]