monkey.utils.PQueue: fix size tracking

1 files changed, 4 insertions, 5 deletions

diff --git a/monkey/util.py b/monkey/util.py
index da04828..655cca9 100644
--- a/monkey/util.py
+++ b/monkey/util.py
@@ -15,7 +15,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 from heapq import heappush, heappop
-import itertools
+from itertools import chain, count
 import math
 
 # A simple priority queue based on the heapq class.
@@ -25,18 +25,17 @@ class PQueue(object):
     def __init__(self):
         self.pq = []                         # list of entries arranged in a heap
         self.entry_finder = {}               # mapping of tasks to entries
-        self.counter = itertools.count()     # unique sequence count
+        self.counter = count()               # unique sequence count
         self.size = 0
 
     def push(self, task, priority=0):
         'Add a new task or update the priority of an existing task'
         if task in self.entry_finder:
             self.remove(task)
-        else:
-            self.size += 1
         entry = [priority, next(self.counter), task]
         self.entry_finder[task] = entry
         heappush(self.pq, entry)
+        self.size += 1
 
     def remove(self, task):
         'Mark an existing task as REMOVED.  Raise KeyError if not found.'
@@ -82,7 +81,7 @@ def damerau_levenshtein(a, b):
     # Construct an alphabet out of characters in [a] and [b] with a sequential
     # identifier for each word.
     alphabet = {}
-    for char in itertools.chain(a, b):
+    for char in chain(a, b):
         if char not in alphabet:
             alphabet[char] = len(alphabet)