feat(gfql ast): type edge by direction

Author: lmeyerov

graphistry/compute/ast.py

@@ -400,6 +400,22 @@ def __init__(self,
             edge_query=edge_query
         )
 
+    @classmethod
+    def from_json(cls, d: dict) -> 'ASTEdge':
+        out = ASTEdgeForward(
+            edge_match=maybe_filter_dict_from_json(d, 'edge_match'),
+            hops=d['hops'] if 'hops' in d else None,
+            to_fixed_point=d['to_fixed_point'] if 'to_fixed_point' in d else DEFAULT_FIXED_POINT,
+            source_node_match=maybe_filter_dict_from_json(d, 'source_node_match'),
+            destination_node_match=maybe_filter_dict_from_json(d, 'destination_node_match'),
+            source_node_query=d['source_node_query'] if 'source_node_query' in d else None,
+            destination_node_query=d['destination_node_query'] if 'destination_node_query' in d else None,
+            edge_query=d['edge_query'] if 'edge_query' in d else None,
+            name=d['name'] if 'name' in d else None
+        )
+        out.validate()
+        return out
+
 e_forward = ASTEdgeForward  # noqa: E305
 
 class ASTEdgeReverse(ASTEdge):
@@ -430,6 +446,22 @@ def __init__(self,
             edge_query=edge_query
         )
 
+    @classmethod
+    def from_json(cls, d: dict) -> 'ASTEdge':
+        out = ASTEdgeReverse(
+            edge_match=maybe_filter_dict_from_json(d, 'edge_match'),
+            hops=d['hops'] if 'hops' in d else None,
+            to_fixed_point=d['to_fixed_point'] if 'to_fixed_point' in d else DEFAULT_FIXED_POINT,
+            source_node_match=maybe_filter_dict_from_json(d, 'source_node_match'),
+            destination_node_match=maybe_filter_dict_from_json(d, 'destination_node_match'),
+            source_node_query=d['source_node_query'] if 'source_node_query' in d else None,
+            destination_node_query=d['destination_node_query'] if 'destination_node_query' in d else None,
+            edge_query=d['edge_query'] if 'edge_query' in d else None,
+            name=d['name'] if 'name' in d else None
+        )
+        out.validate()
+        return out
+
 e_reverse = ASTEdgeReverse  # noqa: E305
 
 class ASTEdgeUndirected(ASTEdge):
@@ -460,6 +492,22 @@ def __init__(self,
             edge_query=edge_query
         )
 
+    @classmethod
+    def from_json(cls, d: dict) -> 'ASTEdge':
+        out = ASTEdgeUndirected(
+            edge_match=maybe_filter_dict_from_json(d, 'edge_match'),
+            hops=d['hops'] if 'hops' in d else None,
+            to_fixed_point=d['to_fixed_point'] if 'to_fixed_point' in d else DEFAULT_FIXED_POINT,
+            source_node_match=maybe_filter_dict_from_json(d, 'source_node_match'),
+            destination_node_match=maybe_filter_dict_from_json(d, 'destination_node_match'),
+            source_node_query=d['source_node_query'] if 'source_node_query' in d else None,
+            destination_node_query=d['destination_node_query'] if 'destination_node_query' in d else None,
+            edge_query=d['edge_query'] if 'edge_query' in d else None,
+            name=d['name'] if 'name' in d else None
+        )
+        out.validate()
+        return out
+
 e_undirected = ASTEdgeUndirected  # noqa: E305
 e = ASTEdgeUndirected  # noqa: E305
 
@@ -472,7 +520,17 @@ def from_json(o: JSONVal) -> Union[ASTNode, ASTEdge]:
     if o['type'] == 'Node':
         out = ASTNode.from_json(o)
     elif o['type'] == 'Edge':
-        out = ASTEdge.from_json(o)
+        if 'direction' in o:
+            if o['direction'] == 'forward':
+                out = ASTEdgeForward.from_json(o)
+            elif o['direction'] == 'reverse':
+                out = ASTEdgeReverse.from_json(o)
+            elif o['direction'] == 'undirected':
+                out = ASTEdgeUndirected.from_json(o)
+            else:
+                raise ValueError(f'Edge has unknown direction {o["direction"]}')
+        else:
+            raise ValueError('Edge missing direction')
     else:
         raise ValueError(f'Unknown type {o["type"]}')
     return out

graphistry/tests/compute/test_chain.py

@@ -1,11 +1,11 @@
 import os
 import pandas as pd
-from graphistry.compute.predicates.is_in import is_in
-from graphistry.compute.predicates.numeric import gt
 import pytest
 
-from graphistry.compute.ast import ASTNode, ASTEdge, n, e, e_undirected, e_forward
+from graphistry.compute.ast import ASTEdgeUndirected, ASTNode, ASTEdge, n, e, e_undirected, e_forward
 from graphistry.compute.chain import Chain
+from graphistry.compute.predicates.is_in import IsIn, is_in
+from graphistry.compute.predicates.numeric import gt
 from graphistry.tests.test_compute import CGFull
 
 
@@ -298,6 +298,25 @@ def test_chain_serialization_pred():
     o2 = d.to_json()
     assert o == o2
 
+def test_chain_serialize_pred_is_in():
+
+    #from graphistry.compute.chain import Chain
+    #from graphistry import e_undirected, is_in
+    o = Chain([
+        e_undirected(
+            hops=1,
+            edge_match={"source": is_in(options=[
+                "Oakville Square",
+                "Maplewood Square"
+            ])})
+    ]).to_json()
+    d = Chain.from_json(o)
+    assert isinstance(d.chain[0], ASTEdgeUndirected), f'got: {type(d.chain[0])}'
+    assert d.chain[0].direction == 'undirected'
+    assert d.chain[0].hops == 1
+    assert isinstance(d.chain[0].edge_match['source'], IsIn)
+    assert d.chain[0].edge_match['source'].options == ['Oakville Square', 'Maplewood Square']
+
 def test_chain_simple_cudf_pd():
     nodes_df = pd.DataFrame({'id': [0, 1, 2], 'label': ['a', 'b', 'c']})
     edges_df = pd.DataFrame({'src': [0, 1, 2], 'dst': [1, 2, 0]})