Merge branch 'master' of https://github.com/WolfgangFahl/gremlin-python-tutorial

Author: WolfgangFahl

gremlin/draw.py

@@ -3,44 +3,148 @@
 
 @author: jv
 '''
+from collections.abc import Iterable
+from typing import Any, List, Union
+
 import graphviz
+
+from gremlin_python.process.anonymous_traversal import GraphTraversalSource
+from gremlin_python.process.graph_traversal import GraphTraversal
+from gremlin_python.process.traversal import T
+from gremlin_python.structure.graph import Vertex, Edge, Path
+
+from aenum import Enum
+
+
 class GremlinDraw:
-    @classmethod
-    def show(cls,g,title:str="Gremlin", v_limit:int=10, e_limit:int=10):
-        G = graphviz.Digraph(title,format="pdf")
+    @staticmethod
+    def __draw_vertex(digraph: graphviz.Digraph, g: GraphTraversalSource, vertex: Vertex) -> graphviz.Digraph:
+        """
+        draw a single given vertex
+        """
+        # developer note: see https://github.com/apache/tinkerpop/blob/master/gremlin-python/src/main/python/gremlin_python/structure/graph.py#LL58C23-L58C23
+        # when gremlin-python 3.7.0 is released, the following code can be improved (get the properties using vertex.properties)
+        # then, g can also be removed as a parameter
+        
+        # get the properties of the vertex
+        kvp_list = list(next(g.V(vertex).element_map()).items())
+        # non-proerty items are of type aenum
+        properties = [item for item in kvp_list if not isinstance(item[0], Enum)]
+        assert len(properties) == len(kvp_list) - 2 # ID and label are not properties
+    
+        properties_label = "\n".join(f"{key}: {value}" for key, value in properties)
+
+        # draw the vertex
+        digraph.node(
+            name=str(vertex.id),
+            label=f"{str(vertex.id)}\n{vertex.label}\n{'─' * 5}\n{properties_label}",
+            fillcolor = "#ADE1FE",
+            style = "filled",
+            fontname = "arial"
+        )
+
+        return digraph
+
+    
+    @staticmethod
+    def __draw_edge(digraph: graphviz.Digraph, g: GraphTraversalSource, edge: Edge) -> graphviz.Digraph:
+        """
+        draw a single given edge
+        """
+        # developer note: see https://github.com/apache/tinkerpop/blob/master/gremlin-python/src/main/python/gremlin_python/structure/graph.py#L66
+        # when gremlin-python 3.7.0 is released, the following code can be improved (get the properties using edge.properties)
+        # then, g can also be removed as a parameter
+
+        # get the properties of the edge
+        #kvp_list = list(next(g.E(edge).element_map()).items())
+        # Workaround, because the above line does not work due to inconsistencies / bugs in the gremlin-python library
+        kvp_list = [edge_element_map for edge_element_map in g.E().element_map().to_list() if edge_element_map[T.id] == edge.id][0].items()
+        # non-proerty items are of type aenum
+        properties = [item for item in kvp_list if not isinstance(item[0], Enum)]
+        assert len(properties) == len(kvp_list) - 4 # ID, label, in, and out are not properties
+        
+        properties_label = "\n".join(f"{key}: {value}" for key, value in properties)
+        
+        # get the image of the edge by id
+        in_vertex_id = edge.inV.id
+        out_vertex_id = edge.outV.id
+
+        # draw the edge
+        digraph.edge(
+            tail_name = str(in_vertex_id),
+            head_name = str(out_vertex_id),
+            label = f"{str(edge.id)}\n{edge.label}\n{'─' * 5}\n{properties_label}",
+            style = "setlinewidth(3)",
+            fontname = "arial"
+        )
+
+        return digraph
+
+
+    @staticmethod
+    def show(g: GraphTraversalSource, title:str="Gremlin", v_limit:int=10, e_limit:int=10) -> graphviz.Digraph:
+        """
+        draw the given graph
+        """
+        G: graphviz.Digraph = graphviz.Digraph(title, format="pdf")
 
         # draw vertices
-        vlist = g.V().toList()
+        vlist = g.V().to_list()
         vlist = vlist[:v_limit]
 
         for v in vlist:
-            kvp_list = list(g.V(v).elementMap().next().items())
-            kvp_list.pop(0) # pop ID
-            kvp_list.pop(0) # pop label
-            vertex_properties = ""
-            for key,val in kvp_list:
-                vertex_properties += "\n " + str(key) + ": " + str(val);
-            
-            G.node(name=str(v.id), label=str(v.id) + "\n" + v.label + "\n" + '─' * 5 + vertex_properties, fillcolor = "#ADE1FE", style = "filled", fontname = "arial")
-            
+            G = GremlinDraw.__draw_vertex(G, g, v)
+
         #draw edges
-        elist = g.E().elementMap().toList()
+        elist = g.E().to_list()
         elist = elist[:e_limit]
 
         for e in elist:
-            kvp_list = list(e.items())
-            kvp_list.pop(0) # pop ID
-            kvp_list.pop(0) # pop label
-            kvp_list.pop(0) # pop in
-            kvp_list.pop(0) # pop out
-            edge_properties = ""
-            for key,val in kvp_list:
-                edge_properties += "\n " + str(key) + ": " + str(val);
-                
-                # list(list(e.values())[2].values())[0]
-                # str(g.E(e).inV().next().id)
-                # str(g.E(e).outV().next().id)
-                # list(e.values())[0]
-            
-            G.edge(tail_name = str(list(list(e.values())[3].values())[0]), head_name = str(list(list(e.values())[2].values())[0]), style = "setlinewidth(3)", label = str(list(e.values())[0]) + "\n" + str(list(e.values())[1]) +  "\n" + '─' * 5 + edge_properties, fontname = "arial")
-        return G
+            G = GremlinDraw.__draw_edge(G, g, e)
+        
+        return G
+
+    @staticmethod
+    def show_graph_traversal(g: GraphTraversalSource, gt: Union[GraphTraversal, Any], title: str="Gremlin") -> graphviz.Digraph:
+        """
+        draw the given graph traversal
+        """
+        # developer note: when moving the minium supported version up to 3.10, the following code can be greatly improved by using match statements
+
+        G: graphviz.Digraph = graphviz.Digraph(title, format="pdf")
+
+        worklist: List[Any] = gt.to_list() if isinstance(gt, GraphTraversal) else list(gt) if isinstance(gt, Iterable) else [gt]
+
+        while len(worklist) > 0:
+            # move any vertices to the front of the worklist (draw them first)
+            worklist = [item for item in worklist if not isinstance(item, Vertex)] + [item for item in worklist if isinstance(item, Vertex)]
+
+            result = worklist.pop(0)
+
+            if isinstance(result, Vertex):
+                G = GremlinDraw.__draw_vertex(G, g, result)
+            elif isinstance(result, Edge):
+                G = GremlinDraw.__draw_edge(G, g, result)
+            elif isinstance(result, Path):
+                for item in result.objects:
+                    worklist.append(item)
+            elif isinstance(result, dict):
+                if T.id in result:
+                    # check if the id is a vertex or an edge
+                    if g.V(result[T.id]).hasNext():
+                        G = GremlinDraw.__draw_vertex(G, g, next(g.V(result[T.id])))
+                    elif g.E(result[T.id]).hasNext():
+                        G = GremlinDraw.__draw_edge(G, g, g.E(result[T.id]).next())
+                    else:
+                        #raise Exception("id not found")
+                        pass # silent skip
+                else:
+                    #raise Exception("id not found")
+                    pass # silent skip
+            else:
+                #raise Exception(f"unknown type: {type(result)}")
+                pass # silent skip
+
+        return G
+
+

gremlin/examples.py

@@ -1,6 +1,9 @@
+from os.path import abspath, dirname
 import urllib.request
 import os
 from pathlib import Path
+
+from gremlin_python.process.anonymous_traversal import GraphTraversalSource
 from gremlin.remote import RemoteTraversal
 from dataclasses import dataclass
 
@@ -39,8 +42,8 @@ def remote(self,file_name:str):
         return path
 
 
-    @classmethod
-    def docker(cls)->"Volume":
+    @staticmethod
+    def docker()->"Volume":
         """
         get the default docker volume mapping
         
@@ -54,30 +57,46 @@ def docker(cls)->"Volume":
         volume=Volume(local_path=local_path,remote_path=remote_path)
         return volume
 
+    @staticmethod
+    def local() -> "Volume":
+        """
+        get the default local volume mapping
+
+        Returns:
+            Volume: the local_path/remote_path mapping
+        """
+        home = str(Path.home())
+        local_path=f"{home}/.gremlin-examples"
+        os.makedirs(local_path, exist_ok=True)
+        remote_path=str(abspath(f"{dirname(abspath(__file__))}/data"))
+        volume=Volume(local_path=local_path,remote_path=remote_path)
+        return volume
+    
 @dataclass
 class Example:
     name:str
     url:str
 
-    def __post_init__(self):
-        """
-        """
-    
-    def load(self,g,volume:Volume,force:bool=False,debug:bool=False):
+    def load(
+        self,
+        g: GraphTraversalSource,
+        volume: Volume,
+        force: bool = False,
+        debug: bool = False,
+    ) -> None:
         """
         download graph from remote_path to local_path depending on force flag
         and load graph into g
         
         Args:
-            g(GraphTraversal): the target 
+            g(GraphTraversalSource): the target graph (inout)
             volume:Volume
             force(bool): if True download even if local copy already exists
             debug(bool): if True show debugging information
         """
         self.download(volume.local_path, force=force,debug=debug)
         graph_xml=f"{volume.remote_path}/{self.name}.xml"
         RemoteTraversal.load(g, graph_xml)
-        pass
 
     def download(self,path,force:bool=False,debug:bool=False)->str:
         """
@@ -113,7 +132,7 @@ class Examples:
     Examples 
     """
 
-    def __init__(self,volume,debug:bool=False):
+    def __init__(self,volume:Volume,debug:bool=False):
         """
         Constructor
         
@@ -129,19 +148,23 @@ def __init__(self,volume,debug:bool=False):
             Example(name="tinkerpop-modern",url="https://raw.githubusercontent.com/apache/tinkerpop/master/data/tinkerpop-modern.xml"),
             Example(name="grateful-dead",url="https://raw.githubusercontent.com/apache/tinkerpop/master/data/grateful-dead.xml"),
             Example(name="air-routes-small",url="https://raw.githubusercontent.com/krlawrence/graph/master/sample-data/air-routes-small.graphml"),
-            Example(name="air-routes-latest",url="https://raw.githubusercontent.com/krlawrence/graph/master/sample-data/air-routes-latest.graphml")        ]:
+            Example(name="air-routes-latest",url="https://raw.githubusercontent.com/krlawrence/graph/master/sample-data/air-routes-latest.graphml")
+        ]:
             self.examples_by_name[example.name]=example
 
-    def load_by_name(self,g,name:str):
+    def load_by_name(self, g: GraphTraversalSource, name: str) -> None:
         """
         load an example by name to the given graph
         
         Args:
-            g: the target graph
+            g(GraphTraversalSource): the target graph (inout)
             name(str): the name of the example
+
+        Raises:
+            Exception: if the example does not exist
         """
         if name in self.examples_by_name:
             example=self.examples_by_name[name]
             example.load(g,self.volume,debug=self.debug)
         else:
-            raise Exception(f"invalid example {name}")
+            raise Exception(f"invalid example {name}")