--- a/IDEFrame.py	Thu May 16 23:01:03 2013 +0200
+++ b/IDEFrame.py	Fri May 17 01:06:58 2013 +0200
@@ -1443,7 +1443,7 @@
         # Refresh treectrl items according to project infos
         infos = self.Controler.GetProjectInfos()
         root = self.ProjectTree.GetRootItem()
-        if not root.IsOk():
+        if root is None or not root.IsOk():
             root = self.ProjectTree.AddRoot(infos["name"])
         self.GenerateProjectTreeBranch(root, infos)
         self.ProjectTree.Expand(root)
@@ -1457,12 +1457,12 @@
     def ResetSelectedItem(self):
         self.SelectedItem = None
 
-    def GenerateProjectTreeBranch(self, root, infos):
+    def GenerateProjectTreeBranch(self, root, infos, item_alone=False):
         to_delete = []
         item_name = infos["name"]
         if infos["type"] in ITEMS_UNEDITABLE:
             if len(infos["values"]) == 1:
-                return self.GenerateProjectTreeBranch(root, infos["values"][0])
+                return self.GenerateProjectTreeBranch(root, infos["values"][0], True)
             item_name = _(item_name)
         self.ProjectTree.SetItemText(root, item_name)
         self.ProjectTree.SetPyData(root, infos)
@@ -1470,7 +1470,10 @@
         self.ProjectTree.SetItemBackgroundColour(root, highlight_colours[0])
         self.ProjectTree.SetItemTextColour(root, highlight_colours[1])
         if infos["type"] == ITEM_POU:
-            self.ProjectTree.SetItemImage(root, self.TreeImageDict[self.Controler.GetPouBodyType(infos["name"])])
+            self.ProjectTree.SetItemImage(root, 
+                self.TreeImageDict[self.Controler.GetPouBodyType(infos["name"])])
+            if item_alone:
+                self.ProjectTree.SetItemExtraImage(root, self.Controler.GetPouType(infos["name"]))
         elif infos.has_key("icon") and infos["icon"] is not None:
             icon_name = infos["icon"]
             if not self.TreeImageDict.has_key(icon_name):
@@ -1479,19 +1482,16 @@
         elif self.TreeImageDict.has_key(infos["type"]):
             self.ProjectTree.SetItemImage(root, self.TreeImageDict[infos["type"]])      
         
-        if wx.VERSION >= (2, 6, 0):
-            item, root_cookie = self.ProjectTree.GetFirstChild(root)
-        else:
-            item, root_cookie = self.ProjectTree.GetFirstChild(root, 0)
+        item, root_cookie = self.ProjectTree.GetFirstChild(root)
         for values in infos["values"]:
             if values["type"] not in ITEMS_UNEDITABLE or len(values["values"]) > 0:
-                if not item.IsOk():
+                if item is None or not item.IsOk():
                     item = self.ProjectTree.AppendItem(root, "")
                     if wx.Platform != '__WXMSW__':
                         item, root_cookie = self.ProjectTree.GetNextChild(root, root_cookie)
                 self.GenerateProjectTreeBranch(item, values)
                 item, root_cookie = self.ProjectTree.GetNextChild(root, root_cookie)
-        while item.IsOk():
+        while item is not None and item.IsOk():
             to_delete.append(item)
             item, root_cookie = self.ProjectTree.GetNextChild(root, root_cookie)
         for item in to_delete:
@@ -1521,11 +1521,8 @@
 
     def RecursiveProjectTreeItemSelection(self, root, items):
         found = False
-        if wx.VERSION >= (2, 6, 0):
-            item, root_cookie = self.ProjectTree.GetFirstChild(root)
-        else:
-            item, root_cookie = self.ProjectTree.GetFirstChild(root, 0)
-        while item.IsOk() and not found:
+        item, root_cookie = self.ProjectTree.GetFirstChild(root)
+        while item is not None and item.IsOk() and not found:
             item_infos = self.ProjectTree.GetPyData(item)
             if (item_infos["name"].split(":")[-1].strip(), item_infos["type"]) == items[0]:
                 if len(items) == 1:
@@ -1686,7 +1683,7 @@
             event.Skip()
     
     def ProjectTreeItemSelect(self, select_item):
-        if select_item.IsOk():
+        if select_item is not None and select_item.IsOk():
             name = self.ProjectTree.GetItemText(select_item)
             item_infos = self.ProjectTree.GetPyData(select_item)
             if item_infos["type"] in [ITEM_DATATYPE, ITEM_POU,
@@ -1706,7 +1703,7 @@
         if not event.Dragging():
             pt = wx.Point(event.GetX(), event.GetY())
             item, flags = self.ProjectTree.HitTest(pt)
-            if item.IsOk() and flags & wx.TREE_HITTEST_ONITEMLABEL:
+            if item is not None and item.IsOk() and flags & wx.TREE_HITTEST_ONITEMLABEL:
                 item_infos = self.ProjectTree.GetPyData(item)
                 if item != self.LastToolTipItem and self.LastToolTipItem is not None:
                     self.ProjectTree.SetToolTip(None)

--- a/controls/CustomTree.py	Thu May 16 23:01:03 2013 +0200
+++ b/controls/CustomTree.py	Fri May 17 01:06:58 2013 +0200
@@ -16,11 +16,29 @@
 #Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
 import wx
+import wx.lib.agw.customtreectrl as CT
 
-class CustomTree(wx.TreeCtrl):
+from util.BitmapLibrary import GetBitmap
+
+# Customize CustomTreeItem for adding icon on item left
+CT.GenericTreeItem._ExtraImage = None
+
+def SetExtraImage(self, image):
+    self._type = 1
+    self._ExtraImage = image
+CT.GenericTreeItem.SetExtraImage = SetExtraImage
+
+_DefaultGetCurrentCheckedImage = CT.GenericTreeItem.GetCurrentCheckedImage
+def GetCurrentCheckedImage(self):
+    if self._ExtraImage is not None:
+        return self._ExtraImage
+    return _DefaultGetCurrentCheckedImage(self)
+CT.GenericTreeItem.GetCurrentCheckedImage = GetCurrentCheckedImage
+
+class CustomTree(CT.CustomTreeCtrl):
     
     def __init__(self, *args, **kwargs):
-        wx.TreeCtrl.__init__(self, *args, **kwargs)
+        CT.CustomTreeCtrl.__init__(self, *args, **kwargs)
         self.SetBackgroundStyle(wx.BG_STYLE_CUSTOM)
         
         self.BackgroundBitmap = None
@@ -29,18 +47,28 @@
         self.AddMenu = None
         self.Enabled = False
         
-        if wx.Platform == '__WXMSW__':
-            self.Bind(wx.EVT_ERASE_BACKGROUND, self.OnEraseBackground)
-        else:
-            self.Bind(wx.EVT_PAINT, self.OnPaint)
-            self.Bind(wx.EVT_SIZE, self.OnResize)
-            self.Bind(wx.EVT_SCROLL, self.OnScroll)
+        self.Bind(wx.EVT_SCROLLWIN, self.OnScroll)
         self.Bind(wx.EVT_LEFT_UP, self.OnLeftUp)
     
     def SetBackgroundBitmap(self, bitmap, align):
         self.BackgroundBitmap = bitmap
         self.BackgroundAlign = align
     
+    def SetImageListCheck(self, sizex, sizey, imglist=None):
+        CT.CustomTreeCtrl.SetImageListCheck(self, sizex, sizey, imglist=None)
+        
+        self.ExtraImages = {}
+        for image in ["function", "functionBlock", "program"]:
+            self.ExtraImages[image] = self._imageListCheck.Add(GetBitmap(image.upper()))
+    
+    def SetItemExtraImage(self, item, bitmap):
+        image = self.ExtraImages.get(bitmap)
+        if image is not None:
+            dc = wx.ClientDC(self)
+            item.SetExtraImage(image)
+            self.CalculateSize(item, dc)
+            self.RefreshLine(item)
+    
     def SetAddMenu(self, add_menu):
         self.AddMenu = add_menu
     
@@ -67,19 +95,6 @@
         
         return wx.Rect(x, y, bitmap_size[0], bitmap_size[1])
     
-    def RefreshBackground(self, refresh_base=False):
-        dc = wx.ClientDC(self)
-        dc.Clear()
-        
-        bitmap_rect = self.GetBitmapRect()
-        dc.DrawBitmap(self.BackgroundBitmap, bitmap_rect.x, bitmap_rect.y)
-        
-        if refresh_base:
-            self.Refresh(False)
-    
-    def OnEraseBackground(self, event):
-        self.RefreshBackground(True)
-    
     def OnLeftUp(self, event):
         if self.Enabled:
             pos = event.GetPosition()
@@ -92,13 +107,18 @@
         event.Skip()
 
     def OnScroll(self, event):
-        self.RefreshBackground(True)
+        wx.CallAfter(self.Refresh)
         event.Skip()
 
-    def OnResize(self, event):
-        self.RefreshBackground(True)
-        event.Skip()
+    def OnSize(self, event):
+        CT.CustomTreeCtrl.OnSize(self, event)
+        self.Refresh()
     
     def OnPaint(self, event):
-        self.RefreshBackground()
-        event.Skip()
\ No newline at end of file
+        dc = wx.ClientDC(self)
+        dc.Clear()
+        
+        bitmap_rect = self.GetBitmapRect()
+        dc.DrawBitmap(self.BackgroundBitmap, bitmap_rect.x, bitmap_rect.y)
+        
+        CT.CustomTreeCtrl.OnPaint(self, event)
\ No newline at end of file