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']