MCPcopy Index your code
hub / github.com/SalesforceAIResearch/perfcodegen / CodeProcessor

Class CodeProcessor

src/sanitize.py:185–269  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

183
184
185class CodeProcessor(ast.NodeTransformer):
186 def __init__(self, code, entry_point = None, force_rename = False):
187 self.code = code
188 self.entry_point = entry_point
189 self.classname = "global"
190 self.mode = "funcname"
191 self.ori_name = None
192 self.force_rename = force_rename
193
194
195 def visit_ClassDef(self, node):
196 self.classname = node.name
197 self.generic_visit(node)
198 self.classname = "global"
199
200 def visit_Name(self, node):
201 if self.mode == "funcname" and self.ori_name != None and node.id == self.ori_name:
202 node.id = "solution"
203
204 return node
205
206
207 def visit_FunctionDef(self, node):
208 #rename the first function generated as LLMs tend to generate extra useless code in the end of response
209 if not self.entry_point and self.mode == "funcname" and node.name == self.visitor.funcs[-1] and self.classname == "global":
210 self.ori_name = node.name
211 node.name = "solution"
212 elif self.entry_point and self.mode == "funcname" and node.name == self.entry_point and self.classname == "global":
213 self.ori_name = node.name
214 node.name = "solution"
215
216 self.generic_visit(node)
217
218 return node
219
220 def visit_Call(self, node):
221 if self.mode == "input" and isinstance(node.func, ast.Name) and node.func.id == "input":
222 new_node = ast.Name(id = "inputs")
223 ast.fix_missing_locations(new_node)
224 return new_node
225 else:
226 self.generic_visit(node)
227
228
229 return node
230
231 def run(self, no_modify = False):
232 try:
233 remover = CommentRemover(self.code)
234 new_code = remover.run().strip()
235 if len(new_code) == 0:
236 return -1, False
237 self.visitor = CodeVisitor(new_code)
238 self.visitor.run()
239 self.root = ast.parse(self.code)
240 if self.visitor.all_func_in_class:
241 if no_modify:
242 return ast.unparse(self.root), False

Callers 8

get_solutionsMethod · 0.90
get_groundtruthMethod · 0.90
process_solutionsMethod · 0.90

Calls

no outgoing calls

Tested by

no test coverage detected