improve formatting and name conventions in query.py

Author: jkasiraj

src/sagemaker/lineage/query.py

@@ -99,7 +99,8 @@ def __str__(self):
 
         Format:
             {
-                'source_arn': 'string', 'destination_arn': 'string',
+                'source_arn': 'string',
+                'destination_arn': 'string',
                 'association_type': 'string'
             }
 
@@ -111,7 +112,8 @@ def __repr__(self):
 
         Format:
             {
-                'source_arn': 'string', 'destination_arn': 'string',
+                'source_arn': 'string',
+                'destination_arn': 'string',
                 'association_type': 'string'
             }
 
@@ -161,7 +163,8 @@ def __str__(self):
 
         Format:
             {
-                'arn': 'string', 'lineage_entity': 'string',
+                'arn': 'string',
+                'lineage_entity': 'string',
                 'lineage_source': 'string',
                 '_session': <sagemaker.session.Session object>
             }
@@ -174,7 +177,8 @@ def __repr__(self):
 
         Format:
             {
-                'arn': 'string', 'lineage_entity': 'string',
+                'arn': 'string',
+                'lineage_entity': 'string',
                 'lineage_source': 'string',
                 '_session': <sagemaker.session.Session object>
             }
@@ -238,7 +242,7 @@ def __init__(self, graph_styles):
             graph_styles: A dictionary that contains graph style for node and edges by their type.
                 Example: Display the nodes with different color by their lineage entity / different
                     shape by start arn.
-                        lineage_graph = {
+                        lineage_graph_styles = {
                             "TrialComponent": {
                                 "name": "Trial Component",
                                 "style": {"background-color": "#f6cf61"},
@@ -253,7 +257,7 @@ def __init__(self, graph_styles):
                                 "name": "StartArn",
                                 "style": {"shape": "star"},
                                 "isShape": "True",
-                                "symbol": "★",  # shape symbol for legend
+                                "symbol": "★", # shape symbol for legend
                             },
                         }
 
@@ -297,7 +301,7 @@ def __init__(self, graph_styles):
                 "solver": "hierarchicalRepulsion"
             }
         }
-            """
+        """
 
     def _import_visual_modules(self):
         """Import modules needed for visualization."""
@@ -358,7 +362,7 @@ def _add_legend(self, path):
         with open(path, "w", encoding="utf8") as file:
             file.write(html)
 
-    def render(self, elements, path="pyvisExample.html"):
+    def render(self, elements, path="lineage_graph_pyvis.html.html"):
         """Render graph for lineage query result.
 
         Args:
@@ -370,7 +374,7 @@ def render(self, elements, path="pyvisExample.html"):
                     elements["edges"] contains list of tuples, each tuple represents an edge
                         format: (edge source arn, edge destination arn, edge association type)
 
-            path(optional): The path/filemname of the rendered graph html file. (default path: "pyvisExample.html")
+            path(optional): The path/filename of the rendered graph html file. (default path: "lineage_graph_pyvis.html.html")
 
         Returns:
             display graph: The interactive visualization is presented as a static HTML file.
@@ -457,13 +461,21 @@ def __str__(self):
         Format:
         {
             'edges':[
-                {'source_arn': 'string', 'destination_arn': 'string', 'association_type': 'string'},
-                ...],
+                {
+                    'source_arn': 'string',
+                    'destination_arn': 'string',
+                    'association_type': 'string'
+                },
+            ],
 
             'vertices':[
-                {'arn': 'string', 'lineage_entity': 'string', 'lineage_source': 'string',
-                '_session': <sagemaker.session.Session object>},
-                ...],
+                {
+                    'arn': 'string',
+                    'lineage_entity': 'string',
+                    'lineage_source': 'string',
+                    '_session': <sagemaker.session.Session object>
+                },
+            ],
 
             'startarn':['string', ...]
         }
@@ -503,7 +515,7 @@ def _get_visualization_elements(self):
         elements = {"nodes": verts, "edges": edges}
         return elements
 
-    def visualize(self, path="pyvisExample.html"):
+    def visualize(self, path: Optional[str] = "lineage_graph_pyvis.html.html"):
         """Visualize lineage query result.
 
         Creates a PyvisVisualizer object to render network graph with Pyvis library.
@@ -512,12 +524,12 @@ def visualize(self, path="pyvisExample.html"):
         PyvisVisualizer for rendering graph.
 
         Args:
-            path(optional): The path/filemname of the rendered graph html file. (default path: "pyvisExample.html")
+            path(optional): The path/filemname of the rendered graph html file. (default path: "lineage_graph_pyvis.html.html")
 
         Returns:
             display graph: The interactive visualization is presented as a static HTML file.
         """
-        lineage_graph = {
+        lineage_graph_styles = {
             # nodes can have shape / color
             "TrialComponent": {
                 "name": "Trial Component",
@@ -547,7 +559,7 @@ def visualize(self, path="pyvisExample.html"):
             },
         }
 
-        pyvis_vis = PyvisVisualizer(lineage_graph)
+        pyvis_vis = PyvisVisualizer(lineage_graph_styles)
         elements = self._get_visualization_elements()
         return pyvis_vis.render(elements=elements, path=path)