#!/usr/bin/env python3 """ AST Parser for Ritz Framework Uses Python's ast module to parse components and generate IR """ import ast import inspect import json from typing import Any, Dict, List class RitzASTTransformer(ast.NodeVisitor): """Transform Python AST to Ritz IR""" def __init__(self): self.current_component = None self.ir = [] def visit_FunctionDef(self, node): """Process function definitions""" # Check if it has @component decorator if self.has_component_decorator(node): self.current_component = { 'type': 'component', 'name': node.name, 'args': [arg.arg for arg in node.args.args], 'body': [] } # Process function body for stmt in node.body: ir_node = self.process_statement(stmt) if ir_node: self.current_component['body'].append(ir_node) self.ir.append(self.current_component) self.current_component = None def has_component_decorator(self, node): """Check if function has @component decorator""" for decorator in node.decorator_list: if isinstance(decorator, ast.Name) and decorator.id == 'component': return True return False def process_statement(self, stmt): """Process a statement in component body""" if isinstance(stmt, ast.With): return self.process_with_statement(stmt) elif isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Yield): return self.process_yield(stmt.value) elif isinstance(stmt, ast.Assign): return self.process_assignment(stmt) elif isinstance(stmt, ast.FunctionDef): return self.process_inner_function(stmt) elif isinstance(stmt, ast.For): return self.process_for_loop(stmt) elif isinstance(stmt, ast.If): return self.process_if_statement(stmt) return None def process_with_statement(self, node): """Process 'with Element():' statements""" context = node.items[0].context_expr with_node = { 'type': 'with', 'element': self.process_expression(context), 'body': [] } # Process body for stmt in node.body: body_node = self.process_statement(stmt) if body_node: with_node['body'].append(body_node) return with_node def process_yield(self, node): """Process yield statements""" return { 'type': 'yield', 'value': self.process_expression(node.value) if node.value else None } def process_expression(self, expr): """Process an expression""" if isinstance(expr, ast.Call): return { 'type': 'call', 'func': self.get_name(expr.func), 'args': [self.process_expression(arg) for arg in expr.args], 'kwargs': { kw.arg: self.process_expression(kw.value) for kw in expr.keywords } } elif isinstance(expr, ast.Name): return { 'type': 'name', 'id': expr.id } elif isinstance(expr, ast.Constant): return { 'type': 'constant', 'value': expr.value } elif isinstance(expr, ast.JoinedStr): # f-strings return { 'type': 'fstring', 'parts': [self.process_expression(v) for v in expr.values] } elif isinstance(expr, ast.FormattedValue): return { 'type': 'formatted_value', 'value': self.process_expression(expr.value) } elif isinstance(expr, ast.Attribute): return { 'type': 'attribute', 'value': self.process_expression(expr.value), 'attr': expr.attr } else: return { 'type': 'unknown', 'ast_type': expr.__class__.__name__ } def process_assignment(self, node): """Process variable assignments (e.g., count = signal(0))""" target = node.targets[0] if isinstance(target, ast.Name): value = self.process_expression(node.value) # Check if it's a signal if (isinstance(value, dict) and value.get('type') == 'call' and value.get('func') == 'signal'): return { 'type': 'signal', 'name': target.id, 'initial': value['args'][0] if value['args'] else {'type': 'constant', 'value': None} } else: return { 'type': 'assign', 'target': target.id, 'value': value } def process_inner_function(self, node): """Process inner function definitions (e.g., event handlers)""" return { 'type': 'function', 'name': node.name, 'args': [arg.arg for arg in node.args.args], 'body': [self.process_statement(stmt) for stmt in node.body] } def process_for_loop(self, node): """Process for loops""" return { 'type': 'for', 'target': self.get_name(node.target), 'iter': self.process_expression(node.iter), 'body': [self.process_statement(stmt) for stmt in node.body] } def process_if_statement(self, node): """Process if statements""" return { 'type': 'if', 'test': self.process_expression(node.test), 'body': [self.process_statement(stmt) for stmt in node.body], 'else': [self.process_statement(stmt) for stmt in node.orelse] if node.orelse else [] } def get_name(self, node): """Get name from Name or Attribute node""" if isinstance(node, ast.Name): return node.id elif isinstance(node, ast.Attribute): return f"{self.get_name(node.value)}.{node.attr}" return 'unknown' def parse_component_source(source: str) -> List[Dict]: """Parse component source code and return IR""" tree = ast.parse(source) transformer = RitzASTTransformer() transformer.visit(tree) return transformer.ir def parse_component_function(func) -> Dict: """Parse a component function and return its IR""" source = inspect.getsource(func) ir = parse_component_source(source) return ir[0] if ir else None def visualize_ast(source: str): """Visualize Python AST for debugging""" tree = ast.parse(source) return ast.dump(tree, indent=2) # Test it! if __name__ == '__main__': test_source = ''' @component def Counter(): count = signal(0) def increment(): count += 1 with Div(class_='counter'): yield H1('Counter') yield P(f'Count: {count}') yield Button('Click me', onclick=increment) ''' print("=== Python AST (for reference) ===") print(visualize_ast(test_source)[:500] + "...\n") print("=== Ritz IR ===") ir = parse_component_source(test_source) print(json.dumps(ir, indent=2)) print("\n=== Test with actual component ===") # Import our framework from framework import component, signal # Mock elements for testing def Div(**kwargs): return {'element': 'div', 'props': kwargs} def H1(*args): return {'element': 'h1', 'children': args} def P(*args): return {'element': 'p', 'children': args} def Button(*args, **kwargs): return {'element': 'button', 'children': args, 'props': kwargs} @component def TestComponent(): count = signal(0) with Div(class_='test'): yield H1('Test') yield P(f'Value: {count}') for i in range(3): yield Button(f'Button {i}') ir = parse_component_function(TestComponent) print("\nParsed from real function:") print(json.dumps(ir, indent=2)[:1000] + "...")