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: 0.9.1
  29 Date: 26 March 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 __init__(self, *args, **kwargs):
  44         # CustomTreeCtrl uses a different keyword for the window style
  45         # argument ('ctstyle'). To hide this, we replace the 'style' keyword
  46         # by 'ctstyle' if we're mixed in with CustomTreeCtrl.
  47         if isinstance(self, wx.lib.customtreectrl.CustomTreeCtrl):
  48             kwargs['ctstyle'] = kwargs.pop('style', wx.TR_DEFAULT_STYLE)
  49         super(TreeAPIHarmonizer, self).__init__(*args, **kwargs)
  50
  51     def __callSuper(self, methodName, default, *args, **kwargs):
  52         # If our super class has a method called methodName, call it,
  53         # otherwise return the default value.
  54         superClass = super(TreeAPIHarmonizer, self)
  55         if hasattr(superClass, methodName):
  56             return getattr(superClass, methodName)(*args, **kwargs)
  57         else:
  58             return default
  59
  60     def GetColumnCount(self, *args, **kwargs):
  61         # Only TreeListCtrl has columns, return 0 if we are mixed in
  62         # with another tree control.
  63         return self.__callSuper('GetColumnCount', 0, *args, **kwargs)
  64
  65     def GetItemType(self, *args, **kwargs):
  66         # Only CustomTreeCtrl has different item types, return the
  67         # default item type if we are mixed in with another tree control.
  68         return self.__callSuper('GetItemType', 0, *args, **kwargs)
  69
  70     def SetItemType(self, item, newType):
  71         # CustomTreeCtrl doesn't support changing the item type on the fly,
  72         # so we create a new item and delete the old one. We currently only
  73         # keep the item text, would be nicer to also retain other attributes.
  74         text = self.GetItemText(item)
  75         newItem = self.InsertItem(self.GetItemParent(item), item, text,
  76                                   ct_type=newType)
  77         self.Delete(item)
  78         return newItem
  79
  80     def IsItemChecked(self, *args, **kwargs):
  81         # Only CustomTreeCtrl supports checkable items, return False if
  82         # we are mixed in with another tree control.
  83         return self.__callSuper('IsItemChecked', False, *args, **kwargs)
  84
  85     def GetItemChecked(self, *args, **kwargs):
  86         # For consistency's sake, provide a 'Get' and 'Set' method for
  87         # checkable items.
  88         return self.IsItemChecked(*args, **kwargs)
  89
  90     def SetItemChecked(self, *args, **kwargs):
  91         # For consistency's sake, provide a 'Get' and 'Set' method for
  92         # checkable items.
  93         return self.CheckItem(*args, **kwargs)
  94
  95     def GetMainWindow(self, *args, **kwargs):
  96         # Only TreeListCtrl has a separate main window, return self if we are
  97         # mixed in with another tree control.
  98         return self.__callSuper('GetMainWindow', self, *args, **kwargs)
  99
 100     def GetItemImage(self, item, which=wx.TreeItemIcon_Normal, column=-1):
 101         # CustomTreeCtrl always wants the which argument, so provide it
 102         # TreeListCtr.GetItemImage has a different order of arguments than
 103         # the other tree controls. Hide the differenes.
 104         if self.GetColumnCount():
 105             args = (item, column, which)
 106         else:
 107             args = (item, which)
 108         return super(TreeAPIHarmonizer, self).GetItemImage(*args)
 109
 110     def SetItemImage(self, item, imageIndex, which=wx.TreeItemIcon_Normal,
 111                      column=-1):
 112         # The SetItemImage signature is different for TreeListCtrl and
 113         # other tree controls. This adapter method hides the differences.
 114         if self.GetColumnCount():
 115             args = (item, imageIndex, column, which)
 116         else:
 117             args = (item, imageIndex, which)
 118         super(TreeAPIHarmonizer, self).SetItemImage(*args)
 119
 120     def UnselectAll(self):
 121         # Unselect all items, regardless of whether we are in multiple
 122         # selection mode or not.
 123         if self.HasFlag(wx.TR_MULTIPLE):
 124             super(TreeAPIHarmonizer, self).UnselectAll()
 125         else:
 126             # CustomTreeCtrl Unselect() doesn't seem to work in all cases,
 127             # also invoke UnselectAll just to be sure.
 128             self.Unselect()
 129             super(TreeAPIHarmonizer, self).UnselectAll()
 130
 131     def GetCount(self):
 132         # TreeListCtrl correctly ignores the root item when it is hidden,
 133         # but doesn't count the root item when it is visible
 134         itemCount = super(TreeAPIHarmonizer, self).GetCount()
 135         if self.GetColumnCount() and not self.HasFlag(wx.TR_HIDE_ROOT):
 136             itemCount += 1
 137         return itemCount
 138
 139     def GetSelections(self):
 140         # Always return a list of selected items, regardless of whether
 141         # we are in multiple selection mode or not.
 142         if self.HasFlag(wx.TR_MULTIPLE):
 143             selections = super(TreeAPIHarmonizer, self).GetSelections()
 144         else:
 145             selection = self.GetSelection()
 146             if selection:
 147                 selections = [selection]
 148             else:
 149                 selections = []
 150         # If the root item is hidden, it should never be selected,
 151         # unfortunately, CustomTreeCtrl and TreeCtrl allow it to be selected.
 152         if self.HasFlag(wx.TR_HIDE_ROOT):
 153             rootItem = self.GetRootItem()
 154             if rootItem and rootItem in selections:
 155                 selections.remove(rootItem)
 156         return selections
 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 doesn'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         item = self.GetItemByIndex(index)
 356         hasChildren = bool(self.OnGetChildrenCount(index))
 357         self.DoRefreshItem(item, index, hasChildren)
 358
 359     def RefreshChildrenRecursively(self, item, itemIndex=None):
 360         ''' Refresh the children of item, reusing as much of the
 361         existing items in the tree as possible. '''
 362         if itemIndex is None:
 363             itemIndex = self.GetIndexOfItem(item)
 364         reusableChildren = self.GetItemChildren(item)
 365         for childIndex in self.ChildIndices(itemIndex):
 366             if reusableChildren:
 367                 child = reusableChildren.pop(0)
 368             else:
 369                 child = self.AppendItem(item, '')
 370             self.RefreshItemRecursively(child, childIndex)
 371         for child in reusableChildren:
 372             self.Delete(child)
 373
 374     def RefreshItemRecursively(self, item, itemIndex):
 375         ''' Refresh the item and its children recursively. '''
 376         hasChildren = bool(self.OnGetChildrenCount(itemIndex))
 377         item = self.DoRefreshItem(item, itemIndex, hasChildren)
 378         # We need to refresh the children when the item is expanded and
 379         # when the item has no children, because in the latter case we
 380         # might have to delete old children from the tree:
 381         if self.IsExpanded(item) or not hasChildren:
 382             self.RefreshChildrenRecursively(item, itemIndex)
 383         self.SetItemHasChildren(item, hasChildren)
 384
 385     def DoRefreshItem(self, item, index, hasChildren):
 386         ''' Refresh one item. '''
 387         item = self.RefreshItemType(item, index)
 388         self.RefreshItemText(item, index)
 389         self.RefreshColumns(item, index)
 390         self.RefreshItemFont(item, index)
 391         self.RefreshTextColour(item, index)
 392         self.RefreshBackgroundColour(item, index)
 393         self.RefreshItemImage(item, index, hasChildren)
 394         self.RefreshCheckedState(item, index)
 395         return item
 396
 397     def RefreshItemText(self, item, index):
 398         self.__refreshAttribute(item, index, 'ItemText')
 399
 400     def RefreshColumns(self, item, index):
 401         for columnIndex in range(1, self.GetColumnCount()):
 402             self.__refreshAttribute(item, index, 'ItemText', columnIndex)
 403
 404     def RefreshItemFont(self, item, index):
 405         self.__refreshAttribute(item, index, 'ItemFont')
 406
 407     def RefreshTextColour(self, item, index):
 408         self.__refreshAttribute(item, index, 'ItemTextColour')
 409
 410     def RefreshBackgroundColour(self, item, index):
 411         self.__refreshAttribute(item, index, 'ItemBackgroundColour')
 412
 413     def RefreshItemImage(self, item, index, hasChildren):
 414         regularIcons = [wx.TreeItemIcon_Normal, wx.TreeItemIcon_Selected]
 415         expandedIcons = [wx.TreeItemIcon_Expanded,
 416                          wx.TreeItemIcon_SelectedExpanded]
 417         # Refresh images in first column:
 418         for icon in regularIcons:
 419             self.__refreshAttribute(item, index, 'ItemImage', icon)
 420         for icon in expandedIcons:
 421             if hasChildren:
 422                 imageIndex = self.OnGetItemImage(index, icon)
 423             else:
 424                 imageIndex = -1
 425             if self.GetItemImage(item, icon) != imageIndex or imageIndex == -1:
 426                 self.SetItemImage(item, imageIndex, icon)
 427         # Refresh images in remaining columns, if any:
 428         for columnIndex in range(1, self.GetColumnCount()):
 429             for icon in regularIcons:
 430                 self.__refreshAttribute(item, index, 'ItemImage', icon,
 431                                         columnIndex)
 432
 433     def RefreshItemType(self, item, index):
 434         return self.__refreshAttribute(item, index, 'ItemType')
 435
 436     def RefreshCheckedState(self, item, index):
 437         self.__refreshAttribute(item, index, 'ItemChecked')
 438
 439     def ChildIndices(self, itemIndex):
 440         childrenCount = self.OnGetChildrenCount(itemIndex)
 441         return [itemIndex + (childNumber,) for childNumber \
 442                 in range(childrenCount)]
 443
 444     def OnItemExpanding(self, event):
 445         self.RefreshChildrenRecursively(event.GetItem())
 446         event.Skip()
 447
 448     def OnItemCollapsed(self, event):
 449         parent = self.GetItemParent(event.GetItem())
 450         if not parent:
 451             parent = self.GetRootItem()
 452         self.RefreshChildrenRecursively(parent)
 453         event.Skip()
 454
 455     def __refreshAttribute(self, item, index, attribute, *args):
 456         ''' Refresh the specified attribute if necessary. '''
 457         value = getattr(self, 'OnGet%s'%attribute)(index, *args)
 458         if getattr(self, 'Get%s'%attribute)(item, *args) != value:
 459             return getattr(self, 'Set%s'%attribute)(item, value, *args)
 460         else:
 461             return item
 462
 463
 464 class DragAndDrop(TreeAPIHarmonizer, TreeHelper):
 465     ''' This is a mixin class that can be used to easily implement
 466     dragging and dropping of tree items. It can be mixed in with
 467     wx.TreeCtrl, wx.gizmos.TreeListCtrl, or wx.lib.customtree.CustomTreeCtrl.
 468
 469     To use it derive a new class from this class and one of the tree
 470     controls, e.g.:
 471     class MyTree(DragAndDrop, wx.TreeCtrl):
 472         ...
 473
 474     You *must* implement OnDrop. OnDrop is called when the user has
 475     dropped an item on top of another item. It's up to you to decide how
 476     to handle the drop. If you are using this mixin together with the
 477     VirtualTree mixin, it makes sense to rearrange your underlying data
 478     and then call RefreshItems to let the virtual tree refresh itself. '''
 479
 480     def __init__(self, *args, **kwargs):
 481         kwargs['style'] = kwargs.get('style', wx.TR_DEFAULT_STYLE) | \
 482                           wx.TR_HIDE_ROOT
 483         super(DragAndDrop, self).__init__(*args, **kwargs)
 484         self.Bind(wx.EVT_TREE_BEGIN_DRAG, self.OnBeginDrag)
 485
 486     def OnDrop(self, dropItem, dragItem):
 487         ''' This function must be overloaded in the derived class.
 488         dragItem is the item being dragged by the user. dropItem is the
 489         item dragItem is dropped upon. If the user doesn't drop dragItem
 490         on another item, dropItem equals the (hidden) root item of the
 491         tree control. '''
 492         raise NotImplementedError
 493
 494     def OnBeginDrag(self, event):
 495         # We allow only one item to be dragged at a time, to keep it simple
 496         self._dragItem = event.GetItem()
 497         if self._dragItem and self._dragItem != self.GetRootItem():
 498             self.StartDragging()
 499             event.Allow()
 500         else:
 501             event.Veto()
 502
 503     def OnEndDrag(self, event):
 504         self.StopDragging()
 505         dropTarget = event.GetItem()
 506         if not dropTarget:
 507             dropTarget = self.GetRootItem()
 508         if self.IsValidDropTarget(dropTarget):
 509             self.UnselectAll()
 510             if dropTarget != self.GetRootItem():
 511                 self.SelectItem(dropTarget)
 512             self.OnDrop(dropTarget, self._dragItem)
 513
 514     def OnDragging(self, event):
 515         if not event.Dragging():
 516             self.StopDragging()
 517             return
 518         item, flags, column = self.HitTest(wx.Point(event.GetX(), event.GetY()),
 519                                            alwaysReturnColumn=True)
 520         if not item:
 521             item = self.GetRootItem()
 522         if self.IsValidDropTarget(item):
 523             self.SetCursorToDragging()
 524         else:
 525             self.SetCursorToDroppingImpossible()
 526         if flags & wx.TREE_HITTEST_ONITEMBUTTON:
 527             self.Expand(item)
 528         if self.GetSelections() != [item]:
 529             self.UnselectAll()
 530             if item != self.GetRootItem():
 531                 self.SelectItem(item)
 532         event.Skip()
 533
 534     def StartDragging(self):
 535         self.GetMainWindow().Bind(wx.EVT_MOTION, self.OnDragging)
 536         self.Bind(wx.EVT_TREE_END_DRAG, self.OnEndDrag)
 537         self.SetCursorToDragging()
 538
 539     def StopDragging(self):
 540         self.GetMainWindow().Unbind(wx.EVT_MOTION)
 541         self.Unbind(wx.EVT_TREE_END_DRAG)
 542         self.ResetCursor()
 543         self.UnselectAll()
 544         self.SelectItem(self._dragItem)
 545
 546     def SetCursorToDragging(self):
 547         self.GetMainWindow().SetCursor(wx.StockCursor(wx.CURSOR_HAND))
 548
 549     def SetCursorToDroppingImpossible(self):
 550         self.GetMainWindow().SetCursor(wx.StockCursor(wx.CURSOR_NO_ENTRY))
 551
 552     def ResetCursor(self):
 553         self.GetMainWindow().SetCursor(wx.NullCursor)
 554
 555     def IsValidDropTarget(self, dropTarget):
 556         if dropTarget:
 557             allChildren = self.GetItemChildren(self._dragItem, recursively=True)
 558             parent = self.GetItemParent(self._dragItem)
 559             return dropTarget not in [self._dragItem, parent] + allChildren
 560         else:
 561             return True
 562
 563
 564 class ExpansionState(TreeAPIHarmonizer, TreeHelper):
 565     ''' This is a mixin class that can be used to save and restore
 566     the expansion state (i.e. which items are expanded and which items
 567     are collapsed) of a tree. It can be mixed in with wx.TreeCtrl,
 568     wx.gizmos.TreeListCtrl, or wx.lib.customtree.CustomTreeCtrl.
 569
 570     To use it derive a new class from this class and one of the tree
 571     controls, e.g.:
 572     class MyTree(ExpansionState, wx.TreeCtrl):
 573         ...
 574
 575     By default, ExpansionState uses the position of tree items in the tree
 576     to keep track of which items are expanded. This should be sufficient
 577     for the simple scenario where you save the expansion state of the tree
 578     when the user closes the application or file so that you can restore
 579     the expansion state when the user start the application or loads that
 580     file for the next session.
 581
 582     If you need to add or remove items between the moments of saving and
 583     restoring the expansion state (e.g. in case of a multi-user application)
 584     you must override GetItemIdentity so that saving and loading of the
 585     expansion doesn't depend on the position of items in the tree, but
 586     rather on some more stable characteristic of the underlying domain
 587     object, e.g. a social security number in case of persons or an isbn
 588     number in case of books. '''
 589
 590     def GetItemIdentity(self, item):
 591         ''' Return a hashable object that represents the identity of the
 592         item. By default this returns the position of the item in the
 593         tree. You may want to override this to return the item label
 594         (if you know that labels are unique and don't change), or return
 595         something that represents the underlying domain object, e.g.
 596         a database key. '''
 597         return self.GetIndexOfItem(item)
 598
 599     def GetExpansionState(self):
 600         ''' GetExpansionState() -> list of expanded items. Expanded items
 601         are coded as determined by the result of GetItemIdentity(item). '''
 602         root = self.GetRootItem()
 603         if not root:
 604             return []
 605         if self.HasFlag(wx.TR_HIDE_ROOT):
 606             return self.GetExpansionStateOfChildren(root)
 607         else:
 608             return self.GetExpansionStateOfItem(root)
 609
 610     def SetExpansionState(self, listOfExpandedItems):
 611         ''' SetExpansionState(listOfExpandedItems). Expands all tree items
 612         whose identity, as determined by GetItemIdentity(item), is present
 613         in the list and collapses all other tree items. '''
 614         root = self.GetRootItem()
 615         if not root:
 616             return
 617         if self.HasFlag(wx.TR_HIDE_ROOT):
 618             self.SetExpansionStateOfChildren(listOfExpandedItems, root)
 619         else:
 620             self.SetExpansionStateOfItem(listOfExpandedItems, root)
 621
 622     ExpansionState = property(GetExpansionState, SetExpansionState)
 623
 624     def GetExpansionStateOfItem(self, item):
 625         listOfExpandedItems = []
 626         if self.IsExpanded(item):
 627             listOfExpandedItems.append(self.GetItemIdentity(item))
 628             listOfExpandedItems.extend(self.GetExpansionStateOfChildren(item))
 629         return listOfExpandedItems
 630
 631     def GetExpansionStateOfChildren(self, item):
 632         listOfExpandedItems = []
 633         for child in self.GetItemChildren(item):
 634             listOfExpandedItems.extend(self.GetExpansionStateOfItem(child))
 635         return listOfExpandedItems
 636
 637     def SetExpansionStateOfItem(self, listOfExpandedItems, item):
 638         if self.GetItemIdentity(item) in listOfExpandedItems:
 639             self.Expand(item)
 640             self.SetExpansionStateOfChildren(listOfExpandedItems, item)
 641         else:
 642             self.Collapse(item)
 643
 644     def SetExpansionStateOfChildren(self, listOfExpandedItems, item):
 645         for child in self.GetItemChildren(item):
 646             self.SetExpansionStateOfItem(listOfExpandedItems, child)