aravagarwal/CodeCloakPII

1from anoninterface import parser, analyzer, anonymizer, PYTHON_LANGUAGE, MAX_STRINGS_IN_SAMPLE2from copy import deepcopy3import gradio as gr4 5TREE_SITTER_QUERY_TO_QUERY_NAME = {6  "(comment) @comment": "comment", 7  "(string_content) @string_content": "string",8  "(identifier) @identifier": "name"9}10 11TREE_SITTER_QUERY_NAME_TO_QUERY = {value:key for key, value in TREE_SITTER_QUERY_TO_QUERY_NAME.items()}12 13 14def get_captures(node_text, querystr, tree=None):15  if tree is None:16    tree = parser.parse(node_text.encode("utf8"))17  query = PYTHON_LANGUAGE.query(querystr)18  return query.captures(tree.root_node), tree19 20def update_tree_and_source(node, new_node_text, source_text, source_tree, captured_nodes = []):21  new_elem_bytes = len(new_node_text.encode('utf8'))22  source_bytes = source_text.encode('utf8')23  node_bstart = node.start_byte24  node_estart = node.end_byte25  node_nestart = node.start_byte + new_elem_bytes26  node_pbstart = node.start_point27  node_pestart = node.end_point28  node_pnestart = (node.end_point[0], node.end_point[1]+new_elem_bytes)29  new_source_text = source_bytes[:node.byte_range[0]] + new_node_text.encode('utf8') + source_bytes[node.byte_range[1]:]30  source_tree.edit(node_bstart, node_estart, node_nestart, node_pbstart, node_pestart, node_pnestart)31  return parser.parse(new_source_text, source_tree), new_source_text.decode('utf8')32 33 34def syntax_aware_replace_with_lambda(solution_text, query, replace_fn, filter_fn):35  captures, tree = get_captures(solution_text, query)36  number_of_edits = len(captures)37  #print("PRE_EDIT:",solution_text,sep="\n")38  for elem in range(number_of_edits):39    print(captures)40    old_text = deepcopy(captures[elem][0].text.decode('utf8'))41    new_text = filter_fn(old_text, anon_fn = replace_fn)42    new_tree, solution_text = update_tree_and_source(captures[elem][0], new_text, solution_text, tree)43    tree = new_tree44    captures, tree = get_captures(solution_text, query, tree)45  redacted_text = solution_text46  return redacted_text47 48 49def visualize_tree_sitter_matches(solution_text, queries, unused=[]):50  queries = [TREE_SITTER_QUERY_NAME_TO_QUERY[query] for query in queries]51  solution_text_bytes = solution_text.encode('utf8')52  total_capture_list = []53  for query in queries:54    captures, tree = get_captures(solution_text, query)55    captures = [(elem[0], query) for elem in captures]56    total_capture_list.extend(captures)57  total_capture_list = sorted(total_capture_list, key=lambda x: x[0].byte_range[0])58  labelled_text = []59  read_loc = 060  for capture in total_capture_list:61    capture_start, capture_end = capture[0].byte_range62    if capture_start >= read_loc:63      labelled_text.append((solution_text_bytes[read_loc:capture_start].decode('utf8'), None))64      labelled_text.append((solution_text_bytes[capture_start:capture_end].decode('utf8'), TREE_SITTER_QUERY_TO_QUERY_NAME[capture[1]]))65      read_loc = capture_end66    else:67      raise ValueError("Overlapping Query Spans Are Currently Not Supported")68  spans_list = [elem[0].text.decode('utf8') for elem in total_capture_list]69  if len(total_capture_list) > MAX_STRINGS_IN_SAMPLE:70    gradio.Warning("Your code snippet has too many identifiers.")71  return labelled_text if labelled_text != [] else [solution_text], total_capture_list72 73