apply/ast_pattern.py - master - socialgraph/dataset

apply/ast_pattern.py
130 строк3 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.java_ast import get_normalized_ast
from tools.utils import load_examples, save_examples, load_snippets


def contains(ast, pattern):
    i = 0
    j = 0
    k = 1

    n = len(ast)
    m = len(pattern)

    C = list(range(n))
    S = []
    for s, c in enumerate(ast):
        if c == '(':
            S.append(s)
        elif c == ')':
            C[S.pop()] = s

    while i < n and j < m:
        if n - i < m - j:
            i = k
            j = 0
            k += 1

        elif pattern[j] == '$':
            j += 1

            while i < n:
                i = C[i] + 1

                if j == m:
                    return True

                if i == n or pattern[j] == '$' or ast[i] == pattern[j]:
                    break

        elif ast[i] == pattern[j]:
            i += 1
            j += 1

        else:
            i = k
            j = 0
            k += 1

    return j == m


def get_pattern(q):
    if not q['pattern']:
        raise IOError

    pattern = q['pattern'].lower()
    pattern = re.sub(r'\[call_[^]]*]', '_call_', pattern)
    pattern = re.sub(r'\$[^ ]*', '$', pattern)

    return tuple(pattern.split())


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 contains(tuple(s.split()), pattern):
                    yield {
                        'ast_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'])
            snippets = {k: get_normalized_ast(v) for k, v in snippets.items()}

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

                patterns = [p for p in dst['ast_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)