change: change visualization to using pyvis library

Author: Yi-Ting Lee

src/sagemaker/lineage/query.py

@@ -201,218 +201,47 @@ def _artifact_to_lineage_object(self):
         return Artifact.load(artifact_arn=self.arn, sagemaker_session=self._session)
 
 
-class DashVisualizer(object):
-    """Create object used for visualizing graph using Dash library."""
+class PyvisVisualizer(object):
+    """Create object used for visualizing graph using Pyvis library."""
 
     def __init__(self, graph_styles):
-        """Init for DashVisualizer."""
+        """Init for PyvisVisualizer."""
         # import visualization packages
         (
-            self.cyto,
-            self.JupyterDash,
-            self.html,
-            self.Input,
-            self.Output,
+            self.pyvis,
+            self.Network,
+            self.Options,
         ) = self._import_visual_modules()
 
         self.graph_styles = graph_styles
 
     def _import_visual_modules(self):
         """Import modules needed for visualization."""
         try:
-            import dash_cytoscape as cyto
+            import pyvis
         except ImportError as e:
             print(e)
-            print("Try: pip install dash-cytoscape")
+            print("Try: pip install pyvis")
             raise
 
         try:
-            from jupyter_dash import JupyterDash
+            from pyvis.network import Network
         except ImportError as e:
             print(e)
-            print("Try: pip install jupyter-dash")
+            print("Try: pip install pyvis")
             raise
 
         try:
-            from dash import html
+            from pyvis.options import Options
         except ImportError as e:
             print(e)
-            print("Try: pip install dash")
+            print("Try: pip install pyvis")
             raise
 
-        try:
-            from dash.dependencies import Input, Output
-        except ImportError as e:
-            print(e)
-            print("Try: pip install dash")
-            raise
-
-        return cyto, JupyterDash, html, Input, Output
-
-    def _create_legend_component(self, style):
-        """Create legend component div."""
-        text = style["name"]
-        symbol = ""
-        color = "#ffffff"
-        if style["isShape"] == "False":
-            color = style["style"]["background-color"]
-        else:
-            symbol = style["symbol"]
-        return self.html.Div(
-            [
-                self.html.Div(
-                    symbol,
-                    style={
-                        "background-color": color,
-                        "width": "1.5vw",
-                        "height": "1.5vw",
-                        "display": "inline-block",
-                        "font-size": "1.5vw",
-                    },
-                ),
-                self.html.Div(
-                    style={
-                        "width": "0.5vw",
-                        "height": "1.5vw",
-                        "display": "inline-block",
-                    }
-                ),
-                self.html.Div(
-                    text,
-                    style={"display": "inline-block", "font-size": "1.5vw"},
-                ),
-            ]
-        )
-
-    def _create_entity_selector(self, entity_name, style):
-        """Create selector for each lineage entity."""
-        return {"selector": "." + entity_name, "style": style["style"]}
-
-    def _get_app(self, elements):
-        """Create JupyterDash app for interactivity on Jupyter notebook."""
-        app = self.JupyterDash(__name__)
-        self.cyto.load_extra_layouts()
-
-        app.layout = self.html.Div(
-            [
-                # graph section
-                self.cyto.Cytoscape(
-                    id="cytoscape-graph",
-                    elements=elements,
-                    style={
-                        "width": "84%",
-                        "height": "350px",
-                        "display": "inline-block",
-                        "border-width": "1vw",
-                        "border-color": "#232f3e",
-                    },
-                    layout={"name": "klay"},
-                    stylesheet=[
-                        {
-                            "selector": "node",
-                            "style": {
-                                "label": "data(label)",
-                                "font-size": "3.5vw",
-                                "height": "10vw",
-                                "width": "10vw",
-                                "border-width": "0.8",
-                                "border-opacity": "0",
-                                "border-color": "#232f3e",
-                                "font-family": "verdana",
-                            },
-                        },
-                        {
-                            "selector": "edge",
-                            "style": {
-                                "label": "data(label)",
-                                "color": "gray",
-                                "text-halign": "left",
-                                "text-margin-y": "2.5",
-                                "font-size": "3",
-                                "width": "1",
-                                "curve-style": "bezier",
-                                "control-point-step-size": "15",
-                                "target-arrow-color": "gray",
-                                "target-arrow-shape": "triangle",
-                                "line-color": "gray",
-                                "arrow-scale": "0.5",
-                                "font-family": "verdana",
-                            },
-                        },
-                        {"selector": ".select", "style": {"border-opacity": "0.7"}},
-                    ]
-                    + [self._create_entity_selector(k, v) for k, v in self.graph_styles.items()],
-                    responsive=True,
-                ),
-                self.html.Div(
-                    style={
-                        "width": "0.5%",
-                        "display": "inline-block",
-                        "font-size": "1vw",
-                        "font-family": "verdana",
-                        "vertical-align": "top",
-                    },
-                ),
-                # legend section
-                self.html.Div(
-                    [self._create_legend_component(v) for k, v in self.graph_styles.items()],
-                    style={
-                        "display": "inline-block",
-                        "font-size": "1vw",
-                        "font-family": "verdana",
-                        "vertical-align": "top",
-                    },
-                ),
-            ]
-        )
-
-        @app.callback(
-            self.Output("cytoscape-graph", "elements"),
-            self.Input("cytoscape-graph", "tapNodeData"),
-            self.Input("cytoscape-graph", "elements"),
-        )
-        def selectNode(tapData, elements):
-            for n in elements:
-                if tapData is not None and n["data"]["id"] == tapData["id"]:
-                    # if is tapped node, add "select" class to node
-                    n["classes"] += " select"
-                elif "classes" in n:
-                    # remove "select" class in "classes" if node not selected
-                    n["classes"] = n["classes"].replace("select", "")
-
-            return elements
-
-        return app
-
-    def render(self, elements, mode):
-        """Render graph for lineage query result."""
-        app = self._get_app(elements)
-
-        return app.run_server(mode=mode)
-
-class PyvisVisualizer(object):
-    """Create object used for visualizing graph using Pyvis library."""
-
-    def __init__(self, graph_styles):
-        """Init for PyvisVisualizer."""
-        # import visualization packages
-        (
-            self.pyvis,
-            self.Network,
-            self.Options,
-        ) = self._import_visual_modules()
-
-        self.graph_styles = graph_styles
-
-    def _import_visual_modules(self):
-        import pyvis
-        from pyvis.network import Network
-        from pyvis.options import Options
-        # No module named 'pyvis'
-
         return pyvis, Network, Options
 
     def _get_options(self):
+        """Get pyvis graph options."""
         options = """
             var options = {
             "configure":{
@@ -445,24 +274,29 @@ def _get_options(self):
         return options
 
     def _node_color(self, n):
+        """Return node color by background-color specified in graph styles."""
         return self.graph_styles[n[2]]["style"]["background-color"]
 
-    def render(self, elements):
-        net = self.Network(height='500px', width='100%', notebook = True, directed = True)
+    def render(self, elements, path="pyvisExample.html"):
+        """Render graph for lineage query result."""
+        net = self.Network(height="500px", width="100%", notebook=True, directed=True)
         options = self._get_options()
-        net.set_options(options)  
+        net.set_options(options)
 
+        # add nodes to graph
         for n in elements["nodes"]:
-            if(n[3]==True): # startarn
-                net.add_node(n[0], label=n[1], title=n[1], color=self._node_color(n), shape="star")
+            if n[3]:  # startarn
+                net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n), shape="star")
             else:
-                net.add_node(n[0], label=n[1], title=n[1], color=self._node_color(n))
+                net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n))
 
+        # add edges to graph
         for e in elements["edges"]:
             print(e)
-            net.add_edge(e[0], e[1], title=e[2])  
+            net.add_edge(e[0], e[1], title=e[2])
+
+        return net.show(path)
 
-        return net.show('pyvisExample.html')    
 
 class LineageQueryResult(object):
     """A wrapper around the results of a lineage query."""
@@ -522,18 +356,6 @@ def __str__(self):
         result_dict = vars(self)
         return str({k: [str(val) for val in v] for k, v in result_dict.items()})
 
-    def _covert_vertices_to_tuples(self):
-        """Convert vertices to tuple format for visualizer."""
-        verts = []
-        # get vertex info in the form of (id, label, class)
-        for vert in self.vertices:
-            if vert.arn in self.startarn:
-                # add "startarn" class to node if arn is a startarn
-                verts.append((vert.arn, vert.lineage_source, vert.lineage_entity + " startarn"))
-            else:
-                verts.append((vert.arn, vert.lineage_source, vert.lineage_entity))
-        return verts
-
     def _covert_edges_to_tuples(self):
         """Convert edges to tuple format for visualizer."""
         edges = []
@@ -542,7 +364,7 @@ def _covert_edges_to_tuples(self):
             edges.append((edge.source_arn, edge.destination_arn, edge.association_type))
         return edges
 
-    def _pyvis_covert_vertices_to_tuples(self):
+    def _covert_vertices_to_tuples(self):
         """Convert vertices to tuple format for visualizer."""
         verts = []
         # get vertex info in the form of (id, label, class)
@@ -555,38 +377,15 @@ def _pyvis_covert_vertices_to_tuples(self):
         return verts
 
     def _get_visualization_elements(self):
-        """Get elements for visualization."""
-        # get vertices and edges info for graph
+        """Get elements(nodes+edges) for visualization."""
         verts = self._covert_vertices_to_tuples()
         edges = self._covert_edges_to_tuples()
 
-        nodes = [
-            {"data": {"id": id, "label": label}, "classes": classes} for id, label, classes in verts
-        ]
-
-        edges = [
-            {"data": {"source": source, "target": target, "label": label}}
-            for source, target, label in edges
-        ]
-
-        elements = nodes + edges
-
-        return elements
-
-    def _get_pyvis_visualization_elements(self):
-        verts = self._pyvis_covert_vertices_to_tuples()
-        edges = self._covert_edges_to_tuples()
-
-        elements = {
-            "nodes": verts,
-            "edges": edges
-        }
+        elements = {"nodes": verts, "edges": edges}
         return elements
 
     def visualize(self):
         """Visualize lineage query result."""
-        elements = self._get_visualization_elements()
-
         lineage_graph = {
             # nodes can have shape / color
             "TrialComponent": {
@@ -617,15 +416,8 @@ def visualize(self):
             },
         }
 
-        # initialize DashVisualizer instance to render graph & interactive components
-        # dash_vis = DashVisualizer(lineage_graph)
-
-        # dash_server = dash_vis.render(elements=elements, mode="inline")
-
-        # return dash_server
-
         pyvis_vis = PyvisVisualizer(lineage_graph)
-        elements = self._get_pyvis_visualization_elements()
+        elements = self._get_visualization_elements()
         return pyvis_vis.render(elements=elements)