File size: 9,450 Bytes
38562d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | import ast, re, textwrap, json
class Rejected(ValueError):
pass
def strip_docstrings(code):
code = textwrap.dedent(code).strip()
try:
tree = ast.parse(code)
except (SyntaxError, ValueError, RecursionError) as e:
raise Rejected("syntax") from e
if len(tree.body) != 1 or not isinstance(tree.body[0], (ast.FunctionDef, ast.AsyncFunctionDef)):
raise Rejected("not_single_function")
# AST columns are UTF-8 byte offsets, not Python character offsets.
lines = code.encode("utf-8").splitlines(keepends=True)
offsets = [0]
for line in lines:
offsets.append(offsets[-1] + len(line))
edits = []
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)):
if node.body and isinstance(node.body[0], ast.Expr):
first = node.body[0]
if isinstance(first.value, ast.Constant) and isinstance(first.value.value, str):
a = offsets[first.lineno - 1] + first.col_offset
b = offsets[first.end_lineno - 1] + first.end_col_offset
# pass also handles docstring-only functions and one-line semicolons.
edits.append((a, b, b"pass"))
raw = code.encode("utf-8")
for a, b, replacement in sorted(edits, reverse=True):
raw = raw[:a] + replacement + raw[b:]
clean = raw.decode("utf-8")
tree = ast.parse(clean)
useful = [n for n in tree.body[0].body if not isinstance(n, ast.Pass)]
if not useful or all(isinstance(n, ast.Expr) and isinstance(n.value, ast.Constant) and n.value.value is Ellipsis for n in useful):
raise Rejected("empty_function")
return clean
def _bound_names(function):
names = {a.arg for a in [*function.args.posonlyargs,*function.args.args,*function.args.kwonlyargs]}
if function.args.vararg:
names.add(function.args.vararg.arg)
if function.args.kwarg:
names.add(function.args.kwarg.arg)
class Bindings(ast.NodeVisitor):
def visit_Name(self,node):
if isinstance(node.ctx,(ast.Store,ast.Del)):
names.add(node.id)
def visit_FunctionDef(self,node): names.add(node.name)
visit_AsyncFunctionDef = visit_FunctionDef
def visit_ClassDef(self,node): names.add(node.name)
def visit_Lambda(self,node): pass
def visit_Import(self,node):
names.update(a.asname or a.name.split('.')[0] for a in node.names)
visit_ImportFrom = visit_Import
visitor = Bindings()
for node in function.body:
visitor.visit(node)
return names
def _rename_tree(tree,new_name,positions=None):
root = tree.body[0]
old = root.name
root.name = new_name
scopes = [_bound_names(root)]
class Canonical(ast.NodeTransformer):
def visit_Name(self,node):
if node.id == old and not any(old in scope for scope in scopes):
if positions is not None:positions.append((node.lineno,node.col_offset,node.end_lineno,node.end_col_offset))
node.id = new_name
return node
def visit_FunctionDef(self,node):
scopes.append(_bound_names(node))
node = self.generic_visit(node)
scopes.pop()
return node
visit_AsyncFunctionDef = visit_FunctionDef
visitor = Canonical()
root.body = [visitor.visit(n) for n in root.body if not isinstance(n,ast.Pass)]
ast.fix_missing_locations(tree)
return tree
def canonical_code(code):
tree = _rename_tree(ast.parse(strip_docstrings(code)),"_function")
return ast.unparse(tree)
def normalize_bindings(code):
tree=ast.parse(canonical_code(code))
if len(tree.body)!=1 or not isinstance(tree.body[0],ast.FunctionDef):
raise ValueError('Expected one Python function')
root=tree.body[0]
parameters=[*root.args.posonlyargs,*root.args.args,*root.args.kwonlyargs]
if root.args.vararg:parameters.append(root.args.vararg)
if root.args.kwarg:parameters.append(root.args.kwarg)
bound={a.arg for a in parameters}|{n.id for n in ast.walk(root) if isinstance(n,ast.Name) and isinstance(n.ctx,ast.Store)}
bound|={n.name for n in ast.walk(root) if isinstance(n,ast.FunctionDef)}
free={n.id for n in ast.walk(root) if isinstance(n,ast.Name) and isinstance(n.ctx,ast.Load)}-bound
prefix=''
while any(re.fullmatch(re.escape(prefix)+r'(?:arg|local)\d+',name) for name in free):prefix+='_'
mapping={a.arg:f'{prefix}arg{i}' for i,a in enumerate(parameters)}
local=0
for node in ast.walk(root):
if isinstance(node,ast.Name) and isinstance(node.ctx,ast.Store) and node.id not in mapping:
mapping[node.id]=f'{prefix}local{local}';local+=1
elif isinstance(node,ast.FunctionDef) and node is not root and node.name not in mapping:
mapping[node.name]=f'{prefix}local{local}';local+=1
# One simultaneous bijection preserves lexical shadowing and comprehensions.
for node in ast.walk(root):
if isinstance(node,ast.Name):node.id=mapping.get(node.id,node.id)
elif isinstance(node,ast.arg):node.arg=mapping.get(node.arg,node.arg)
elif isinstance(node,ast.FunctionDef) and node is not root:node.name=mapping.get(node.name,node.name)
elif isinstance(node,ast.Call) and isinstance(node.func,ast.Name) and node.func.id==root.name:
for keyword in node.keywords:keyword.arg=mapping.get(keyword.arg,keyword.arg)
ast.fix_missing_locations(tree)
return ast.unparse(tree),mapping
def restore_names(text,mapping):
"""Restore identifier words outside string literals, including incomplete streams."""
reverse={v:k for k,v in mapping.items()}
pattern=re.compile(r'\b(?:'+'|'.join(re.escape(name) for name in reverse)+r')\b') if reverse else None
parts=[];index=0;start=0;quote=None;escaped=False
def replace(chunk):
return pattern.sub(lambda m:reverse[m.group()],chunk) if pattern else chunk
while index<len(text):
char=text[index]
if quote is None and char in ['"',"'"]:
parts.append(replace(text[start:index]));start=index;quote=char;escaped=False
elif quote is not None:
if escaped:escaped=False
elif char=='\\':escaped=True
elif char==quote:
parts.append(text[start:index+1]);start=index+1;quote=None
index+=1
parts.append(text[start:] if quote else replace(text[start:]))
return ''.join(parts)
def reasoning_prompt(code, inputs=None):
instruction = ('Explain all branches and track variable dependencies. Write Steps:, '
'Result expression: with equivalent Python, then Summary: with the behavior.')
if inputs is not None:
instruction = ('Execute for these arguments: ' + json.dumps(inputs, sort_keys=True) +
'. Write Trace: with executed lines and values, then Return: with its Python value.')
return 'Python function:\n' + canonical_code(code) + '\n\n' + instruction + '\nAnswer:\n'
"""Optional name normalization; arbitrary input is never rejected as non-Python."""
SNIPPET_ENABLED = True
SNIPPET_INSTRUCTION = 'Explain this Python code step by step. Write Steps:, Result expression: for printed values or the final value, then Summary: with what the code does.'
def normalize_snippet(code):
tree = ast.parse(code)
if not tree.body:
raise ValueError('Empty code')
bound = {n.id for n in ast.walk(tree) if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Store)}
free = {n.id for n in ast.walk(tree) if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load)} - bound
prefix = ''
while any(re.fullmatch(re.escape(prefix)+r'local\d+', n) for n in free):
prefix += '_'
mapping = {}
for node in ast.walk(tree):
if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store) and node.id not in mapping:
mapping[node.id] = f'{prefix}local{len(mapping)}'
for node in ast.walk(tree):
if isinstance(node, ast.Name):
node.id = mapping.get(node.id, node.id)
return ast.unparse(ast.fix_missing_locations(tree)), mapping
def snippet_prompt(code):
return 'Python code:\n'+code+'\n\n'+SNIPPET_INSTRUCTION+'\nAnswer:\n'
"""Optional name normalization; arbitrary input is never rejected as non-Python."""
def prepare(code):
try:
source, mapping = normalize_bindings(code)
prompt = reasoning_prompt(source)
mode = 'normalized_python'
except (SyntaxError, ValueError, RecursionError):
if globals().get('SNIPPET_ENABLED', False):
try:
source, mapping = normalize_snippet(code)
prompt, mode = snippet_prompt(source), 'snippet'
except (SyntaxError, ValueError, RecursionError):
source, mapping, prompt, mode = code, {}, code, 'raw_text'
else:
# Pass non-code, partial code and other text directly to the tokenizer.
source, mapping, prompt, mode = code, {}, code, 'raw_text'
return json.dumps({'prompt': prompt, 'mapping': mapping,
'canonical': source, 'input_mode': mode})
|