class BitKeyMap:
"""Map a string to a limited number of bits"""
def __init__(self, bits=16, load=None):
self.useBits(bits)
self.mapping = {}
self.unused_keys = [] # keys are added to this as they are deleted
self.last_key = 0
self.rwlock = ReadWriteLock()
if load: self.load(load)
def useBits(self, bits):
self.bits = bits
self.max_key = 2 ** bits
def size(self):
return len(self.mapping)
def lastKey(self):
return self.last_key
def get(self, key_string):
self.rwlock.acquire_read()
if key_string not in self.mapping:
self.rwlock.release_read()
self.rwlock.acquire_write()
try:
if self.last_key < self.max_key-1:
self.last_key += 1
self.mapping[key_string] = self.last_key
else:
if len(self.unused_keys) == 0:
raise Exception("Mapping is full; all bits have been exhausted")
self.mapping[key_string] = self.unused_keys.pop()
return self.mapping[key_string]
finally:
self.rwlock.release_write()
#print "BITKEYMAP_NEW_KEY\t%d\t%s\t%d" % (threading.current_thread().ident, key_string, self.mapping[key_string])
k = self.mapping[key_string]
self.rwlock.release_read()
return k
def contains(self, key_string):
self.rwlock.acquire_read()
b = key_string in self.mapping
self.rwlock.release_read()
return b
def free(self, key_string):
self.rwlock.acquire_write()
if key_string in self.mapping:
self.unused_keys.append(self.mapping[key_string])
del self.mapping[key_string]
self.rwlock.release_write()
def dump(self, file_or_str):
d = dict(max_key=self.max_key, mapping=self.mapping, unused_keys=self.unused_keys, last_key=self.last_key)
if type(file_or_str) is str:
with open(file_or_str, 'w') as f:
pickle.dump(d, f)
else: pickle.dump(d, file_or_str)
def load(self, file_or_str):
if type(file_or_str) is str:
with open(file_or_str, 'r') as f:
d = pickle.load(f)
else: d = pickle.load(file_or_str)
self.max_key, self.mapping, self.unused_keys, self.last_key = d['max_key'], d['mapping'], d['unused_keys'], d['last_key']
