Java JTree custom icon changes back to default icon while editing tree cell

In my Code I customize a DefaultTreeCellRenderer to give my nodes a custom icon.

setCellRenderer(new DefaultTreeCellRenderer() {
            private static final long serialVersionUID = -2839238218110688876L;
            private ImageIcon icon = myIcon;

            @Override
            public Component getTreeCellRendererComponent(JTree tree, Object value, boolean selected, boolean expanded, boolean isLeaf, int row, boolean focused) {
                super.getTreeCellRendererComponent(tree, value, selected, expanded, isLeaf, row, hasFocus);
                Component c = super.getTreeCellRendererComponent(tree, value, selected, expanded, isLeaf, row, focused);
                setIcon(icon);
                return c;
            }
        });

All that worked just fine until I made my tree editable so that the user can edit the node identifier. Now the icon switches back to default while editing nodes. It seems like the editing tree cell renderer is not the same as the one I wrote. Does anyone know how to fix this?

Solution

Using the setLeafIcon(Icon), setOpenIcon(Icon) and setClosedIcon(Icon) methods of DefaultTreeCellRenderer might be simpler.

import java.awt.*;
import javax.swing.*;
import javax.swing.tree.*;

public class TreeIconTest {
  private final Icon icon = new Icon() {
    @Override public void paintIcon(Component c, Graphics g, int x, int y) {
      Graphics2D g2 = (Graphics2D) g.create();
      g2.translate(x, y);
      g2.setPaint(Color.RED);
      g2.fillRect(1, 1, getIconWidth() - 2, getIconHeight() - 2);
      g2.dispose();
    }
    @Override public int getIconWidth() {
      return 16;
    }
    @Override public int getIconHeight() {
      return 16;
    }
  };
  public JComponent makeUI() {
    JTree tree1 = new JTree();
    tree1.setEditable(true);
    tree1.setCellRenderer(new DefaultTreeCellRenderer() {
      @Override
      public Component getTreeCellRendererComponent(
          JTree tree, Object value, boolean selected, boolean expanded,
          boolean isLeaf, int row, boolean focused) {
        JLabel c = (JLabel) super.getTreeCellRendererComponent(
            tree, value, selected, expanded, isLeaf, row, focused);
        c.setIcon(icon);
        return c;
      }
    });

    JTree tree2 = new JTree();
    tree2.setEditable(true);
    DefaultTreeCellRenderer renderer2 = new DefaultTreeCellRenderer();
    renderer2.setOpenIcon(icon);
    renderer2.setClosedIcon(icon);
    renderer2.setLeafIcon(icon);
    tree2.setCellRenderer(renderer2);

    JTree tree3 = new JTree();
    tree3.setEditable(true);
    tree3.setCellRenderer(new DefaultTreeCellRenderer() {
      @Override public Icon getOpenIcon() {
        return icon;
      }
      @Override public Icon getClosedIcon() {
        return icon;
      }
      @Override public Icon getLeafIcon() {
        return icon;
      }
    });
    tree3.setCellEditor(new DefaultTreeCellEditor(
        tree3, (DefaultTreeCellRenderer) tree3.getCellRenderer()));

    JPanel p = new JPanel(new GridLayout(1, 3));
    p.add(new JScrollPane(tree1));
    p.add(new JScrollPane(tree2));
    p.add(new JScrollPane(tree3));
    return p;
  }
  public static void main(String... args) {
    EventQueue.invokeLater(() -> {
      JFrame f = new JFrame();
      f.setDefaultCloseOperation(WindowConstants.EXIT_ON_CLOSE);
      f.getContentPane().add(new TreeIconTest().makeUI());
      f.setSize(320, 240);
      f.setLocationRelativeTo(null);
      f.setVisible(true);
    });
  }
}