datastructures.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. # -*- coding: utf-8 -*-
  2. # Licensed under the Apache License, Version 2.0 (the "License");
  3. # you may not use this file except in compliance with the License.
  4. # You may obtain a copy of the License at
  5. #
  6. # http://www.apache.org/licenses/LICENSE-2.0
  7. #
  8. # Unless required by applicable law or agreed to in writing, software
  9. # distributed under the License is distributed on an "AS IS" BASIS,
  10. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
  11. # implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import copy
  15. from types import GeneratorType
  16. # taken from Django.utils.datastructures
  17. class SortedDict(dict):
  18. """
  19. A dictionary that keeps its keys in the order in which they're inserted.
  20. """
  21. def __new__(cls, *args, **kwargs):
  22. instance = super(SortedDict, cls).__new__(cls, *args, **kwargs)
  23. instance.keyOrder = []
  24. return instance
  25. def __init__(self, data=None):
  26. if data is None:
  27. data = {}
  28. elif isinstance(data, GeneratorType):
  29. # Unfortunately we need to be able to read a generator twice. Once
  30. # to get the data into self with our super().__init__ call and a
  31. # second time to setup keyOrder correctly
  32. data = list(data)
  33. super(SortedDict, self).__init__(data)
  34. if isinstance(data, dict):
  35. self.keyOrder = data.keys()
  36. else:
  37. self.keyOrder = []
  38. seen = set()
  39. for key, value in data:
  40. if key not in seen:
  41. self.keyOrder.append(key)
  42. seen.add(key)
  43. def __deepcopy__(self, memo):
  44. return self.__class__([(key, copy.deepcopy(value, memo))
  45. for key, value in self.iteritems()])
  46. def __setitem__(self, key, value):
  47. if key not in self:
  48. self.keyOrder.append(key)
  49. super(SortedDict, self).__setitem__(key, value)
  50. def __delitem__(self, key):
  51. super(SortedDict, self).__delitem__(key)
  52. self.keyOrder.remove(key)
  53. def __iter__(self):
  54. return iter(self.keyOrder)
  55. def pop(self, k, *args):
  56. result = super(SortedDict, self).pop(k, *args)
  57. try:
  58. self.keyOrder.remove(k)
  59. except ValueError:
  60. # Key wasn't in the dictionary in the first place. No problem.
  61. pass
  62. return result
  63. def popitem(self):
  64. result = super(SortedDict, self).popitem()
  65. self.keyOrder.remove(result[0])
  66. return result
  67. def items(self):
  68. return zip(self.keyOrder, self.values())
  69. def iteritems(self):
  70. for key in self.keyOrder:
  71. yield key, self[key]
  72. def keys(self):
  73. return self.keyOrder[:]
  74. def iterkeys(self):
  75. return iter(self.keyOrder)
  76. def values(self):
  77. return map(self.__getitem__, self.keyOrder)
  78. def itervalues(self):
  79. for key in self.keyOrder:
  80. yield self[key]
  81. def update(self, dict_):
  82. for k, v in dict_.iteritems():
  83. self[k] = v
  84. def setdefault(self, key, default):
  85. if key not in self:
  86. self.keyOrder.append(key)
  87. return super(SortedDict, self).setdefault(key, default)
  88. def value_for_index(self, index):
  89. """Returns the value of the item at the given zero-based index."""
  90. return self[self.keyOrder[index]]
  91. def insert(self, index, key, value):
  92. """Inserts the key, value pair before the item with the given index."""
  93. if key in self.keyOrder:
  94. n = self.keyOrder.index(key)
  95. del self.keyOrder[n]
  96. if n < index:
  97. index -= 1
  98. self.keyOrder.insert(index, key)
  99. super(SortedDict, self).__setitem__(key, value)
  100. def copy(self):
  101. """Returns a copy of this object."""
  102. # This way of initializing the copy means it works for subclasses, too.
  103. obj = self.__class__(self)
  104. obj.keyOrder = self.keyOrder[:]
  105. return obj
  106. def __repr__(self):
  107. """
  108. Replaces the normal dict.__repr__ with a version that returns the keys
  109. in their sorted order.
  110. """
  111. return '{%s}' % ', '.join(['%r: %r' % (k, v) for k, v in self.items()])
  112. def clear(self):
  113. super(SortedDict, self).clear()
  114. self.keyOrder = []