codekingpro/portable-devtools

1import re2import sqlparse3from sqlparse.tokens import Name4from collections import defaultdict5 6 7white_space_regex = re.compile("\\s+", re.MULTILINE)8 9 10def _compile_regex(keyword):11    # Surround the keyword with word boundaries and replace interior whitespace12    # with whitespace wildcards13    pattern = "\\b" + white_space_regex.sub(r"\\s+", keyword) + "\\b"14    return re.compile(pattern, re.MULTILINE | re.IGNORECASE)15 16 17class PrevalenceCounter:18    def __init__(self, keywords):19        self.keyword_counts = defaultdict(int)20        self.name_counts = defaultdict(int)21        self.keyword_regexs = dict((kw, _compile_regex(kw)) for kw in keywords)22 23    def update(self, text):24        self.update_keywords(text)25        self.update_names(text)26 27    def update_names(self, text):28        for parsed in sqlparse.parse(text):29            for token in parsed.flatten():30                if token.ttype in Name:31                    self.name_counts[token.value] += 132 33    def clear_names(self):34        self.name_counts = defaultdict(int)35 36    def update_keywords(self, text):37        # Count keywords. Can't rely for sqlparse for this, because it's38        # database agnostic39        for keyword, regex in self.keyword_regexs.items():40            for _ in regex.finditer(text):41                self.keyword_counts[keyword] += 142 43    def keyword_count(self, keyword):44        return self.keyword_counts[keyword]45 46    def name_count(self, name):47        return self.name_counts[name]48