wxPython/wx/lib/mixins/treemixin.py

   1 '''
   2 treemixin.py
   3
   4 This module provides three mixin classes that can be used with tree
   5 controls:
   6
   7 - VirtualTree is a class that, when mixed in with a tree control,
   8   makes the tree control virtual, similar to a ListCtrl in virtual mode.
   9   A virtual tree control builds the tree itself by means of callbacks,
  10   so the programmer is freed from the burden of building the tree herself.
  11
  12 - DragAndDrop is a mixin class that helps with dragging and dropping of
  13   items. The graphical part of dragging and dropping tree items is done by
  14   this mixin class. You only need to implement the OnDrop method that is
  15   called when the drop happens.
  16
  17 - ExpansionState is a mixin that can be queried for the expansion state of
  18   all items in the tree to restore it later.
  19
  20 All mixin classes work with wx.TreeCtrl, wx.gizmos.TreeListCtrl,
  21 and wx.lib.customtree.CustomTreeCtrl. They can be used together or
  22 separately.
  23
  24 The VirtualTree and DragAndDrop mixins force the wx.TR_HIDE_ROOT style.
  25
  26 Author: Frank Niessink <frank@niessink.com>
  27 License: wxWidgets license
  28 Version: 1.0
  29 Date: 15 April 2007
  30
  31 ExpansionState is based on code and ideas from Karsten Hilbert.
  32 Andrea Gavana provided help with the CustomTreeCtrl integration.
  33 '''
  34
  35
  36 import wx, wx.lib.customtreectrl
  37
  38
  39 class TreeAPIHarmonizer(object):
  40     ''' This class attempts to hide the differences in API between the
  41     different tree controls that are part of wxPython. '''
  42
  43     def __callSuper(self, methodName, default, *args, **kwargs):
  44         # If our super class has a method called methodName, call it,
  45         # otherwise return the default value.
  46         superClass = super(TreeAPIHarmonizer, self)
  47         if hasattr(superClass, methodName):
  48             return getattr(superClass, methodName)(*args, **kwargs)
  49         else:
  50             return default
  51
  52     def GetColumnCount(self, *args, **kwargs):
  53         # Only TreeListCtrl has columns, return 0 if we are mixed in
  54         # with another tree control.
  55         return self.__callSuper('GetColumnCount', 0, *args, **kwargs)
  56
  57     def GetItemType(self, *args, **kwargs):
  58         # Only CustomTreeCtrl has different item types, return the
  59         # default item type if we are mixed in with another tree control.
  60         return self.__callSuper('GetItemType', 0, *args, **kwargs)
  61
  62     def SetItemType(self, item, newType):
  63         # CustomTreeCtrl doesn't support changing the item type on the fly,
  64         # so we create a new item and delete the old one. We currently only
  65         # keep the item text, would be nicer to also retain other attributes.
  66         text = self.GetItemText(item)
  67         newItem = self.InsertItem(self.GetItemParent(item), item, text,
  68                                   ct_type=newType)
  69         self.Delete(item)
  70         return newItem
  71
  72     def IsItemChecked(self, *args, **kwargs):
  73         # Only CustomTreeCtrl supports checkable items, return False if
  74         # we are mixed in with another tree control.
  75         return self.__callSuper('IsItemChecked', False, *args, **kwargs)
  76
  77     def GetItemChecked(self, *args, **kwargs):
  78         # For consistency's sake, provide a 'Get' and 'Set' method for
  79         # checkable items.
  80         return self.IsItemChecked(*args, **kwargs)
  81
  82     def SetItemChecked(self, *args, **kwargs):
  83         # For consistency's sake, provide a 'Get' and 'Set' method for
  84         # checkable items.
  85         return self.CheckItem(*args, **kwargs)
  86
  87     def GetMainWindow(self, *args, **kwargs):
  88         # Only TreeListCtrl has a separate main window, return self if we are
  89         # mixed in with another tree control.
  90         return self.__callSuper('GetMainWindow', self, *args, **kwargs)
  91
  92     def GetItemImage(self, item, which=wx.TreeItemIcon_Normal, column=-1):
  93         # CustomTreeCtrl always wants the which argument, so provide it
  94         # TreeListCtr.GetItemImage has a different order of arguments than
  95         # the other tree controls. Hide the differenes.
  96         if self.GetColumnCount():
  97             args = (item, column, which)
  98         else:
  99             args = (item, which)
 100         return super(TreeAPIHarmonizer, self).GetItemImage(*args)
 101
 102     def SetItemImage(self, item, imageIndex, which=wx.TreeItemIcon_Normal,
 103                      column=-1):
 104         # The SetItemImage signature is different for TreeListCtrl and
 105         # other tree controls. This adapter method hides the differences.
 106         if self.GetColumnCount():
 107             args = (item, imageIndex, column, which)
 108         else:
 109             args = (item, imageIndex, which)
 110         super(TreeAPIHarmonizer, self).SetItemImage(*args)
 111
 112     def UnselectAll(self):
 113         # Unselect all items, regardless of whether we are in multiple
 114         # selection mode or not.
 115         if self.HasFlag(wx.TR_MULTIPLE):
 116             super(TreeAPIHarmonizer, self).UnselectAll()
 117         else:
 118             # CustomTreeCtrl Unselect() doesn't seem to work in all cases,
 119             # also invoke UnselectAll just to be sure.
 120             self.Unselect()
 121             super(TreeAPIHarmonizer, self).UnselectAll()
 122
 123     def GetCount(self):
 124         # TreeListCtrl correctly ignores the root item when it is hidden,
 125         # but doesn't count the root item when it is visible
 126         itemCount = super(TreeAPIHarmonizer, self).GetCount()
 127         if self.GetColumnCount() and not self.HasFlag(wx.TR_HIDE_ROOT):
 128             itemCount += 1
 129         return itemCount
 130
 131     def GetSelections(self):
 132         # Always return a list of selected items, regardless of whether
 133         # we are in multiple selection mode or not.
 134         if self.HasFlag(wx.TR_MULTIPLE):
 135             selections = super(TreeAPIHarmonizer, self).GetSelections()
 136         else:
 137             selection = self.GetSelection()
 138             if selection:
 139                 selections = [selection]
 140             else:
 141                 selections = []
 142         # If the root item is hidden, it should never be selected,
 143         # unfortunately, CustomTreeCtrl allows it to be selected.
 144         if self.HasFlag(wx.TR_HIDE_ROOT):
 145             rootItem = self.GetRootItem()
 146             if rootItem and rootItem in selections:
 147                 selections.remove(rootItem)
 148         return selections
 149
 150     def GetFirstVisibleItem(self):
 151         # TreeListCtrl raises an exception or even crashes when invoking
 152         # GetFirstVisibleItem on an empty tree.
 153         if self.GetRootItem():
 154             return super(TreeAPIHarmonizer, self).GetFirstVisibleItem()
 155         else:
 156             return wx.TreeItemId()
 157
 158     def SelectItem(self, item, *args, **kwargs):
 159         # Prevent the hidden root from being selected, otherwise TreeCtrl
 160         # crashes
 161         if self.HasFlag(wx.TR_HIDE_ROOT) and item == self.GetRootItem():
 162             return
 163         else:
 164             return super(TreeAPIHarmonizer, self).SelectItem(item, *args,
 165                                                              **kwargs)
 166
 167     def HitTest(self, *args, **kwargs):
 168         ''' HitTest returns a two-tuple (item, flags) for tree controls
 169         without columns and a three-tuple (item, flags, column) for tree
 170         controls with columns. Our caller can indicate this method to
 171         always return a three-tuple no matter what tree control we're mixed
 172         in with by specifying the optional argument 'alwaysReturnColumn'
 173         to be True. '''
 174         alwaysReturnColumn = kwargs.pop('alwaysReturnColumn', False)
 175         hitTestResult = super(TreeAPIHarmonizer, self).HitTest(*args, **kwargs)
 176         if len(hitTestResult) == 2 and alwaysReturnColumn:
 177             hitTestResult += (0,)
 178         return hitTestResult
 179
 180     def ExpandAll(self, item=None):
 181         # TreeListCtrl wants an item as argument. That's an inconsistency with
 182         # the TreeCtrl API. Also, TreeCtrl doesn't allow invoking ExpandAll
 183         # on a tree with hidden root node, so prevent that.
 184         if self.HasFlag(wx.TR_HIDE_ROOT):
 185             rootItem = self.GetRootItem()
 186             if rootItem:
 187                 child, cookie = self.GetFirstChild(rootItem)
 188                 while child:
 189                     self.ExpandAllChildren(child)
 190                     child, cookie = self.GetNextChild(rootItem, cookie)
 191         else:
 192             try:
 193                 super(TreeAPIHarmonizer, self).ExpandAll()
 194             except TypeError:
 195                 if item is None:
 196                     item = self.GetRootItem()
 197                 super(TreeAPIHarmonizer, self).ExpandAll(item)
 198
 199     def ExpandAllChildren(self, item):
 200         # TreeListCtrl and CustomTreeCtrl don't have ExpandallChildren
 201         try:
 202             super(TreeAPIHarmonizer, self).ExpandAllChildren(item)
 203         except AttributeError:
 204             self.Expand(item)
 205             child, cookie = self.GetFirstChild(item)
 206             while child:
 207                 self.ExpandAllChildren(child)
 208                 child, cookie = self.GetNextChild(item, cookie)
 209
 210
 211 class TreeHelper(object):
 212     ''' This class provides methods that are not part of the API of any
 213     tree control, but are convenient to have available. '''
 214
 215     def GetItemChildren(self, item=None, recursively=False):
 216         ''' Return the children of item as a list. '''
 217         if not item:
 218             item = self.GetRootItem()
 219             if not item:
 220                 return []
 221         children = []
 222         child, cookie = self.GetFirstChild(item)
 223         while child:
 224             children.append(child)
 225             if recursively:
 226                 children.extend(self.GetItemChildren(child, True))
 227             child, cookie = self.GetNextChild(item, cookie)
 228         return children
 229
 230     def GetIndexOfItem(self, item):
 231         ''' Return the index of item. '''
 232         parent = self.GetItemParent(item)
 233         if parent:
 234             parentIndices = self.GetIndexOfItem(parent)
 235             ownIndex = self.GetItemChildren(parent).index(item)
 236             return parentIndices + (ownIndex,)
 237         else:
 238             return ()
 239
 240     def GetItemByIndex(self, index):
 241         ''' Return the item specified by index. '''
 242         item = self.GetRootItem()
 243         for i in index:
 244             children = self.GetItemChildren(item)
 245             item = children[i]
 246         return item
 247
 248
 249 class VirtualTree(TreeAPIHarmonizer, TreeHelper):
 250     ''' This is a mixin class that can be used to allow for virtual tree
 251     controls. It can be mixed in with wx.TreeCtrl, wx.gizmos.TreeListCtrl,
 252     wx.lib.customtree.CustomTreeCtrl.
 253
 254     To use it derive a new class from this class and one of the tree
 255     controls, e.g.:
 256     class MyTree(VirtualTree, wx.TreeCtrl):
 257         ...
 258
 259     VirtualTree uses several callbacks (such as OnGetItemText) to
 260     retrieve information needed to construct the tree and render the
 261     items. To specify what item the callback needs information about,
 262     the callback passes an item index. Whereas for list controls a simple
 263     integer index can be used, for tree controls indicating a specific
 264     item is a little bit more complicated. See below for a more detailed
 265     explanation of the how index works.
 266
 267     Note that VirtualTree forces the wx.TR_HIDE_ROOT style.
 268
 269     In your subclass you *must* override OnGetItemText and
 270     OnGetChildrenCount. These two methods are the minimum needed to
 271     construct the tree and render the item labels. If you want to add
 272     images, change fonts our colours, etc., you need to override the
 273     appropriate OnGetXXX method as well.
 274
 275     About indices: your callbacks are passed a tuple of integers that
 276     identifies the item the VirtualTree wants information about. An
 277     empty tuple, i.e. (), represents the hidden root item.  A tuple with
 278     one integer, e.g. (3,), represents a visible root item, in this case
 279     the fourth one. A tuple with two integers, e.g. (3,0), represents a
 280     child of a visible root item, in this case the first child of the
 281     fourth root item.
 282     '''
 283
 284     def __init__(self, *args, **kwargs):
 285         kwargs['style'] = kwargs.get('style', wx.TR_DEFAULT_STYLE) | \
 286                           wx.TR_HIDE_ROOT
 287         super(VirtualTree, self).__init__(*args, **kwargs)
 288         self.Bind(wx.EVT_TREE_ITEM_EXPANDING, self.OnItemExpanding)
 289         self.Bind(wx.EVT_TREE_ITEM_COLLAPSED, self.OnItemCollapsed)
 290
 291     def OnGetChildrenCount(self, index):
 292         ''' This function *must* be overloaded in the derived class.
 293         It should return the number of child items of the item with the
 294         provided index. If index == () it should return the number of
 295         root items. '''
 296         raise NotImplementedError
 297
 298     def OnGetItemText(self, index, column=0):
 299         ''' This function *must* be overloaded in the derived class. It
 300         should return the string containing the text of the specified
 301         item. '''
 302         raise NotImplementedError
 303
 304     def OnGetItemFont(self, index):
 305         ''' This function may be overloaded in the derived class. It
 306         should return the wx.Font to be used for the specified item. '''
 307         return wx.NullFont
 308
 309     def OnGetItemTextColour(self, index):
 310         ''' This function may be overloaded in the derived class. It
 311         should return the wx.Colour to be used as text colour for the
 312         specified item. '''
 313         return wx.NullColour
 314
 315     def OnGetItemBackgroundColour(self, index):
 316         ''' This function may be overloaded in the derived class. It
 317         should return the wx.Colour to be used as background colour for
 318         the specified item. '''
 319         return wx.NullColour
 320
 321     def OnGetItemImage(self, index, which=wx.TreeItemIcon_Normal, column=0):
 322         ''' This function may be overloaded in the derived class. It
 323         should return the index of the image to be used. Don't forget
 324         to associate an ImageList with the tree control. '''
 325         return -1
 326
 327     def OnGetItemType(self, index):
 328         ''' This function may be overloaded in the derived class, but
 329         that only makes sense when this class is mixed in with a tree
 330         control that supports checkable items, i.e. CustomTreeCtrl.
 331         This method should return whether the item is to be normal (0,
 332         the default), a checkbox (1) or a radiobutton (2).
 333         Note that OnGetItemChecked needs to be implemented as well; it
 334         should return whether the item is actually checked. '''
 335         return 0
 336
 337     def OnGetItemChecked(self, index):
 338         ''' This function may be overloaded in the derived class, but
 339         that only makes sense when this class is mixed in with a tree
 340         control that supports checkable items, i.e. CustomTreeCtrl.
 341         This method should return whether the item is to be checked.
 342         Note that OnGetItemType should return 1 (checkbox) or 2
 343         (radiobutton) for this item. '''
 344         return False
 345
 346     def RefreshItems(self):
 347         ''' Redraws all visible items. '''
 348         rootItem = self.GetRootItem()
 349         if not rootItem:
 350             rootItem = self.AddRoot('Hidden root')
 351         self.RefreshChildrenRecursively(rootItem)
 352
 353     def RefreshItem(self, index):
 354         ''' Redraws the item with the specified index. '''
 355         try:
 356             item = self.GetItemByIndex(index)
 357         except IndexError:
 358             # There's no corresponding item for index, because its parent
 359             # has not been expanded yet.
 360             return
 361         hasChildren = bool(self.OnGetChildrenCount(index))
 362         self.DoRefreshItem(item, index, hasChildren)
 363
 364     def RefreshChildrenRecursively(self, item, itemIndex=None):
 365         ''' Refresh the children of item, reusing as much of the
 366         existing items in the tree as possible. '''
 367         if itemIndex is None:
 368             itemIndex = self.GetIndexOfItem(item)
 369         reusableChildren = self.GetItemChildren(item)
 370         for childIndex in self.ChildIndices(itemIndex):
 371             if reusableChildren:
 372                 child = reusableChildren.pop(0)
 373             else:
 374                 child = self.AppendItem(item, '')
 375             self.RefreshItemRecursively(child, childIndex)
 376         for child in reusableChildren:
 377             self.Delete(child)
 378
 379     def RefreshItemRecursively(self, item, itemIndex):
 380         ''' Refresh the item and its children recursively. '''
 381         hasChildren = bool(self.OnGetChildrenCount(itemIndex))
 382         item = self.DoRefreshItem(item, itemIndex, hasChildren)
 383         # We need to refresh the children when the item is expanded and
 384         # when the item has no children, because in the latter case we
 385         # might have to delete old children from the tree:
 386         if self.IsExpanded(item) or not hasChildren:
 387             self.RefreshChildrenRecursively(item, itemIndex)
 388         self.SetItemHasChildren(item, hasChildren)
 389
 390     def DoRefreshItem(self, item, index, hasChildren):
 391         ''' Refresh one item. '''
 392         item = self.RefreshItemType(item, index)
 393         self.RefreshItemText(item, index)
 394         self.RefreshColumns(item, index)
 395         self.RefreshItemFont(item, index)
 396         self.RefreshTextColour(item, index)
 397         self.RefreshBackgroundColour(item, index)
 398         self.RefreshItemImage(item, index, hasChildren)
 399         self.RefreshCheckedState(item, index)
 400         return item
 401
 402     def RefreshItemText(self, item, index):
 403         self.__refreshAttribute(item, index, 'ItemText')
 404
 405     def RefreshColumns(self, item, index):
 406         for columnIndex in range(1, self.GetColumnCount()):
 407             self.__refreshAttribute(item, index, 'ItemText', columnIndex)
 408
 409     def RefreshItemFont(self, item, index):
 410         self.__refreshAttribute(item, index, 'ItemFont')
 411
 412     def RefreshTextColour(self, item, index):
 413         self.__refreshAttribute(item, index, 'ItemTextColour')
 414
 415     def RefreshBackgroundColour(self, item, index):
 416         self.__refreshAttribute(item, index, 'ItemBackgroundColour')
 417
 418     def RefreshItemImage(self, item, index, hasChildren):
 419         regularIcons = [wx.TreeItemIcon_Normal, wx.TreeItemIcon_Selected]
 420         expandedIcons = [wx.TreeItemIcon_Expanded,
 421                          wx.TreeItemIcon_SelectedExpanded]
 422         # Refresh images in first column:
 423         for icon in regularIcons:
 424             self.__refreshAttribute(item, index, 'ItemImage', icon)
 425         for icon in expandedIcons:
 426             if hasChildren:
 427                 imageIndex = self.OnGetItemImage(index, icon)
 428             else:
 429                 imageIndex = -1
 430             if self.GetItemImage(item, icon) != imageIndex or imageIndex == -1:
 431                 self.SetItemImage(item, imageIndex, icon)
 432         # Refresh images in remaining columns, if any:
 433         for columnIndex in range(1, self.GetColumnCount()):
 434             for icon in regularIcons:
 435                 self.__refreshAttribute(item, index, 'ItemImage', icon,
 436                                         columnIndex)
 437
 438     def RefreshItemType(self, item, index):
 439         return self.__refreshAttribute(item, index, 'ItemType')
 440
 441     def RefreshCheckedState(self, item, index):
 442         self.__refreshAttribute(item, index, 'ItemChecked')
 443
 444     def ChildIndices(self, itemIndex):
 445         childrenCount = self.OnGetChildrenCount(itemIndex)
 446         return [itemIndex + (childNumber,) for childNumber \
 447                 in range(childrenCount)]
 448
 449     def OnItemExpanding(self, event):
 450         self.RefreshChildrenRecursively(event.GetItem())
 451         event.Skip()
 452
 453     def OnItemCollapsed(self, event):
 454         parent = self.GetItemParent(event.GetItem())
 455         if not parent:
 456             parent = self.GetRootItem()
 457         self.RefreshChildrenRecursively(parent)
 458         event.Skip()
 459
 460     def __refreshAttribute(self, item, index, attribute, *args):
 461         ''' Refresh the specified attribute if necessary. '''
 462         value = getattr(self, 'OnGet%s'%attribute)(index, *args)
 463         if getattr(self, 'Get%s'%attribute)(item, *args) != value:
 464             return getattr(self, 'Set%s'%attribute)(item, value, *args)
 465         else:
 466             return item
 467
 468
 469 class DragAndDrop(TreeAPIHarmonizer, TreeHelper):
 470     ''' This is a mixin class that can be used to easily implement
 471     dragging and dropping of tree items. It can be mixed in with
 472     wx.TreeCtrl, wx.gizmos.TreeListCtrl, or wx.lib.customtree.CustomTreeCtrl.
 473
 474     To use it derive a new class from this class and one of the tree
 475     controls, e.g.:
 476     class MyTree(DragAndDrop, wx.TreeCtrl):
 477         ...
 478
 479     You *must* implement OnDrop. OnDrop is called when the user has
 480     dropped an item on top of another item. It's up to you to decide how
 481     to handle the drop. If you are using this mixin together with the
 482     VirtualTree mixin, it makes sense to rearrange your underlying data
 483     and then call RefreshItems to let the virtual tree refresh itself. '''
 484
 485     def __init__(self, *args, **kwargs):
 486         kwargs['style'] = kwargs.get('style', wx.TR_DEFAULT_STYLE) | \
 487                           wx.TR_HIDE_ROOT
 488         super(DragAndDrop, self).__init__(*args, **kwargs)
 489         self.Bind(wx.EVT_TREE_BEGIN_DRAG, self.OnBeginDrag)
 490
 491     def OnDrop(self, dropItem, dragItem):
 492         ''' This function must be overloaded in the derived class.
 493         dragItem is the item being dragged by the user. dropItem is the
 494         item dragItem is dropped upon. If the user doesn't drop dragItem
 495         on another item, dropItem equals the (hidden) root item of the
 496         tree control. '''
 497         raise NotImplementedError
 498
 499     def OnBeginDrag(self, event):
 500         # We allow only one item to be dragged at a time, to keep it simple
 501         self._dragItem = event.GetItem()
 502         if self._dragItem and self._dragItem != self.GetRootItem():
 503             self.StartDragging()
 504             event.Allow()
 505         else:
 506             event.Veto()
 507
 508     def OnEndDrag(self, event):
 509         self.StopDragging()
 510         dropTarget = event.GetItem()
 511         if not dropTarget:
 512             dropTarget = self.GetRootItem()
 513         if self.IsValidDropTarget(dropTarget):
 514             self.UnselectAll()
 515             if dropTarget != self.GetRootItem():
 516                 self.SelectItem(dropTarget)
 517             self.OnDrop(dropTarget, self._dragItem)
 518
 519     def OnDragging(self, event):
 520         if not event.Dragging():
 521             self.StopDragging()
 522             return
 523         item, flags, column = self.HitTest(wx.Point(event.GetX(), event.GetY()),
 524                                            alwaysReturnColumn=True)
 525         if not item:
 526             item = self.GetRootItem()
 527         if self.IsValidDropTarget(item):
 528             self.SetCursorToDragging()
 529         else:
 530             self.SetCursorToDroppingImpossible()
 531         if flags & wx.TREE_HITTEST_ONITEMBUTTON:
 532             self.Expand(item)
 533         if self.GetSelections() != [item]:
 534             self.UnselectAll()
 535             if item != self.GetRootItem():
 536                 self.SelectItem(item)
 537         event.Skip()
 538
 539     def StartDragging(self):
 540         self.GetMainWindow().Bind(wx.EVT_MOTION, self.OnDragging)
 541         self.Bind(wx.EVT_TREE_END_DRAG, self.OnEndDrag)
 542         self.SetCursorToDragging()
 543
 544     def StopDragging(self):
 545         self.GetMainWindow().Unbind(wx.EVT_MOTION)
 546         self.Unbind(wx.EVT_TREE_END_DRAG)
 547         self.ResetCursor()
 548         self.UnselectAll()
 549         self.SelectItem(self._dragItem)
 550
 551     def SetCursorToDragging(self):
 552         self.GetMainWindow().SetCursor(wx.StockCursor(wx.CURSOR_HAND))
 553
 554     def SetCursorToDroppingImpossible(self):
 555         self.GetMainWindow().SetCursor(wx.StockCursor(wx.CURSOR_NO_ENTRY))
 556
 557     def ResetCursor(self):
 558         self.GetMainWindow().SetCursor(wx.NullCursor)
 559
 560     def IsValidDropTarget(self, dropTarget):
 561         if dropTarget:
 562             allChildren = self.GetItemChildren(self._dragItem, recursively=True)
 563             parent = self.GetItemParent(self._dragItem)
 564             return dropTarget not in [self._dragItem, parent] + allChildren
 565         else:
 566             return True
 567
 568
 569 class ExpansionState(TreeAPIHarmonizer, TreeHelper):
 570     ''' This is a mixin class that can be used to save and restore
 571     the expansion state (i.e. which items are expanded and which items
 572     are collapsed) of a tree. It can be mixed in with wx.TreeCtrl,
 573     wx.gizmos.TreeListCtrl, or wx.lib.customtree.CustomTreeCtrl.
 574
 575     To use it derive a new class from this class and one of the tree
 576     controls, e.g.:
 577     class MyTree(ExpansionState, wx.TreeCtrl):
 578         ...
 579
 580     By default, ExpansionState uses the position of tree items in the tree
 581     to keep track of which items are expanded. This should be sufficient
 582     for the simple scenario where you save the expansion state of the tree
 583     when the user closes the application or file so that you can restore
 584     the expansion state when the user start the application or loads that
 585     file for the next session.
 586
 587     If you need to add or remove items between the moments of saving and
 588     restoring the expansion state (e.g. in case of a multi-user application)
 589     you must override GetItemIdentity so that saving and loading of the
 590     expansion doesn't depend on the position of items in the tree, but
 591     rather on some more stable characteristic of the underlying domain
 592     object, e.g. a social security number in case of persons or an isbn
 593     number in case of books. '''
 594
 595     def GetItemIdentity(self, item):
 596         ''' Return a hashable object that represents the identity of the
 597         item. By default this returns the position of the item in the
 598         tree. You may want to override this to return the item label
 599         (if you know that labels are unique and don't change), or return
 600         something that represents the underlying domain object, e.g.
 601         a database key. '''
 602         return self.GetIndexOfItem(item)
 603
 604     def GetExpansionState(self):
 605         ''' GetExpansionState() -> list of expanded items. Expanded items
 606         are coded as determined by the result of GetItemIdentity(item). '''
 607         root = self.GetRootItem()
 608         if not root:
 609             return []
 610         if self.HasFlag(wx.TR_HIDE_ROOT):
 611             return self.GetExpansionStateOfChildren(root)
 612         else:
 613             return self.GetExpansionStateOfItem(root)
 614
 615     def SetExpansionState(self, listOfExpandedItems):
 616         ''' SetExpansionState(listOfExpandedItems). Expands all tree items
 617         whose identity, as determined by GetItemIdentity(item), is present
 618         in the list and collapses all other tree items. '''
 619         root = self.GetRootItem()
 620         if not root:
 621             return
 622         if self.HasFlag(wx.TR_HIDE_ROOT):
 623             self.SetExpansionStateOfChildren(listOfExpandedItems, root)
 624         else:
 625             self.SetExpansionStateOfItem(listOfExpandedItems, root)
 626
 627     ExpansionState = property(GetExpansionState, SetExpansionState)
 628
 629     def GetExpansionStateOfItem(self, item):
 630         listOfExpandedItems = []
 631         if self.IsExpanded(item):
 632             listOfExpandedItems.append(self.GetItemIdentity(item))
 633             listOfExpandedItems.extend(self.GetExpansionStateOfChildren(item))
 634         return listOfExpandedItems
 635
 636     def GetExpansionStateOfChildren(self, item):
 637         listOfExpandedItems = []
 638         for child in self.GetItemChildren(item):
 639             listOfExpandedItems.extend(self.GetExpansionStateOfItem(child))
 640         return listOfExpandedItems
 641
 642     def SetExpansionStateOfItem(self, listOfExpandedItems, item):
 643         if self.GetItemIdentity(item) in listOfExpandedItems:
 644             self.Expand(item)
 645             self.SetExpansionStateOfChildren(listOfExpandedItems, item)
 646         else:
 647             self.Collapse(item)
 648
 649     def SetExpansionStateOfChildren(self, listOfExpandedItems, item):
 650         for child in self.GetItemChildren(item):
 651             self.SetExpansionStateOfItem(listOfExpandedItems, child)