From d0a74bdf42ee5eaa51f22bad10c710ddeb4f69ba Mon Sep 17 00:00:00 2001
From: gabornemeth <gabor.nemeth.dev@gmail.com>
Date: Wed, 3 Nov 2021 20:55:04 +0100
Subject: [PATCH] type of the TreeView's items can be customized

---
 src/Avalonia.Controls/TreeView.cs             | 14 ++++--
 src/Avalonia.Controls/TreeViewItem.cs         |  8 +++-
 .../TreeViewTests.cs                          | 46 +++++++++++++++++--
 3 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/src/Avalonia.Controls/TreeView.cs b/src/Avalonia.Controls/TreeView.cs
index 78cd22ae329..03f6fd3aaac 100644
--- a/src/Avalonia.Controls/TreeView.cs
+++ b/src/Avalonia.Controls/TreeView.cs
@@ -392,14 +392,22 @@ private void UnsubscribeFromSelectedItems()
         /// <inheritdoc/>
         protected override IItemContainerGenerator CreateItemContainerGenerator()
         {
-            var result = new TreeItemContainerGenerator<TreeViewItem>(
+            var result = CreateTreeItemContainerGenerator();
+            result.Index.Materialized += ContainerMaterialized;
+            return result;
+        }
+
+        protected virtual ITreeItemContainerGenerator CreateTreeItemContainerGenerator() =>
+            CreateTreeItemContainerGenerator<TreeViewItem>();
+
+        protected virtual ITreeItemContainerGenerator CreateTreeItemContainerGenerator<TVItem>() where TVItem: TreeViewItem, new()
+        {
+            return new TreeItemContainerGenerator<TVItem>(
                 this,
                 TreeViewItem.HeaderProperty,
                 TreeViewItem.ItemTemplateProperty,
                 TreeViewItem.ItemsProperty,
                 TreeViewItem.IsExpandedProperty);
-            result.Index.Materialized += ContainerMaterialized;
-            return result;
         }
 
         /// <inheritdoc/>
diff --git a/src/Avalonia.Controls/TreeViewItem.cs b/src/Avalonia.Controls/TreeViewItem.cs
index 8ce258b5464..88fb3cd9f17 100644
--- a/src/Avalonia.Controls/TreeViewItem.cs
+++ b/src/Avalonia.Controls/TreeViewItem.cs
@@ -92,9 +92,13 @@ public int Level
             (ITreeItemContainerGenerator)base.ItemContainerGenerator;
 
         /// <inheritdoc/>
-        protected override IItemContainerGenerator CreateItemContainerGenerator()
+        protected override IItemContainerGenerator CreateItemContainerGenerator() => CreateTreeItemContainerGenerator<TreeViewItem>();
+
+        /// <inheritdoc/>
+        protected virtual ITreeItemContainerGenerator CreateTreeItemContainerGenerator<TVItem>()
+            where TVItem: TreeViewItem, new()
         {
-            return new TreeItemContainerGenerator<TreeViewItem>(
+            return new TreeItemContainerGenerator<TVItem>(
                 this,
                 TreeViewItem.HeaderProperty,
                 TreeViewItem.ItemTemplateProperty,
diff --git a/tests/Avalonia.Controls.UnitTests/TreeViewTests.cs b/tests/Avalonia.Controls.UnitTests/TreeViewTests.cs
index 72ba3ab273a..2169b15cadc 100644
--- a/tests/Avalonia.Controls.UnitTests/TreeViewTests.cs
+++ b/tests/Avalonia.Controls.UnitTests/TreeViewTests.cs
@@ -6,6 +6,7 @@
 using System.Linq;
 using System.Runtime.CompilerServices;
 using Avalonia.Collections;
+using Avalonia.Controls.Generators;
 using Avalonia.Controls.Presenters;
 using Avalonia.Controls.Templates;
 using Avalonia.Data;
@@ -24,7 +25,7 @@ namespace Avalonia.Controls.UnitTests
     public class TreeViewTests
     {
         MouseTestHelper _mouse = new MouseTestHelper();
-        
+
         [Fact]
         public void Items_Should_Be_Created()
         {
@@ -675,7 +676,7 @@ public void Keyboard_Navigation_Should_Move_To_Last_Selected_Node()
                 Assert.Same(node, focus.Current);
             }
         }
-        
+
         [Fact]
         public void Keyboard_Navigation_Should_Not_Crash_If_Selected_Item_Is_not_In_Tree()
         {
@@ -1166,6 +1167,34 @@ public void Clearing_TreeView_Items_Clears_Index()
             Assert.Empty(target.ItemContainerGenerator.Index.Containers);
         }
 
+        [Fact]
+        public void Can_Use_Derived_TreeViewItem()
+        {
+            var tree = CreateTestTreeData();
+            var target = new DerivedTreeViewWithDerivedTreeViewItems
+            {
+                Template = CreateTreeViewTemplate(),
+                Items = tree,
+            };
+
+            ApplyTemplates(target);
+
+            // Verify that all items are DerivedTreeViewItem
+            VerifyItemType(target.ItemContainerGenerator);
+
+            void VerifyItemType(ITreeItemContainerGenerator containerGenerator)
+            {
+                foreach (var container in containerGenerator.Index.Containers)
+                {
+                    var item = Assert.IsType<DerivedTreeViewItem>(container);
+                    if (item.ItemCount > 0)
+                    {
+                        VerifyItemType(item.ItemContainerGenerator);
+                    }
+                }
+            }
+        }
+
         private void ApplyTemplates(TreeView tree)
         {
             tree.ApplyTemplate();
@@ -1376,6 +1405,17 @@ private class DerivedTreeView : TreeView
         {
         }
 
+        private class DerivedTreeViewWithDerivedTreeViewItems : TreeView
+        {
+            protected override ITreeItemContainerGenerator CreateTreeItemContainerGenerator() =>
+                CreateTreeItemContainerGenerator<DerivedTreeViewItem>();
+        }
+
+        private class DerivedTreeViewItem : TreeViewItem
+        {
+            protected override IItemContainerGenerator CreateItemContainerGenerator() => CreateTreeItemContainerGenerator<DerivedTreeViewItem>();
+        }
+
         private class TestDataContext : INotifyPropertyChanged
         {
             private string _selectedItem;
@@ -1398,7 +1438,7 @@ public string SelectedItem
             }
 
             public event PropertyChangedEventHandler PropertyChanged;
-            
+
         }
     }
 }