Use class icon in class tree (ClassSelector) (#2)

1 files changed, 6 insertions, 2 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/ClassSelector.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/ClassSelector.java
index 488d04ed..b27832be 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/ClassSelector.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/ClassSelector.java
@@ -27,6 +27,7 @@ import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
 import cuchaz.enigma.gui.node.ClassSelectorClassNode;
 import cuchaz.enigma.gui.node.ClassSelectorPackageNode;
+import cuchaz.enigma.gui.util.GuiUtil;
 import cuchaz.enigma.translation.Translator;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.utils.validation.ValidationContext;
@@ -69,10 +70,13 @@ public class ClassSelector extends JTree {
                         }
                 });
 
+                final DefaultTreeCellRenderer renderer = new DefaultTreeCellRenderer();
+                renderer.setLeafIcon(GuiUtil.CLASS_ICON);
+                setCellRenderer(renderer);
+
                 final JTree tree = this;
 
-                final DefaultTreeCellEditor editor = new DefaultTreeCellEditor(tree,
-                                (DefaultTreeCellRenderer) tree.getCellRenderer()) {
+                final DefaultTreeCellEditor editor = new DefaultTreeCellEditor(tree, renderer) {
                         @Override
                         public boolean isCellEditable(EventObject event) {
                                 return isRenamable && !(event instanceof MouseEvent) && super.isCellEditable(event);