# parser.py import ast def get_category(node): """Determine the category of an AST node.""" if isinstance(node, (ast.Import, ast.ImportFrom)): return 'import' elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)): return 'assignment' elif isinstance(node, ast.FunctionDef): return 'function' elif isinstance(node, ast.AsyncFunctionDef): return 'async_function' elif isinstance(node, ast.ClassDef): return 'class' elif isinstance(node, ast.If): return 'if' elif isinstance(node, ast.While): return 'while' elif isinstance(node, ast.For): return 'for' elif isinstance(node, ast.Try): return 'try' elif isinstance(node, ast.Return): return 'return' elif isinstance(node, ast.Expr): return 'expression' else: return 'other' def create_vector(category, level, location, total_lines, parent_path): """Create a vector optimized for role similarity.""" category_map = { 'import': 1, 'assignment': 2, 'function': 3, 'async_function': 4, 'class': 5, 'if': 6, 'while': 7, 'for': 8, 'try': 9, 'expression': 10, 'spacer': 11, 'other': 12, 'elif': 13, 'else': 14, 'except': 15, 'finally': 16, 'return': 17 } category_id = category_map.get(category, 0) start_line, end_line = location span = (end_line - start_line + 1) / total_lines center_pos = ((start_line + end_line) / 2) / total_lines parent_depth = len(parent_path) parent_weight = sum(category_map.get(parent.split('[')[0].lower(), 0) * (1 / (i + 1)) for i, parent in enumerate(parent_path)) / max(1, len(category_map)) return [category_id, level, center_pos, span, parent_depth, parent_weight] def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=None, counters=None): if total_lines is None: total_lines = len(lines) if parent_path is None: parent_path = [] if counters is None: counters = {cat: 0 for cat in ['import', 'assignment', 'function', 'async_function', 'class', 'if', 'while', 'for', 'try', 'return', 'expression', 'other', 'spacer', 'elif', 'else', 'except', 'finally']} parts = [] start_line = getattr(node, 'lineno', prev_end + 1) end_line = getattr(node, 'end_lineno', start_line) category = get_category(node) counters[category] += 1 node_id = f"{category.capitalize()}[{counters[category]}]" # Spacer before node if start_line > prev_end + 1: spacer_lines = lines[prev_end:start_line - 1] counters['spacer'] += 1 spacer_node_id = f"Spacer[{counters['spacer']}]" parts.append({ 'category': 'spacer', 'source': ''.join(spacer_lines), 'location': (prev_end + 1, start_line - 1), 'level': level, 'vector': create_vector('spacer', level, (prev_end + 1, start_line - 1), total_lines, parent_path), 'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level', 'node_id': spacer_node_id }) # Current node's header (e.g., 'def', 'if') current_path = parent_path + [node_id] node_lines = lines[start_line - 1:end_line] parts.append({ 'category': category, 'source': node_lines[0], # Only the header line 'location': (start_line, start_line), 'level': level, 'vector': create_vector(category, level, (start_line, start_line), total_lines, current_path), 'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level', 'node_id': node_id }) # Process nested bodies nested_prev_end = start_line for attr in ('body', 'orelse', 'handlers', 'finalbody'): if hasattr(node, attr) and getattr(node, attr): for child in getattr(node, attr): if attr == 'orelse' and isinstance(node, ast.If) and child.lineno != start_line: sub_category = 'elif' if 'elif' in lines[child.lineno - 1] else 'else' counters[sub_category] += 1 sub_node_id = f"{sub_category.capitalize()}[{counters[sub_category]}]" sub_start = child.lineno sub_end = getattr(child, 'end_lineno', sub_start) parts.append({ 'category': sub_category, 'source': lines[sub_start - 1], 'location': (sub_start, sub_start), 'level': level, 'vector': create_vector(sub_category, level, (sub_start, sub_start), total_lines, current_path), 'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level', 'node_id': sub_node_id }) child_parts = parse_node(child, lines, sub_start, level + 1, total_lines, current_path, counters) parts.extend(child_parts) nested_prev_end = max(nested_prev_end, child_parts[-1]['location'][1] if child_parts else sub_start) elif attr == 'handlers' and isinstance(child, ast.ExceptHandler): counters['except'] += 1 sub_node_id = f"Except[{counters['except']}]" sub_start = child.lineno sub_end = getattr(child, 'end_lineno', sub_start) parts.append({ 'category': 'except', 'source': lines[sub_start - 1], 'location': (sub_start, sub_start), 'level': level, 'vector': create_vector('except', level, (sub_start, sub_start), total_lines, current_path), 'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level', 'node_id': sub_node_id }) child_parts = parse_node(child, lines, sub_start, level + 1, total_lines, current_path, counters) parts.extend(child_parts) nested_prev_end = max(nested_prev_end, child_parts[-1]['location'][1] if child_parts else sub_start) elif attr == 'finalbody': counters['finally'] += 1 sub_node_id = f"Finally[{counters['finally']}]" sub_start = child.lineno sub_end = getattr(child, 'end_lineno', sub_start) parts.append({ 'category': 'finally', 'source': lines[sub_start - 1], 'location': (sub_start, sub_start), 'level': level, 'vector': create_vector('finally', level, (sub_start, sub_start), total_lines, current_path), 'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level', 'node_id': sub_node_id }) child_parts = parse_node(child, lines, sub_start, level + 1, total_lines, current_path, counters) parts.extend(child_parts) nested_prev_end = max(nested_prev_end, child_parts[-1]['location'][1] if child_parts else sub_start) else: child_parts = parse_node(child, lines, nested_prev_end, level + 1, total_lines, current_path, counters) parts.extend(child_parts) nested_prev_end = child_parts[-1]['location'][1] if child_parts else nested_prev_end # Update end_line of the parent node to reflect its full scope if nested_prev_end > start_line: parts[-1]['location'] = (start_line, nested_prev_end) parts[-1]['source'] = ''.join(lines[start_line - 1:nested_prev_end]) parts[-1]['vector'] = create_vector(category, level, (start_line, nested_prev_end), total_lines, current_path) return parts def parse_python_code(code): lines = code.splitlines(keepends=True) total_lines = len(lines) try: tree = ast.parse(code) except SyntaxError: return [{'category': 'error', 'source': 'Invalid Python code', 'location': (1, 1), 'level': 0, 'vector': [0, 0, 1.0, 0.0, 0, 0], 'parent_path': 'Top-Level', 'node_id': 'Error[1]'}] parts = [] prev_end = 0 for stmt in tree.body: stmt_parts = parse_node(stmt, lines, prev_end, total_lines=total_lines) parts.extend(stmt_parts) prev_end = stmt_parts[-1]['location'][1] if prev_end < total_lines: remaining_lines = lines[prev_end:] counters = {'spacer': 0} counters['spacer'] += 1 spacer_node_id = f"Spacer[{counters['spacer']}]" parts.append({ 'category': 'spacer', 'source': ''.join(remaining_lines), 'location': (prev_end + 1, total_lines + 1), 'level': 0, 'vector': create_vector('spacer', 0, (prev_end + 1, total_lines + 1), total_lines, []), 'parent_path': 'Top-Level', 'node_id': spacer_node_id }) return parts