import sys import os import random import gzip import pickle import optparse import time import ahocorasick from optparse import OptionParser def main(): options = parse_args() app = TestApplication(options) app.run() chars = 'abcdefghijklmnopqestuvwxyzABCDEFGHIJKLMNOPQESTUVWXYZ0123456789.,;:-' class TestApplication(object): def __init__(self, options): self.options = options self.words = set() random.seed(options.seed) def run(self): self.A = ahocorasick.Automaton() if self.options.compare and (not self.options.pickle and not self.options.save): self.generate_words() if self.options.pickle or self.options.save: self.add_words() if self.options.pickle: t1 = time.time() self.pickle() t2 = time.time() print(" time: %0.2fs" % (t2 - t1)) self.A.clear() if self.options.save: t1 = time.time() self.save() t2 = time.time() print(" time: %0.2fs" % (t2 - t1)) self.A.clear() if self.options.unpickle: t1 = time.time() self.unpickle() t2 = time.time() print(" time: %0.2fs" % (t2 - t1)) if self.options.load: t1 = time.time() self.load() t2 = time.time() print(" time: %0.2fs" % (t2 - t1)) if self.options.compare: self.compare() def add_words(self): if self.options.random: self.__add_random_words() else: self.__add_from_file() print("Automaton statistics:") d = self.A.get_stats() print("- nodes_count : %d" % d['nodes_count']) print("- words_count : %d" % d['words_count']) print("- links_count : %d" % d['links_count']) print("- longest_word : %d" % d['longest_word']) print("- sizeof_node : %d" % d['sizeof_node']) print("- total_size : %d" % d['total_size']) def __add_random_words(self): n = self.options.words print("Adding %d words" % n) while n > 0: word = self.generate_random_word() if self.options.compare: self.words.add(word) if self.A.add_word(word, True): n -= 1 def __add_from_file(self): n = self.options.words print("Adding %d words from %s" % (n, self.options.file_gz)) for i, word in enumerate(self.read()): if i > n: return if self.options.compare: self.words.add(word) self.A.add_word(word, True) def generate_words(self): if self.options.random: self.__generate_random_words() else: self.__load_words() def __generate_random_words(self): n = self.options.words print("Generating %d words" % n) while len(self.words) < n: word = self.generate_random_word() self.words.add(word) def __load_words(self): n = self.options.words print ("Loading %d words from %s" % (n, self.options.file_gz)) for i, word in enumerate(self.read()): if i < n: self.words.add(word) else: return def read(self): with gzip.open(self.options.file_gz, "rt", encoding="utf-8") as f: for line in f: yield line.strip() def pickle(self): path = self.options.picklepath print("Pickling automaton in %s" % path) with open(path, 'wb') as f: pickle.dump(self.A, f) size = os.path.getsize(path) print(" file size is %s" % format_size(size)) def unpickle(self): path = self.options.picklepath print("Unpickling automaton from %s" % path) with open(path, 'rb') as f: self.A = pickle.load(f) def save(self): path = self.options.picklepath print("Saving automaton in %s" % path) self.A.save(path, pickle.dumps); size = os.path.getsize(path) print(" file size is %s" % format_size(size)) def load(self): path = self.options.picklepath print("Loading automaton from %s" % path) self.A = ahocorasick.load(path, pickle.loads) def compare(self): print("Comparing added words with restored automaton") for word in self.A: self.words.remove(word) if self.words: print("Not all words were restored (%d missing)" % len(self.words)) def generate_random_word(self): n = random.randint(1, self.options.maxlength + 1) s = '' for i in range(n): s += random.choice(chars) return s def format_size(size): units = [ ('GB', 1024**3), ('MB', 1024**2), ('kB', 1024), ] for suffix, threshold in units: if size > threshold: return '%0.2f %s (%d bytes)' % (float(size)/threshold, suffix, size) return '%d bytes' % size def parse_args(): parser = OptionParser() parser.add_option( "--pickle-path", dest='picklepath', default='pickle_stresstest.pickle', help="path used in pickling/unpickling" ) parser.add_option( "-p", "--pickle", dest='pickle', action='store_true', default=False, help="perform pickle operation on generated/loaded words" ) parser.add_option( "-u", "--unpickle", dest='unpickle', action='store_true', default=False, help="perform unpickle operation on previously pickled data" ) parser.add_option( "-s", "--save", dest='save', action='store_true', default=False, help="perform save operation on generated/loaded words" ) parser.add_option( "-l", "--load", dest='load', action='store_true', default=False, help="perform load operation on previously saved data" ) parser.add_option( "-c", "--compare", action='store_true', default=False, help="compare generated/loaded words with unpickled data" ) parser.add_option( "--max-words", dest='words', type=int, default=50000, metavar='N', help="maximum number of words generated/loaded" ) parser.add_option( "--random", dest='random', action='store_true', default=False, help="generate random words" ) parser.add_option( "--seed", dest='seed', type=int, default=0, metavar='INT', help="random seed" ) parser.add_option( "--random-max-len", dest='maxlength', type=int, default=100, metavar='K', help="maximum count of characters in a word" ) parser.add_option( "--file-gz", metavar='FILE', help="load words from utf8-encoded gz file" ) (options, rest) = parser.parse_args() if not (options.file_gz or options.random): raise parser.error("pass --random or --file-gz option") if (options.pickle or options.unpickle) and (options.save or options.load): raise parser.error("use separately --pickle/--unpickle and --save/--load") return options if __name__ == '__main__': main()