Subclassing Tree in order to provide custom styling for a BranchTree

Author: StanleyGoldman

src/UnityExtension/Assets/Editor/GitHub.Unity/UI/BranchesView.cs

@@ -39,8 +39,8 @@ class BranchesView : Subview
         [NonSerialized] private int listID = -1;
         [NonSerialized] private BranchesMode targetMode;
 
-        [SerializeField] private Tree treeLocals = new Tree();
-        [SerializeField] private Tree treeRemotes = new Tree();
+        [SerializeField] private BranchesTree treeLocals = new BranchesTree();
+        [SerializeField] private BranchesTree treeRemotes = new BranchesTree();
         [SerializeField] private BranchesMode mode = BranchesMode.Default;
         [SerializeField] private string newBranchName;
         [SerializeField] private Vector2 scroll;
@@ -140,9 +140,9 @@ private void BuildTree()
         {
             if (treeLocals == null)
             {
-                treeLocals = new Tree();
+                treeLocals = new BranchesTree();
 
-                treeRemotes = new Tree();
+                treeRemotes = new BranchesTree();
 
                 UpdateTreeIcons();
             }

src/UnityExtension/Assets/Editor/GitHub.Unity/UI/TreeControl.cs

@@ -9,19 +9,44 @@
 namespace GitHub.Unity
 {
     [Serializable]
-    public class Tree
+    public class BranchesTree: Tree
+    {
+        [NonSerialized] public Texture2D ActiveNodeIcon;
+        [NonSerialized] public Texture2D NodeIcon;
+        [NonSerialized] public Texture2D FolderIcon;
+        [NonSerialized] public Texture2D RootFolderIcon;
+
+        protected override Texture2D GetNodeIcon(TreeNode node)
+        {
+            Texture2D nodeIcon;
+            if (node.IsActive)
+            {
+                nodeIcon = ActiveNodeIcon;
+            }
+            else if (node.IsFolder)
+            {
+                if (node.Level == 1)
+                    nodeIcon = RootFolderIcon;
+                else
+                    nodeIcon = FolderIcon;
+            }
+            else
+            {
+                nodeIcon = NodeIcon;
+            }
+            return nodeIcon;
+        }
+    }
+
+    [Serializable]
+    public abstract class Tree
     {
         public static float ItemHeight { get { return EditorGUIUtility.singleLineHeight; } }
         public static float ItemSpacing { get { return EditorGUIUtility.standardVerticalSpacing; } }
 
         [SerializeField] public Rect Margin = new Rect();
         [SerializeField] public Rect Padding = new Rect();
 
-        [NonSerialized] public Texture2D ActiveNodeIcon;
-        [NonSerialized] public Texture2D NodeIcon;
-        [NonSerialized] public Texture2D FolderIcon;
-        [NonSerialized] public Texture2D RootFolderIcon;
-
         [SerializeField] public GUIStyle FolderStyle;
         [SerializeField] public GUIStyle TreeNodeStyle;
         [SerializeField] public GUIStyle ActiveTreeNodeStyle;
@@ -82,7 +107,7 @@ public void Load(IEnumerable<ITreeData> data, string title)
                 Level = 0,
                 IsFolder = true
             };
-            SetNodeIcons(titleNode);
+            SetNodeIcon(titleNode);
             nodes.Add(titleNode);
 
             foreach (var d in data)
@@ -110,7 +135,7 @@ public void Load(IEnumerable<ITreeData> data, string title)
                             activeNode = node;
                         }
 
-                        SetNodeIcons(node);
+                        SetNodeIcon(node);
 
                         nodes.Add(node);
                         if (isFolder)
@@ -369,32 +394,19 @@ private void Unindent()
             indents.Pop();
         }
 
-        private void SetNodeIcons(TreeNode node)
+        private void SetNodeIcon(TreeNode node)
         {
-            if (node.IsActive)
-            {
-                node.Icon = ActiveNodeIcon;
-            }
-            else if (node.IsFolder)
-            {
-                if (node.Level == 1)
-                    node.Icon = RootFolderIcon;
-                else
-                    node.Icon = FolderIcon;
-            }
-            else
-            {
-                node.Icon = NodeIcon;
-            }
-
+            node.Icon = GetNodeIcon(node);
             node.Load();
         }
 
+        protected abstract Texture2D GetNodeIcon(TreeNode node);
+
         public void LoadNodeIcons()
         {
             foreach (var treeNode in nodes)
             {
-                SetNodeIcons(treeNode);
+                SetNodeIcon(treeNode);
             }
         }
     }