206 lines
6.8 KiB
Python
206 lines
6.8 KiB
Python
import ast
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
def find_definition_range(source: str, symbol_path: str) -> tuple[int, int] | None:
|
|
"""
|
|
Finds the start and end line numbers for a given symbol path (e.g., 'ClassName', 'ClassName.method', 'function_name').
|
|
[C: scripts/py_struct_tools.py:py_remove_def, scripts/py_struct_tools.py:py_move_def]
|
|
"""
|
|
try:
|
|
tree = ast.parse(source)
|
|
except SyntaxError:
|
|
return None
|
|
parts = symbol_path.split('.')
|
|
current_nodes = tree.body
|
|
target_node = None
|
|
for i, part in enumerate(parts):
|
|
found = False
|
|
for node in current_nodes:
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name == part:
|
|
if i == len(parts) - 1:
|
|
target_node = node
|
|
found = True
|
|
break
|
|
else:
|
|
if isinstance(node, ast.ClassDef):
|
|
current_nodes = node.body
|
|
found = True
|
|
break
|
|
if not found:
|
|
return None
|
|
if target_node:
|
|
start = target_node.lineno
|
|
if target_node.decorator_list:
|
|
start = target_node.decorator_list[0].lineno
|
|
return (start, target_node.end_lineno)
|
|
return None
|
|
|
|
def shift_indentation(content: str, target_depth: int) -> str:
|
|
"""
|
|
Shifts and normalizes the indentation of a code block to 1-space units.
|
|
Detects the base indentation and scales relative indentation to 1-space.
|
|
[C: scripts/py_struct_tools.py:py_add_def]
|
|
"""
|
|
lines = content.splitlines()
|
|
if not lines:
|
|
return ""
|
|
|
|
# 1. Find min indent of non-empty lines
|
|
min_indent = sys.maxsize
|
|
for line in lines:
|
|
if line.strip():
|
|
indent = len(line) - len(line.lstrip())
|
|
if indent < min_indent:
|
|
min_indent = indent
|
|
if min_indent == sys.maxsize:
|
|
min_indent = 0
|
|
|
|
# 2. Try to detect indentation unit (width)
|
|
indent_unit = 0
|
|
for line in lines:
|
|
if line.strip():
|
|
indent = len(line) - len(line.lstrip())
|
|
rel_indent = indent - min_indent
|
|
if rel_indent > 0:
|
|
if indent_unit == 0:
|
|
indent_unit = rel_indent
|
|
else:
|
|
import math
|
|
indent_unit = math.gcd(indent_unit, rel_indent)
|
|
|
|
if indent_unit == 0:
|
|
indent_unit = 1
|
|
|
|
shifted_lines = []
|
|
for line in lines:
|
|
if line.strip():
|
|
indent = len(line) - len(line.lstrip())
|
|
rel_level = (indent - min_indent) // indent_unit
|
|
shifted_lines.append(" " * (target_depth + rel_level) + line.lstrip())
|
|
else:
|
|
shifted_lines.append("")
|
|
|
|
return "\n".join(shifted_lines) + ("\n" if content.endswith("\n") else "")
|
|
|
|
def py_remove_def(filepath: str, symbol_path: str) -> str:
|
|
"""
|
|
Removes a definition from a file. Returns error message or empty string on success.
|
|
"""
|
|
if not os.path.exists(filepath):
|
|
return f"File not found: {filepath}"
|
|
with open(filepath, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
range_info = find_definition_range(content, symbol_path)
|
|
if not range_info:
|
|
return f"Symbol '{symbol_path}' not found in {filepath}"
|
|
start, end = range_info
|
|
lines = content.splitlines(True)
|
|
del lines[start-1:end]
|
|
with open(filepath, 'w', encoding='utf-8', newline='') as f:
|
|
f.writelines(lines)
|
|
return ""
|
|
|
|
def py_add_def(filepath: str, symbol_path: str, new_content: str, anchor_type: str, anchor_symbol: str = None) -> str:
|
|
"""
|
|
Adds a definition to a file.
|
|
anchor_type: 'before', 'after', 'top', 'bottom'
|
|
anchor_symbol: if 'before' or 'after', the symbol to anchor to.
|
|
symbol_path: context path (e.g. 'ClassName' or '' for module level)
|
|
"""
|
|
if not os.path.exists(filepath):
|
|
return f"File not found: {filepath}"
|
|
with open(filepath, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
lines = content.splitlines(True)
|
|
target_depth = 0
|
|
insert_idx = -1
|
|
if symbol_path:
|
|
range_info = find_definition_range(content, symbol_path)
|
|
if not range_info:
|
|
return f"Context symbol '{symbol_path}' not found"
|
|
c_start, c_end = range_info
|
|
target_depth = 1
|
|
if anchor_type == 'top':
|
|
insert_idx = c_start
|
|
for i in range(c_start, c_end):
|
|
if '"""' in lines[i] or "'''" in lines[i]:
|
|
for j in range(i, c_end):
|
|
if lines[j].count('"""') + lines[j].count("'''") == 2 or (j > i and ('"""' in lines[j] or "'''" in lines[j])):
|
|
insert_idx = j + 1
|
|
break
|
|
break
|
|
elif anchor_type == 'bottom':
|
|
insert_idx = c_end
|
|
else:
|
|
if anchor_type == 'top':
|
|
insert_idx = 0
|
|
elif anchor_type == 'bottom':
|
|
insert_idx = len(lines)
|
|
if anchor_type in ('before', 'after') and anchor_symbol:
|
|
full_anchor_path = f"{symbol_path}.{anchor_symbol}" if symbol_path else anchor_symbol
|
|
range_info = find_definition_range(content, full_anchor_path)
|
|
if not range_info:
|
|
return f"Anchor symbol '{anchor_symbol}' not found"
|
|
a_start, a_end = range_info
|
|
insert_idx = a_start - 1 if anchor_type == 'before' else a_end
|
|
if insert_idx == -1:
|
|
return f"Invalid anchor_type '{anchor_type}'"
|
|
formatted_content = shift_indentation(new_content, target_depth)
|
|
if not formatted_content.endswith('\n'):
|
|
formatted_content += '\n'
|
|
# Ensure at least one blank line between definitions
|
|
if insert_idx > 0 and lines[insert_idx-1].strip():
|
|
formatted_content = '\n' + formatted_content
|
|
if insert_idx < len(lines) and lines[insert_idx].strip():
|
|
formatted_content = formatted_content + '\n'
|
|
lines.insert(insert_idx, formatted_content)
|
|
with open(filepath, 'w', encoding='utf-8', newline='') as f:
|
|
f.writelines(lines)
|
|
return ""
|
|
|
|
def py_move_def(src_file: str, dest_file: str, symbol_path: str, dest_symbol_path: str, anchor_type: str, anchor_symbol: str = None) -> str:
|
|
"""
|
|
Moves a definition from one location to another.
|
|
"""
|
|
if not os.path.exists(src_file):
|
|
return f"Source file not found: {src_file}"
|
|
with open(src_file, 'r', encoding='utf-8') as f:
|
|
src_content = f.read()
|
|
range_info = find_definition_range(src_content, symbol_path)
|
|
if not range_info:
|
|
return f"Symbol '{symbol_path}' not found in {src_file}"
|
|
start, end = range_info
|
|
src_lines = src_content.splitlines(True)
|
|
moved_content = "".join(src_lines[start-1:end])
|
|
err = py_add_def(dest_file, dest_symbol_path, moved_content, anchor_type, anchor_symbol)
|
|
if err:
|
|
return err
|
|
return py_remove_def(src_file, symbol_path)
|
|
|
|
def py_region_wrap(filepath: str, start_line: int, end_line: int, region_name: str) -> str:
|
|
"""
|
|
Wraps a line range in region tags.
|
|
"""
|
|
if not os.path.exists(filepath):
|
|
return f"File not found: {filepath}"
|
|
with open(filepath, 'r', encoding='utf-8') as f:
|
|
lines = f.readlines()
|
|
if start_line < 1 or end_line > len(lines) or start_line > end_line:
|
|
return f"Invalid line range: {start_line}-{end_line}"
|
|
indent = ""
|
|
first_line = lines[start_line-1]
|
|
if first_line.strip():
|
|
indent = first_line[:len(first_line) - len(first_line.lstrip())]
|
|
# Insert bottom first to avoid shifting top
|
|
lines.insert(end_line, f"{indent}#endregion: {region_name}\n")
|
|
lines.insert(start_line-1, f"{indent}#region: {region_name}\n")
|
|
with open(filepath, 'w', encoding='utf-8', newline='') as f:
|
|
f.writelines(lines)
|
|
return ""
|
|
|
|
if __name__ == "__main__":
|
|
# Simple CLI for internal use
|
|
pass
|