wxPython/demo/TreeMixin.py

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