Add expression evaluation protocol

grandcypher/__init__.py

@@ -26,7 +26,7 @@
     Compare as IndexerCompare, OR as IndexerOr,
     AND as IndexerAnd, ArrayAttributeIndexer, IndexerConditionAST,
     UnsupportedOp as IndexerUnsupportedOp, IndexerConditionRunner)
-from .types import EntityRef, AttributeRef, IDRef
+from .types import Expression, EntityRef, AttributeRef, IDRef
 
 
 _GrandCypherGrammar = Lark(
@@ -251,7 +251,7 @@ def __init__(self, left, op: str, right):
         self.op = op
         self.right = right
 
-    def resolve(self, match, host, return_edges):
+    def evaluate(self, match, host, return_edges, scope=None):
         left_val = _resolve_operand(self.left, match, host, return_edges)
         right_val = _resolve_operand(self.right, match, host, return_edges)
         try:
@@ -264,28 +264,8 @@ def __str__(self):
 
 
 def _resolve_operand(operand, match, host, return_edges):
-    if isinstance(operand, ArithmeticExpression):
-        return operand.resolve(match, host, return_edges)
-    if isinstance(operand, IDRef):
-        if operand.entity_name in match.node_mappings:
-            return match.node_mappings[operand.entity_name]
-        raise IndexError(f"Entity {operand.entity_name} not in match.")
-    if isinstance(operand, AttributeRef):
-        entity_name = operand.entity_name
-        attribute = operand.attribute
-        if entity_name in match.node_mappings:
-            host_node_id = match.node_mappings[entity_name]
-            return get_node_from_host(host, host_node_id, attribute)
-        if entity_name in return_edges:
-            edge_mapping = return_edges[entity_name]
-            host_edges = match.mth.edge(*edge_mapping).edges
-            return get_edge_from_host(host, host_edges, attribute)
-        raise IndexError(f"Entity {operand} not in graph.")
-    if isinstance(operand, EntityRef):
-        raise TypeError(
-            f"Cannot use bare entity '{operand}' in a comparison. "
-            f"Use a property like '{operand}.attribute' or ID({operand})."
-        )
+    if isinstance(operand, Expression):
+        return operand.evaluate(match, host, return_edges)
     return operand
 
 

grandcypher/test_types.py

@@ -1,9 +1,11 @@
 import pickle
 
 import networkx as nx
+import pytest
 
-from grandcypher import GrandCypher
-from grandcypher.types import AttributeRef, IDRef, EntityRef
+from grandcypher import ArithmeticExpression, GrandCypher
+from grandcypher.struct import EdgeHopKey, EdgeMapping, HopSpec, Match
+from grandcypher.types import AttributeRef, Expression, IDRef, EntityRef
 
 
 def test_attribute_ref_pickle_roundtrip():
@@ -31,6 +33,54 @@ def test_entity_ref_pickle_roundtrip():
     assert restored.entity_name == 'a'
 
 
+def test_typed_references_implement_expression_protocol():
+    assert isinstance(AttributeRef("a", "name"), Expression)
+    assert isinstance(IDRef("a"), Expression)
+    assert isinstance(EntityRef("a"), Expression)
+
+
+def test_typed_reference_evaluation():
+    host = nx.DiGraph()
+    host.add_node("alice", age=30)
+    match = Match(node_mappings={"a": "alice"}, where_results=None, edge_mapping=None)
+
+    assert AttributeRef("a", "age").evaluate(match, host, {}) == 30
+    assert IDRef("a").evaluate(match, host, {}) == "alice"
+    with pytest.raises(TypeError, match="Cannot use bare entity"):
+        EntityRef("a").evaluate(match, host, {})
+
+
+@pytest.mark.parametrize("graph_type", [nx.DiGraph, nx.MultiDiGraph])
+def test_edge_attribute_reference_evaluation(graph_type):
+    host = graph_type()
+    edge_key = host.add_edge("alice", "bob", since=2020)
+    hop = HopSpec(edge_id=("a", "b"), nodes=("a", "b"), hop_count=1)
+    match = Match(
+        node_mappings={"a": "alice", "b": "bob"},
+        where_results=None,
+        edge_mapping=EdgeMapping(
+            edge_hop_map={("a", "b"): hop},
+            edge_key_map={
+                ("a", "b"): EdgeHopKey(("a", "b"), (edge_key,)),
+            },
+        ),
+    )
+
+    assert AttributeRef("r", "since").evaluate(
+        match, host, {"r": ("a", "b")}
+    ) == 2020
+
+
+def test_arithmetic_expression_implements_protocol():
+    host = nx.DiGraph()
+    host.add_node("alice", age=30)
+    match = Match(node_mappings={"a": "alice"}, where_results=None, edge_mapping=None)
+    expression = ArithmeticExpression(AttributeRef("a", "age"), "+", 2)
+
+    assert isinstance(expression, Expression)
+    assert expression.evaluate(match, host, {}) == 32
+
+
 def test_query_results_pickle_roundtrip():
     host = nx.DiGraph()
     host.add_node("x", name="Alice", age=30)

grandcypher/types.py

@@ -1,5 +1,15 @@
 """Typed operand classes for GrandCypher query processing."""
 
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Expression(Protocol):
+    """A query expression that can be evaluated for one graph match."""
+
+    def evaluate(self, match, host, return_edges, scope=None) -> Any:
+        ...
+
 
 class EntityRef(str):
     """Bare node/edge reference, e.g. A.
@@ -15,6 +25,12 @@ def __new__(cls, entity_name):
     def __getnewargs__(self):
         return (self.entity_name,)
 
+    def evaluate(self, match, host, return_edges, scope=None):
+        raise TypeError(
+            f"Cannot use bare entity '{self}' in a comparison. "
+            f"Use a property like '{self}.attribute' or ID({self})."
+        )
+
 
 class AttributeRef(str):
     """Node/edge attribute reference, e.g. A.age.
@@ -31,6 +47,20 @@ def __new__(cls, entity_name, attribute):
     def __getnewargs__(self):
         return (self.entity_name, self.attribute)
 
+    def evaluate(self, match, host, return_edges, scope=None):
+        if self.entity_name in match.node_mappings:
+            host_node_id = match.node_mappings[self.entity_name]
+            return host.nodes[host_node_id].get(self.attribute)
+        if self.entity_name in return_edges:
+            edge_mapping = return_edges[self.entity_name]
+            host_edges = match.mth.edge(*edge_mapping).edges
+            if len(host_edges) != 1:
+                raise TypeError("Cannot get edge attribute from multiple edges")
+            edge = host_edges[0]
+            edge_id = (edge.u, edge.v, edge.k) if host.is_multigraph() else (edge.u, edge.v)
+            return host.edges[edge_id].get(self.attribute)
+        raise IndexError(f"Entity {self} not in graph.")
+
 
 class IDRef(str):
     """Reference to ID(A) in WHERE clauses.
@@ -45,3 +75,8 @@ def __new__(cls, entity_name):
 
     def __getnewargs__(self):
         return (self.entity_name,)
+
+    def evaluate(self, match, host, return_edges, scope=None):
+        if self.entity_name in match.node_mappings:
+            return match.node_mappings[self.entity_name]
+        raise IndexError(f"Entity {self.entity_name} not in match.")