99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
import tree_sitter
|
|
import tree_sitter_python
|
|
|
|
class ASTParser:
|
|
def __init__(self, language: str):
|
|
self.language = tree_sitter.Language(tree_sitter_python.language())
|
|
self.parser = tree_sitter.Parser(self.language)
|
|
|
|
def parse(self, code: str) -> tree_sitter.Tree:
|
|
return self.parser.parse(bytes(code, "utf8"))
|
|
|
|
def get_curated_view(self, code: str) -> str:
|
|
tree = self.parse(code)
|
|
edits = []
|
|
|
|
def is_docstring(node):
|
|
if node.type == "expression_statement" and node.child_count > 0:
|
|
if node.children[0].type == "string":
|
|
return True
|
|
return False
|
|
|
|
def has_core_logic_decorator(node):
|
|
parent = node.parent
|
|
if parent and parent.type == "decorated_definition":
|
|
for child in parent.children:
|
|
if child.type == "decorator":
|
|
if "@core_logic" in code[child.start_byte:child.end_byte]:
|
|
return True
|
|
return False
|
|
|
|
def has_hot_comment(func_node):
|
|
print(f"Checking {code[func_node.start_byte:func_node.start_byte+20].strip()}...")
|
|
stack = [func_node]
|
|
while stack:
|
|
curr = stack.pop()
|
|
if curr.type == "comment":
|
|
comment_text = code[curr.start_byte:curr.end_byte]
|
|
print(f" Found comment: {comment_text}")
|
|
if "[HOT]" in comment_text:
|
|
print(" [HOT] FOUND!")
|
|
return True
|
|
for child in curr.children:
|
|
stack.append(child)
|
|
return False
|
|
|
|
def walk(node):
|
|
if node.type == "function_definition":
|
|
body = node.child_by_field_name("body")
|
|
if body and body.type == "block":
|
|
preserve = has_core_logic_decorator(node) or has_hot_comment(node)
|
|
print(f"Function {code[node.start_byte:node.start_byte+20].strip()}, preserve={preserve}")
|
|
if not preserve:
|
|
indent = " " * body.start_point.column
|
|
first_stmt = None
|
|
for child in body.children:
|
|
if child.type != "comment":
|
|
first_stmt = child
|
|
break
|
|
if first_stmt and is_docstring(first_stmt):
|
|
start_byte = first_stmt.end_byte
|
|
end_byte = body.end_byte
|
|
if end_byte > start_byte:
|
|
edits.append((start_byte, end_byte, "\\n" + indent + "..."))
|
|
else:
|
|
start_byte = body.start_byte
|
|
end_byte = body.end_byte
|
|
edits.append((start_byte, end_byte, "..."))
|
|
for child in node.children:
|
|
walk(child)
|
|
walk(tree.root_node)
|
|
edits.sort(key=lambda x: x[0], reverse=True)
|
|
code_bytes = bytearray(code, "utf8")
|
|
for start, end, replacement in edits:
|
|
code_bytes[start:end] = bytes(replacement, "utf8")
|
|
return code_bytes.decode("utf8")
|
|
|
|
parser = ASTParser("python")
|
|
code = '''
|
|
@core_logic
|
|
def core_func():
|
|
"""Core logic doc."""
|
|
print("this should be preserved")
|
|
return True
|
|
|
|
def hot_func():
|
|
# [HOT]
|
|
print("this should also be preserved")
|
|
return 42
|
|
|
|
def normal_func():
|
|
"""Normal doc."""
|
|
print("this should be stripped")
|
|
return None
|
|
'''
|
|
|
|
result = parser.get_curated_view(code)
|
|
print("\\n--- RESULT ---\\n")
|
|
print(result)
|