from   wxPython      import wx

class ImageList(wx.wxImageList):
    """
    This is a class to make the standard wxImageList
    much easier to work with.  It adds two main features
    to the wxImageList:

    1. It contains a dictionary and saves images with
       keys in order to easily retrieve the indexes.
    2. It allows wxImages to be stored in addition to
       bitmaps.
    """
    def __init__(self, w, h):
        wx.wxImageList.__init__(self, w, h)
        self.lookup = {}

    def Add(self, image):
        """
        Add a bitmap or an image to the list.
        Return the image's new index.
        """
        return self.add(image, image)

    def add(self, key, image):
        """
        Add a bitmap or an image to the list using
        the given key.  The image may have the value None,
        in which case it will get the index -1.
        Return the image's new index.
        """
        if image is None:
            self.lookup[key] = -1
            return -1
        try:
            bitmap = image.ConvertToBitmap()
        except AttributeError:
            bitmap = image
        return self.__add(key, bitmap)

    def getIndex(self, key):
        return self.lookup[key]

    def remove(self, key):
        index = self.getIndex(key)
        del self.lookup[key]
        if index != -1:
            wx.wxImageList.Remove(self, index)

    def RemoveAll(self):
        wx.wxImageList.RemoveAll(self)
        self.lookup.clear()

    def __add(self, key, bitmap):
        index = wx.wxImageList.Add(self, bitmap)
        self.lookup[key] = index
        return index

    #
    # Implement the dict's contain protocol.
    #

    def has_key(self, key):
        return self.lookup.has_key(key)

    def __contains__(self, key):
        return self.has_key(key)