257 lines
8.3 KiB
Python
257 lines
8.3 KiB
Python
#!/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] + "...")
|