########################################################################
# Copyright (C) 2010-2018 VMWare, Inc.                                 #
# All Rights Reserved                                                  #
########################################################################

import copy

class IndexedDict(dict):
   """IndexedDict adds indexing of attributes of the dict's objects.
      Each of the dict's values is assumed to be a subclass of Object and
      getattr is used to obtain attribute values.
      The index is updated when values are assigned into the dict, at
      dict creation time, when update() is called, and when keys are removed.
      Attributes:
         * indexedfields - A list of attributes to index.  If the attribute is
                           not defined in a value instance, it will not be
                           indexed.
         * indexes       - A dict of the form <attr>: <attr-index>, where
                           the key <attr> is the name of the attribute being
                           indexed and is one of indexedfields, and <attr-index>
                           is a dict mapping each different value of <attr> to a
                           set of the keys in the IndexedDict whose values
                           have the corresponding attribute.  An example:
                              d = {'San Francisco': Object(timezone = 'PST',
                                                     country = 'USA'),
                                   'New York':      Object(timezone = 'EST',
                                                     country = 'USA'),
                                   'London':        Object(timezone = 'GMT',
                                                     country = 'England'),
                                  }
                              mydict = IndexedDict(d, indexedfields=['timezone', 'country'])
                              mydict.indexes
                                {'timezone': {'PST': set(['San Francisco']),
                                              'EST': set(['New York']),
                                              'GMT': set(['London']) },
                                 'country': {'USA': set(['San Francisco',
                                                         'New York']),
                                             'England': set(['London']) }
                                }
   """
   def __init__(self, *args, **kwargs):
      """Constructs a new instance of IndexedDict.
         Parameters:
            The same as that for dict(), with this new keyword param-
            * indexedfields - A list of attributes to index.
      """
      self.indexedfields = kwargs.pop('indexedfields', [])
      self.indexes = dict()
      dict.__init__(self, *args, **kwargs)
      self.RebuildIndexes()

   def update(self, other):
      dict.update(self, other)
      self.RebuildIndexes()

   def __setitem__(self, key, value):
      if key in self:
         self._removefromindex(key, self[key])
      dict.__setitem__(self, key, value)
      self._addtoindex(key, value)

   def __delitem__(self, key):
      if key in self:
         self._removefromindex(key, self[key])
      dict.__delitem__(self, key)

   def __copy__(self):
      # The default copy method does not work since indexedfields is not set
      newone = self.__class__()
      newone.indexedfields = list(self.indexedfields)
      newone.update(self)
      return newone

   def __deepcopy__(self, memo):
      newone = self.__class__()
      newone.indexedfields = list(self.indexedfields)
      newone.RebuildIndexes()
      memo[id(self)] = newone
      for key, val in self.items():
         newone[copy.deepcopy(key)] = copy.deepcopy(val)
      return newone

   def RebuildIndexes(self):
      """Rebuilds all indexes based on current values in the dict.
         Parameters:
            None
      """
      for attr in self.indexedfields:
         self.indexes[attr] = dict()
      for key, val in self.items():
         self._addtoindex(key, val)

   def _addtoindex(self, key, val):
      # Adds attributes from val into our index using key
      for attr in self.indexedfields:
         attrval = getattr(val, attr, None)
         if attrval:
            if isinstance(attrval, list) or isinstance(attrval, tuple):
               values = attrval
            else:
               values = [attrval,]

            for token in values:
               self.indexes[attr].setdefault(token, set()).add(key)

   def _removefromindex(self, key, val):
      # Removes key from indexes
      for attr in self.indexedfields:
         attrval = getattr(val, attr, None)
         if attrval:
            if isinstance(attrval, list) or isinstance(attrval, tuple):
               values = attrval
            else:
               values = [attrval,]

            for token in values:
               try:
                  self.indexes[attr][token].remove(key)
                  if len(self.indexes[attr][token]) == 0:
                     del self.indexes[attr][token]
               except Exception:
                  pass
