|
@@ -5,6 +5,7 @@ class EventQueue(object):
|
|
|
def __init__(self):
|
|
|
self.event_list = []
|
|
|
self.event_time_numbers = {}
|
|
|
+ self.removed = set()
|
|
|
|
|
|
def __str__(self):
|
|
|
return str(self.event_list)
|
|
@@ -18,15 +19,19 @@ class EventQueue(object):
|
|
|
def add(self, event_time, event):
|
|
|
self.event_time_numbers[event_time] = self.event_time_numbers.setdefault(event_time, 0) + 1
|
|
|
heappush(self.event_list, (event_time, self.event_time_numbers[event_time], event))
|
|
|
- return id(event)
|
|
|
+ return event
|
|
|
|
|
|
- def remove(self, event_id):
|
|
|
- self.event_list = sorted([e for e in self.event_list if id(e) != event_id])
|
|
|
+ def remove(self, event):
|
|
|
+ self.removed.add(event)
|
|
|
|
|
|
def pop(self):
|
|
|
- item = heappop(self.event_list)
|
|
|
- event_time = item[0]
|
|
|
- self.event_time_numbers[event_time] -= 1
|
|
|
- if not self.event_time_numbers[event_time]:
|
|
|
- del self.event_time_numbers[event_time]
|
|
|
- return item[2]
|
|
|
+ while 1:
|
|
|
+ item = heappop(self.event_list)
|
|
|
+ event_time = item[0]
|
|
|
+ self.event_time_numbers[event_time] -= 1
|
|
|
+ if not self.event_time_numbers[event_time]:
|
|
|
+ del self.event_time_numbers[event_time]
|
|
|
+ if item not in self.removed:
|
|
|
+ return item[2]
|
|
|
+ else:
|
|
|
+ self.removed.remove(item)
|