vex-ast 0.1.0__py3-none-any.whl

vex_ast/ast/literals.py

@@ -0,0 +1,80 @@
+"""Literal value nodes for the AST."""
+
+from typing import Any, List, Optional, Union
+
+from .interfaces import IAstNode, IVisitor, T_VisitorResult, ILiteral
+from .core import Expression
+from ..utils.source_location import SourceLocation
+
+class Literal(Expression, ILiteral):
+    """Base class for all literal values."""
+    
+    _fields = ('value',)
+    
+    def __init__(self, value: Any, location: Optional[SourceLocation] = None):
+        super().__init__(location)
+        self.value = value
+    
+    def get_children(self) -> List[IAstNode]:
+        """Literals have no children."""
+        return []
+    
+    def get_value(self) -> Any:
+        """Get the literal value."""
+        return self.value
+
+class NumberLiteral(Literal):
+    """A numeric literal (integer or float)."""
+    
+    def __init__(self, value: Union[int, float], location: Optional[SourceLocation] = None):
+        super().__init__(value, location)
+    
+    def accept(self, visitor: IVisitor[T_VisitorResult]) -> T_VisitorResult:
+        return visitor.visit_numberliteral(self)
+    
+    def is_integer(self) -> bool:
+        """Check if this is an integer literal."""
+        return isinstance(self.value, int)
+    
+    def is_float(self) -> bool:
+        """Check if this is a float literal."""
+        return isinstance(self.value, float)
+
+class StringLiteral(Literal):
+    """A string literal."""
+    
+    def __init__(self, value: str, location: Optional[SourceLocation] = None):
+        super().__init__(value, location)
+    
+    def accept(self, visitor: IVisitor[T_VisitorResult]) -> T_VisitorResult:
+        return visitor.visit_stringliteral(self)
+    
+    def get_length(self) -> int:
+        """Get the length of the string."""
+        return len(self.value)
+
+class BooleanLiteral(Literal):
+    """A boolean literal (True or False)."""
+    
+    def __init__(self, value: bool, location: Optional[SourceLocation] = None):
+        super().__init__(value, location)
+    
+    def accept(self, visitor: IVisitor[T_VisitorResult]) -> T_VisitorResult:
+        return visitor.visit_booleanliteral(self)
+    
+    def is_true(self) -> bool:
+        """Check if this is a True literal."""
+        return self.value is True
+    
+    def is_false(self) -> bool:
+        """Check if this is a False literal."""
+        return self.value is False
+
+class NoneLiteral(Literal):
+    """A None literal."""
+    
+    def __init__(self, location: Optional[SourceLocation] = None):
+        super().__init__(None, location)
+    
+    def accept(self, visitor: IVisitor[T_VisitorResult]) -> T_VisitorResult:
+        return visitor.visit_noneliteral(self)

vex_ast/ast/navigator.py

@@ -0,0 +1,213 @@
+"""AST navigation utilities.
+
+This module provides utilities for navigating the AST structure,
+hiding implementation details and providing a more stable interface.
+"""
+
+from typing import Dict, List, Optional, Set, Iterator, TypeVar, Generic, Type, cast, Any, Union
+
+from .interfaces import IAstNode, IExpression, IStatement, ILiteral, IIdentifier, IFunctionCall
+from .core import Program
+from .expressions import (
+    Identifier, VariableReference, AttributeAccess, 
+    BinaryOperation, UnaryOperation, FunctionCall
+)
+from .statements import (
+    ExpressionStatement, Assignment, FunctionDefinition,
+    IfStatement, WhileLoop, ForLoop, ReturnStatement
+)
+from .literals import NumberLiteral, StringLiteral, BooleanLiteral
+from .vex_nodes import VexAPICall
+
+T = TypeVar('T', bound=IAstNode)
+
+class AstNavigator:
+    """Navigator for traversing and querying the AST."""
+    
+    def __init__(self, root: IAstNode):
+        """Initialize with a root node."""
+        self.root = root
+    
+    def find_all(self, node_type: Type[T]) -> List[T]:
+        """Find all nodes of a specific type in the AST.
+        
+        Args:
+            node_type: The type of node to find
+            
+        Returns:
+            List of nodes matching the type
+        """
+        result: List[T] = []
+        self._find_all_recursive(self.root, node_type, result)
+        return result
+    
+    def _find_all_recursive(self, node: IAstNode, node_type: Type[T], result: List[T]) -> None:
+        """Recursively find all nodes of a specific type."""
+        if isinstance(node, node_type):
+            result.append(cast(T, node))
+        
+        for child in node.get_children():
+            self._find_all_recursive(child, node_type, result)
+    
+    def find_first(self, node_type: Type[T]) -> Optional[T]:
+        """Find the first node of a specific type in the AST.
+        
+        Args:
+            node_type: The type of node to find
+            
+        Returns:
+            The first node matching the type, or None if not found
+        """
+        return next(self.find_iter(node_type), None)
+    
+    def find_iter(self, node_type: Type[T]) -> Iterator[T]:
+        """Find all nodes of a specific type as an iterator.
+        
+        Args:
+            node_type: The type of node to find
+            
+        Returns:
+            Iterator of nodes matching the type
+        """
+        for node in self._traverse_iter(self.root):
+            if isinstance(node, node_type):
+                yield cast(T, node)
+    
+    def _traverse_iter(self, node: IAstNode) -> Iterator[IAstNode]:
+        """Traverse the AST in pre-order."""
+        yield node
+        for child in node.get_children():
+            yield from self._traverse_iter(child)
+    
+    def find_parent(self, node: IAstNode) -> Optional[IAstNode]:
+        """Find the parent of a node.
+        
+        Args:
+            node: The node to find the parent of
+            
+        Returns:
+            The parent node, or None if not found or if the node is the root
+        """
+        return node.get_parent()
+    
+    def find_ancestors(self, node: IAstNode) -> List[IAstNode]:
+        """Find all ancestors of a node.
+        
+        Args:
+            node: The node to find ancestors of
+            
+        Returns:
+            List of ancestor nodes, from immediate parent to root
+        """
+        result: List[IAstNode] = []
+        current = node.get_parent()
+        while current:
+            result.append(current)
+            current = current.get_parent()
+        return result
+    
+    def find_siblings(self, node: IAstNode) -> List[IAstNode]:
+        """Find all siblings of a node.
+        
+        Args:
+            node: The node to find siblings of
+            
+        Returns:
+            List of sibling nodes, excluding the node itself
+        """
+        parent = node.get_parent()
+        if not parent:
+            return []
+        
+        return [child for child in parent.get_children() if child is not node]
+    
+    # Specific node type finders
+    
+    def find_identifiers(self) -> List[IIdentifier]:
+        """Find all identifiers in the AST."""
+        return self.find_all(Identifier)
+    
+    def find_function_calls(self) -> List[IFunctionCall]:
+        """Find all function calls in the AST."""
+        return self.find_all(FunctionCall)
+    
+    def find_vex_api_calls(self) -> List[VexAPICall]:
+        """Find all VEX API calls in the AST."""
+        return self.find_all(VexAPICall)
+    
+    def find_assignments(self) -> List[Assignment]:
+        """Find all assignments in the AST."""
+        return self.find_all(Assignment)
+    
+    def find_function_definitions(self) -> List[FunctionDefinition]:
+        """Find all function definitions in the AST."""
+        return self.find_all(FunctionDefinition)
+    
+    def find_if_statements(self) -> List[IfStatement]:
+        """Find all if statements in the AST."""
+        return self.find_all(IfStatement)
+    
+    def find_loops(self) -> List[Union[WhileLoop, ForLoop]]:
+        """Find all loops in the AST."""
+        while_loops = self.find_all(WhileLoop)
+        for_loops = self.find_all(ForLoop)
+        return cast(List[Union[WhileLoop, ForLoop]], while_loops + for_loops)
+    
+    def find_return_statements(self) -> List[ReturnStatement]:
+        """Find all return statements in the AST."""
+        return self.find_all(ReturnStatement)
+    
+    def find_literals(self) -> List[ILiteral]:
+        """Find all literals in the AST."""
+        return cast(List[ILiteral], 
+                   self.find_all(NumberLiteral) + 
+                   self.find_all(StringLiteral) + 
+                   self.find_all(BooleanLiteral))
+    
+    # Program-specific methods
+    
+    def get_statements(self) -> List[IStatement]:
+        """Get all top-level statements in the program."""
+        if isinstance(self.root, Program):
+            return self.root.get_statements()
+        return []
+    
+    def get_function_by_name(self, name: str) -> Optional[FunctionDefinition]:
+        """Find a function definition by name.
+        
+        Args:
+            name: The function name
+            
+        Returns:
+            The function definition, or None if not found
+        """
+        for func in self.find_function_definitions():
+            if func.get_name() == name:
+                return func
+        return None
+    
+    def get_variable_references(self, name: str) -> List[VariableReference]:
+        """Find all references to a variable.
+        
+        Args:
+            name: The variable name
+            
+        Returns:
+            List of variable references
+        """
+        return [ref for ref in self.find_all(VariableReference) if ref.name == name]
+    
+    def get_attribute_accesses(self, object_name: str) -> List[AttributeAccess]:
+        """Find all attribute accesses on an object.
+        
+        Args:
+            object_name: The object name
+            
+        Returns:
+            List of attribute accesses
+        """
+        result = []
+        for access in self.find_all(AttributeAccess):
+            if hasattr(access.object, 'name') and getattr(access.object, 'name') == object_name:
+                result.append(access)
+        return result

vex_ast/ast/operators.py

@@ -0,0 +1,136 @@
+"""Operator definitions for the AST."""
+
+from enum import Enum, auto
+from typing import Dict, Set, cast
+
+class OperatorType(Enum):
+    """Categories of operators."""
+    BINARY = auto()
+    UNARY = auto()
+    COMPARISON = auto()
+    LOGICAL = auto()
+
+class Operator(Enum):
+    """Enumeration of supported operators."""
+    # Binary Arithmetic
+    ADD = '+'
+    SUBTRACT = '-'
+    MULTIPLY = '*'
+    DIVIDE = '/'
+    MODULO = '%'
+    POWER = '**'
+    FLOOR_DIVIDE = '//'
+    
+    # Binary Bitwise
+    BITWISE_AND = '&'
+    BITWISE_OR = '|'
+    BITWISE_XOR = '^'
+    LEFT_SHIFT = '<<'
+    RIGHT_SHIFT = '>>'
+    
+    # Binary Comparison
+    EQUAL = '=='
+    NOT_EQUAL = '!='
+    LESS_THAN = '<'
+    LESS_EQUAL = '<='
+    GREATER_THAN = '>'
+    GREATER_EQUAL = '>='
+    
+    # Binary Logical
+    LOGICAL_AND = 'and'
+    LOGICAL_OR = 'or'
+    
+    # Membership/Identity
+    IN = 'in'
+    NOT_IN = 'not in'
+    IS = 'is'
+    IS_NOT = 'is not'
+
+    # Object access
+    MEMBER_ACCESS = '.'  # Add this line
+    
+    # Unary Arithmetic
+    UNARY_PLUS = '+ (unary)'
+    UNARY_MINUS = '- (unary)'
+    
+    # Unary Bitwise
+    BITWISE_NOT = '~'
+    
+    # Unary Logical
+    LOGICAL_NOT = 'not'
+    
+    @classmethod
+    def get_type(cls, op: 'Operator') -> OperatorType:
+        """Get the type of an operator."""
+        if op in UNARY_OPERATORS:
+            return OperatorType.UNARY
+        if op in COMPARISON_OPERATORS:
+            return OperatorType.COMPARISON
+        if op in LOGICAL_OPERATORS:
+            return OperatorType.LOGICAL
+        return OperatorType.BINARY
+
+# Operator sets for validation
+UNARY_OPERATORS: Set[Operator] = {
+    Operator.UNARY_PLUS,
+    Operator.UNARY_MINUS,
+    Operator.BITWISE_NOT,
+    Operator.LOGICAL_NOT
+}
+
+COMPARISON_OPERATORS: Set[Operator] = {
+    Operator.EQUAL,
+    Operator.NOT_EQUAL,
+    Operator.LESS_THAN,
+    Operator.LESS_EQUAL,
+    Operator.GREATER_THAN,
+    Operator.GREATER_EQUAL,
+    Operator.IN,
+    Operator.NOT_IN,
+    Operator.IS,
+    Operator.IS_NOT
+}
+
+LOGICAL_OPERATORS: Set[Operator] = {
+    Operator.LOGICAL_AND,
+    Operator.LOGICAL_OR
+}
+
+# Mappings for parser use
+PYTHON_BINARY_OP_MAP: Dict[str, Operator] = {
+    '+': Operator.ADD,
+    '-': Operator.SUBTRACT,
+    '*': Operator.MULTIPLY,
+    '/': Operator.DIVIDE,
+    '//': Operator.FLOOR_DIVIDE,
+    '%': Operator.MODULO,
+    '**': Operator.POWER,
+    '<<': Operator.LEFT_SHIFT,
+    '>>': Operator.RIGHT_SHIFT,
+    '|': Operator.BITWISE_OR,
+    '^': Operator.BITWISE_XOR,
+    '&': Operator.BITWISE_AND,
+    '@': Operator.MULTIPLY  # Matrix multiplication, maps to regular multiply for now
+}
+
+PYTHON_BINARY_OP_MAP['.'] = Operator.MEMBER_ACCESS
+
+PYTHON_UNARY_OP_MAP: Dict[str, Operator] = {
+    '+': Operator.UNARY_PLUS,
+    '-': Operator.UNARY_MINUS,
+    'not': Operator.LOGICAL_NOT,
+    '~': Operator.BITWISE_NOT
+}
+
+PYTHON_COMP_OP_MAP: Dict[str, Operator] = {
+    '==': Operator.EQUAL,
+    '!=': Operator.NOT_EQUAL,
+    '<': Operator.LESS_THAN,
+    '<=': Operator.LESS_EQUAL,
+    '>': Operator.GREATER_THAN,
+    '>=': Operator.GREATER_EQUAL,
+    'in': Operator.IN,
+    'not in': Operator.NOT_IN,
+    'is': Operator.IS,
+    'is not': Operator.IS_NOT
+}