generate legend divs programmatically

Author: Yi-Ting Lee

src/sagemaker/lineage/query.py

@@ -314,6 +314,25 @@ def _node_color(self, entity):
         """Return node color by background-color specified in graph styles."""
         return self.graph_styles[entity]["style"]["background-color"]
 
+    def _get_legend_line(self, component_name):
+        """Generate lengend div line for each graph component in graph_styles."""
+        if self.graph_styles[component_name]["isShape"] == "False":
+            return '<div><div style="background-color: {color}; width: 1.6vw; height: 1.6vw;\
+                display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>\
+                <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>\
+                <div style="display: inline-block; font-size: 1.5vw;">{name}</div></div>'.format(
+                color=self.graph_styles[component_name]["style"]["background-color"],
+                name=self.graph_styles[component_name]["name"],
+            )
+        else:
+            return '<div style="background-color: #ffffff; width: 1.6vw; height: 1.6vw;\
+                display: inline-block; font-size: 0.9vw; vertical-align: -0.2em;">{shape}</div>\
+                <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>\
+                <div style="display: inline-block; font-size: 1.5vw;">{name}</div></div>'.format(
+                shape=self.graph_styles[component_name]["style"]["shape"],
+                name=self.graph_styles[component_name]["name"],
+            )
+
     def _add_legend(self, path):
         """Embed legend to html file generated by pyvis."""
         f = open(path, "r")
@@ -322,38 +341,13 @@ def _add_legend(self, path):
         legend = """
             <div style="display: inline-block; font-size: 1vw; font-family: verdana;
                 vertical-align: top; padding: 1vw;">
-                <div>
-                    <div style="background-color: rgb(246, 207, 97); width: 1.6vw; height: 1.6vw;
-                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
-                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
-                    <div style="display: inline-block; font-size: 1.5vw;">Trial Component</div>
-                </div>
-                <div>
-                    <div style="background-color: rgb(255, 153, 0); width: 1.6vw; height: 1.6vw;
-                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
-                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
-                    <div style="display: inline-block; font-size: 1.5vw;">Context</div>
-                </div>
-                <div>
-                    <div style="background-color: rgb(136, 195, 150); width: 1.6vw; height: 1.6vw;
-                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
-                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
-                    <div style="display: inline-block; font-size: 1.5vw;">Action</div>
-                </div>
-                <div>
-                    <div style="background-color: rgb(20, 110, 180); width: 1.6vw; height: 1.6vw;
-                        display: inline-block; font-size: 1.5vw; vertical-align: -0.2em;"></div>
-                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
-                    <div style="display: inline-block; font-size: 1.5vw;">Artifact</div>
-                </div>
-                <div>
-                    <div style="background-color: rgb(255, 255, 255); width: 1.6vw; height: 1.6vw;
-                        display: inline-block; font-size: 0.9vw;">star</div>
-                    <div style="width: 0.3vw; height: 1.5vw; display: inline-block;"></div>
-                    <div style="display: inline-block; font-size: 1.5vw;">StartArn</div>
-                </div>
-            </div>
         """
+        # iterate through graph styles to get legend
+        for component in self.graph_styles.keys():
+            legend += self._get_legend_line(component_name=component)
+
+        legend += "</div>"
+
         legend_div = self.BeautifulSoup(legend, "html.parser")
 
         content.div.insert_after(legend_div)