apply/str_pattern.py - master - socialgraph/dataset

apply/str_pattern.py
96 строк2 KB
socialgraph
parse
09 авг 2026, 10:22
09 авг 2026, 10:2276cada1
from tqdm import tqdm

import regex as re
from tools.const import is_selected
from tools.utils import load_examples, save_examples, load_snippets

FLAGS = {
    'ASCII': re.ASCII,
    'A': re.ASCII,
    'IGNORECASE': re.IGNORECASE,
    'I': re.IGNORECASE,
    'MULTILINE': re.MULTILINE,
    'M': re.MULTILINE,
    'DOTALL': re.DOTALL,
    'S': re.DOTALL,
    'UNICODE': re.UNICODE,
    'U': re.UNICODE,
    'VERBOSE': re.VERBOSE,
    'X': re.VERBOSE,
    'LOCALE': re.LOCALE,
    'L': re.LOCALE
}


def get_pattern(q):
    flags = 0
    for f in q['flags']:
        flags |= FLAGS.get(f.upper(), 0)
    pattern = q['pattern'].replace('\\\\', '\\')
    return re.compile(rf'{pattern}', flags)


def get_candidates(snippets: dict, patterns: list[dict]) -> list[dict]:
    for q in sorted(patterns, key=lambda q: q.get('confidence', 0)):
        try:
            pattern = get_pattern(q)
            for key, s in snippets.items():
                if pattern.search(s, timeout=1):
                    yield {
                        'str_pattern': q['pattern'],
                        'confidence': q.get('confidence', 0),
                        'key': key,
                    }
        except Exception as e:
            print(e)


def run(language):
    examples = load_examples(language, 'updates')
    src = {q['id']: q for q in examples}

    for q in tqdm(examples):
        related = [r for r in q['related'] if not r.get('error')]
        if not related:
            continue

        if not is_selected(q):
            continue

        for state in ['before', 'after']:
            snippets = load_snippets(q['project'], q[state]['commit'])

            for r in related:
                dst = src[r['id']]
                if not is_selected(dst):
                    continue

                patterns = [p for p in dst['str_patterns'] if not p.get('error')]
                if not patterns:
                    continue

                candidates = list(get_candidates(snippets, patterns))
                dst['results'].append({
                    'id': q['id'],
                    'state': state,
                    'project': q['project'],
                    'commit': q[state]['commit'],
                    'candidates': candidates
                })

    save_examples(examples, language)


if __name__ == '__main__':
    import argparse

    args = argparse.ArgumentParser()

    args.add_argument(
        "--language",
        type=str,
        default="java"
    )

    args = args.parse_args()
    run(args.language)