wxPython/demo/TreeMixin.py

   1 import wx, wx.lib.customtreectrl, wx.gizmos
   2 try:
   3     import treemixin
   4 except ImportError:
   5     from wx.lib.mixins import treemixin
   6
   7 overview = treemixin.__doc__
   8
   9 class TreeModel(object):
  10     ''' TreeModel holds the domain objects that are shown in the different
  11     tree controls. Each domain object is simply a two-tuple consisting of
  12     a label and a list of child tuples, i.e. (label, [list of child tuples]).
  13     '''
  14     def __init__(self, *args, **kwargs):
  15         self.items = []
  16         self.itemCounter = 0
  17         super(TreeModel, self).__init__(*args, **kwargs)
  18
  19     def GetItem(self, indices):
  20         text, children = 'Hidden root', self.items
  21         for index in indices:
  22             text, children = children[index]
  23         return text, children
  24
  25     def GetText(self, indices):
  26         return self.GetItem(indices)[0]
  27
  28     def GetChildren(self, indices):
  29         return self.GetItem(indices)[1]
  30
  31     def GetChildrenCount(self, indices):
  32         return len(self.GetChildren(indices))
  33
  34     def SetChildrenCount(self, indices, count):
  35         children = self.GetChildren(indices)
  36         while len(children) > count:
  37             children.pop()
  38         while len(children) < count:
  39             children.append(('item %d'%self.itemCounter, []))
  40             self.itemCounter += 1
  41
  42     def MoveItem(self, itemToMoveIndex, newParentIndex):
  43         itemToMove = self.GetItem(itemToMoveIndex)
  44         newParentChildren = self.GetChildren(newParentIndex)
  45         newParentChildren.append(itemToMove)
  46         oldParentChildren = self.GetChildren(itemToMoveIndex[:-1])
  47         oldParentChildren.remove(itemToMove)
  48
  49
  50 class DemoTreeMixin(treemixin.VirtualTree, treemixin.DragAndDrop,
  51                     treemixin.ExpansionState):
  52     def __init__(self, *args, **kwargs):
  53         self.model = kwargs.pop('treemodel')
  54         self.log = kwargs.pop('log')
  55         super(DemoTreeMixin, self).__init__(*args, **kwargs)
  56         self.CreateImageList()
  57
  58     def CreateImageList(self):
  59         size = (16, 16)
  60         self.imageList = wx.ImageList(*size)
  61         for art in wx.ART_FOLDER, wx.ART_FILE_OPEN, wx.ART_NORMAL_FILE:
  62             self.imageList.Add(wx.ArtProvider.GetBitmap(art, wx.ART_OTHER,
  63                                                         size))
  64         self.AssignImageList(self.imageList)
  65
  66     def OnGetItemText(self, indices):
  67         return self.model.GetText(indices)
  68
  69     def OnGetChildrenCount(self, indices):
  70         return self.model.GetChildrenCount(indices)
  71
  72     def OnGetItemFont(self, indices):
  73         # Show how to change the item font. Here we use a small font for
  74         # items that have children and the default font otherwise.
  75         if self.model.GetChildrenCount(indices) > 0:
  76             return wx.SMALL_FONT
  77         else:
  78             return super(DemoTreeMixin, self).OnGetItemFont(indices)
  79
  80     def OnGetItemTextColour(self, indices):
  81         # Show how to change the item text colour. In this case second level
  82         # items are coloured red and third level items are blue. All other
  83         # items have the default text colour.
  84         if len(indices) % 2 == 0:
  85             return wx.RED
  86         elif len(indices) % 3 == 0:
  87             return wx.BLUE
  88         else:
  89             return super(DemoTreeMixin, self).OnGetItemTextColour(indices)
  90
  91     def OnGetItemBackgroundColour(self, indices):
  92         # Show how to change the item background colour. In this case the
  93         # background colour of each third item is green.
  94         if indices[-1] == 2:
  95             return wx.GREEN
  96         else:
  97             return super(DemoTreeMixin,
  98                          self).OnGetItemBackgroundColour(indices)
  99
 100     def OnGetItemImage(self, indices, which):
 101         # Return the right icon depending on whether the item has children.
 102         if which in [wx.TreeItemIcon_Normal, wx.TreeItemIcon_Selected]:
 103             if self.model.GetChildrenCount(indices):
 104                 return 0
 105             else:
 106                 return 2
 107         else:
 108             return 1
 109
 110     def OnDrop(self, dropTarget, dragItem):
 111         dropIndex = self.GetIndexOfItem(dropTarget)
 112         dropText = self.model.GetText(dropIndex)
 113         dragIndex = self.GetIndexOfItem(dragItem)
 114         dragText = self.model.GetText(dragIndex)
 115         self.log.write('drop %s %s on %s %s'%(dragText, dragIndex,
 116             dropText, dropIndex))
 117         self.model.MoveItem(dragIndex, dropIndex)
 118         self.GetParent().RefreshItems()
 119
 120
 121 class VirtualTreeCtrl(DemoTreeMixin, wx.TreeCtrl):
 122     pass
 123
 124
 125 class VirtualTreeListCtrl(DemoTreeMixin, wx.gizmos.TreeListCtrl):
 126     def __init__(self, *args, **kwargs):
 127         kwargs['style'] = wx.TR_DEFAULT_STYLE | wx.TR_FULL_ROW_HIGHLIGHT
 128         super(VirtualTreeListCtrl, self).__init__(*args, **kwargs)
 129         self.AddColumn('Column 0')
 130         self.AddColumn('Column 1')
 131         for art in wx.ART_TIP, wx.ART_WARNING:
 132             self.imageList.Add(wx.ArtProvider.GetBitmap(art, wx.ART_OTHER,
 133                                                         (16, 16)))
 134
 135     def OnGetItemText(self, indices, column=0):
 136         # Return a different label depending on column.
 137         return '%s, column %d'%\
 138             (super(VirtualTreeListCtrl, self).OnGetItemText(indices), column)
 139
 140     def OnGetItemImage(self, indices, which, column=0):
 141         # Also change the image of the other columns when the item has
 142         # children.
 143         if column == 0:
 144             return super(VirtualTreeListCtrl, self).OnGetItemImage(indices,
 145                                                                    which)
 146         elif self.OnGetChildrenCount(indices):
 147             return 4
 148         else:
 149             return 3
 150
 151
 152 class VirtualCustomTreeCtrl(DemoTreeMixin,
 153                             wx.lib.customtreectrl.CustomTreeCtrl):
 154     def __init__(self, *args, **kwargs):
 155         self.checked = {}
 156         kwargs['style'] = wx.TR_HIDE_ROOT | \
 157             wx.TR_HAS_BUTTONS | wx.TR_FULL_ROW_HIGHLIGHT
 158         super(VirtualCustomTreeCtrl, self).__init__(*args, **kwargs)
 159         self.Bind(wx.lib.customtreectrl.EVT_TREE_ITEM_CHECKED,
 160                   self.OnItemChecked)
 161
 162     def OnGetItemType(self, indices):
 163         if len(indices) == 1:
 164             return 1
 165         elif len(indices) == 2:
 166             return 2
 167         else:
 168             return 0
 169
 170     def OnGetItemChecked(self, indices):
 171         return self.checked.get(indices, False)
 172
 173     def OnItemChecked(self, event):
 174         item = event.GetItem()
 175         itemIndex = self.GetIndexOfItem(item)
 176         if self.GetItemType(item) == 2:
 177             # It's a radio item; reset other items on the same level
 178             for nr in range(self.GetChildrenCount(self.GetItemParent(item))):
 179                 self.checked[itemIndex[:-1]+(nr,)] = False
 180         self.checked[itemIndex] = True
 181
 182
 183
 184 class TreeNotebook(wx.Notebook):
 185     def __init__(self, *args, **kwargs):
 186         treemodel = kwargs.pop('treemodel')
 187         log = kwargs.pop('log')
 188         super(TreeNotebook, self).__init__(*args, **kwargs)
 189         self.trees = []
 190         for class_, title in [(VirtualTreeCtrl, 'TreeCtrl'),
 191                               (VirtualTreeListCtrl, 'TreeListCtrl'),
 192                               (VirtualCustomTreeCtrl, 'CustomTreeCtrl')]:
 193             tree = class_(self, treemodel=treemodel, log=log)
 194             self.trees.append(tree)
 195             self.AddPage(tree, title)
 196         self.Bind(wx.EVT_NOTEBOOK_PAGE_CHANGED, self.OnPageChanged)
 197
 198     def OnPageChanged(self, event):
 199         oldTree = self.GetPage(event.OldSelection)
 200         newTree = self.GetPage(event.Selection)
 201         newTree.RefreshItems()
 202         newTree.SetExpansionState(oldTree.GetExpansionState())
 203         event.Skip()
 204
 205     def GetIndicesOfSelectedItems(self):
 206         tree = self.trees[self.GetSelection()]
 207         if tree.GetSelections():
 208             return [tree.GetIndexOfItem(item) for item in tree.GetSelections()]
 209         else:
 210             return [()]
 211
 212     def RefreshItems(self):
 213         tree = self.trees[self.GetSelection()]
 214         tree.RefreshItems()
 215         tree.UnselectAll()
 216
 217
 218 class TestPanel(wx.Panel):
 219     def __init__(self, parent, log):
 220         self.log = log
 221         super(TestPanel, self).__init__(parent)
 222         self.treemodel = TreeModel()
 223         self.CreateControls()
 224         self.LayoutControls()
 225
 226     def CreateControls(self):
 227         self.notebook = TreeNotebook(self, treemodel=self.treemodel,
 228                                      log=self.log)
 229         self.label = wx.StaticText(self, label='Number of children: ')
 230         self.childrenCountCtrl = wx.SpinCtrl(self, value='0', max=10000)
 231         self.button = wx.Button(self, label='Update children')
 232         self.button.Bind(wx.EVT_BUTTON, self.OnEnter)
 233
 234     def LayoutControls(self):
 235         hSizer = wx.BoxSizer(wx.HORIZONTAL)
 236         options = dict(flag=wx.ALIGN_CENTER_VERTICAL|wx.ALL, border=2)
 237         hSizer.Add(self.label, **options)
 238         hSizer.Add(self.childrenCountCtrl, 2, **options)
 239         hSizer.Add(self.button, **options)
 240         sizer = wx.BoxSizer(wx.VERTICAL)
 241         sizer.Add(self.notebook, 1, wx.EXPAND)
 242         sizer.Add(hSizer, 0, wx.EXPAND)
 243         self.SetSizer(sizer)
 244
 245     def OnEnter(self, event):
 246         indicesList = self.notebook.GetIndicesOfSelectedItems()
 247         newChildrenCount = self.childrenCountCtrl.GetValue()
 248         for indices in indicesList:
 249             text = self.treemodel.GetText(indices)
 250             oldChildrenCount = self.treemodel.GetChildrenCount(indices)
 251             self.log.write('%s %s now has %d children (was %d)'%(text, indices,
 252                 newChildrenCount, oldChildrenCount))
 253             self.treemodel.SetChildrenCount(indices, newChildrenCount)
 254         self.notebook.RefreshItems()
 255
 256
 257 def runTest(frame, nb, log):
 258     win = TestPanel(nb, log)
 259     return win
 260
 261
 262 if __name__ == '__main__':
 263     import sys, os, run
 264     run.main(['', os.path.basename(sys.argv[0])] + sys.argv[1:])
 265