idsedykh/codebleu

1# Copyright (c) Microsoft Corporation. 
2# Licensed under the MIT license.
3
4import os
5from .parser import DFG_python,DFG_java,DFG_ruby,DFG_go,DFG_php,DFG_javascript,DFG_csharp
6from .parser import (remove_comments_and_docstrings,
7                   tree_to_token_index,
8                   index_to_code_token,
9                   tree_to_variable_index)
10from tree_sitter import Language, Parser
11
12dfg_function={
13    'python':DFG_python,
14    'java':DFG_java,
15    'ruby':DFG_ruby,
16    'go':DFG_go,
17    'php':DFG_php,
18    'javascript':DFG_javascript,
19    'c_sharp':DFG_csharp,
20}
21
22def calc_syntax_match(references, candidate, lang):
23    return corpus_syntax_match([references], [candidate], lang)
24
25def corpus_syntax_match(references, candidates, lang):
26    # print(os.listdir())
27    JAVA_LANGUAGE = Language(os.path.abspath(os.path.dirname(__file__)) + '/parser/my-languages.so', lang)
28    parser = Parser()
29    parser.set_language(JAVA_LANGUAGE)
30    match_count = 0
31    total_count = 0
32
33    for i in range(len(candidates)):
34        references_sample = references[i]
35        candidate = candidates[i] 
36        for reference in references_sample:
37            try:
38                candidate=remove_comments_and_docstrings(candidate,'java')
39            except:
40                pass    
41            try:
42                reference=remove_comments_and_docstrings(reference,'java')
43            except:
44                pass  
45
46            candidate_tree = parser.parse(bytes(candidate,'utf8')).root_node
47
48            reference_tree = parser.parse(bytes(reference,'utf8')).root_node
49
50            def get_all_sub_trees(root_node):
51                node_stack = []
52                sub_tree_sexp_list = []
53                depth = 1
54                node_stack.append([root_node, depth])
55                while len(node_stack) != 0:
56                    cur_node, cur_depth = node_stack.pop()
57                    sub_tree_sexp_list.append([cur_node.sexp(), cur_depth])
58                    for child_node in cur_node.children:
59                        if len(child_node.children) != 0:
60                            depth = cur_depth + 1
61                            node_stack.append([child_node, depth])
62                return sub_tree_sexp_list
63            cand_sexps = [x[0] for x in get_all_sub_trees(candidate_tree)]
64            ref_sexps = get_all_sub_trees(reference_tree)
65
66            # print(cand_sexps)
67            # print(ref_sexps)
68            
69            for sub_tree, depth in ref_sexps:
70                if sub_tree in cand_sexps:
71                     match_count += 1
72            total_count += len(ref_sexps)          
73       
74    score = match_count / total_count
75    return score
76