From 4c0254ec10fccea4a0e61254f7af28b7a9a84146 Mon Sep 17 00:00:00 2001
From: Huy Pham <42948410+hpham@users.noreply.github.com>
Date: Mon, 27 Jan 2020 16:27:21 -0800
Subject: [PATCH] reload custom node image when open xml, if the gui doesn't
 know about the custom image, use a default one

---
 daemon/core/gui/graph/graph.py |  7 ++++++-
 daemon/core/gui/images.py      | 24 +++++++++++++++++++---
 daemon/core/gui/nodeutils.py   | 37 +++++++++++++++++++++++++++-------
 3 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/daemon/core/gui/graph/graph.py b/daemon/core/gui/graph/graph.py
index 2b71e7cc..7ec01480 100644
--- a/daemon/core/gui/graph/graph.py
+++ b/daemon/core/gui/graph/graph.py
@@ -12,6 +12,7 @@ from core.gui.graph.enums import GraphMode, ScaleOption
 from core.gui.graph.node import CanvasNode
 from core.gui.graph.shape import Shape
 from core.gui.graph.shapeutils import ShapeType, is_draw_shape, is_marker
+from core.gui.images import ImageEnum, Images
 from core.gui.nodeutils import EdgeUtils, NodeUtils
 
 if TYPE_CHECKING:
@@ -20,6 +21,7 @@ if TYPE_CHECKING:
 
 ZOOM_IN = 1.1
 ZOOM_OUT = 0.9
+ICON_SIZE = 48
 
 
 class CanvasGraph(tk.Canvas):
@@ -217,7 +219,10 @@ class CanvasGraph(tk.Canvas):
             # peer to peer node is not drawn on the GUI
             if NodeUtils.is_ignore_node(core_node.type):
                 continue
-            image = NodeUtils.node_image(core_node)
+            image = NodeUtils.node_image(core_node, self.app.guiconfig)
+            # if the gui can't find node's image, default to the "edit-node" image
+            if not image:
+                image = Images.get(ImageEnum.EDITNODE, ICON_SIZE)
             x = core_node.position.x
             y = core_node.position.y
             node = CanvasNode(self.master, x, y, core_node, image)
diff --git a/daemon/core/gui/images.py b/daemon/core/gui/images.py
index 9b51b96c..6bf23e24 100644
--- a/daemon/core/gui/images.py
+++ b/daemon/core/gui/images.py
@@ -1,4 +1,5 @@
 from enum import Enum
+from tkinter import messagebox
 
 from PIL import Image, ImageTk
 
@@ -22,15 +23,32 @@ class Images:
             cls.images[image.stem] = str(image)
 
     @classmethod
-    def get(cls, image_enum: Enum, width: int, height: int = None):
+    def get(
+        cls, image_enum: Enum, width: int, height: int = None
+    ) -> ImageTk.PhotoImage:
         file_path = cls.images[image_enum.value]
         return cls.create(file_path, width, height)
 
     @classmethod
-    def get_custom(cls, name: str, width: int, height: int = None):
-        file_path = cls.images[name]
+    def get_with_image_file(
+        cls, stem: str, width: int, height: int = None
+    ) -> ImageTk.PhotoImage:
+        file_path = cls.images[stem]
         return cls.create(file_path, width, height)
 
+    @classmethod
+    def get_custom(
+        cls, name: str, width: int, height: int = None
+    ) -> ImageTk.PhotoImage:
+        try:
+            file_path = cls.images[name]
+            return cls.create(file_path, width, height)
+        except KeyError:
+            messagebox.showwarning(
+                "Missing image file",
+                f"{name}.png is missing at daemon/core/gui/data/icons, drop image file at daemon/core/gui/data/icons and restart the gui",
+            )
+
 
 class ImageEnum(Enum):
     SWITCH = "lanswitch"
diff --git a/daemon/core/gui/nodeutils.py b/daemon/core/gui/nodeutils.py
index 9f1b5a7f..c8ddb8fa 100644
--- a/daemon/core/gui/nodeutils.py
+++ b/daemon/core/gui/nodeutils.py
@@ -1,5 +1,5 @@
 import logging
-from typing import TYPE_CHECKING, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union
 
 from core.api.grpc.core_pb2 import NodeType
 from core.gui.images import ImageEnum, Images
@@ -91,14 +91,27 @@ class NodeUtils:
         return node_type in cls.RJ45_NODES
 
     @classmethod
-    def node_icon(cls, node_type: NodeType, model: str) -> "ImageTk.PhotoImage":
+    def node_icon(
+        cls,
+        node_type: NodeType,
+        model: str,
+        gui_config: Dict[str, List[Dict[str, str]]],
+    ) -> "ImageTk.PhotoImage":
         if model == "":
             model = None
-        return cls.NODE_ICONS[(node_type, model)]
+        try:
+            image = cls.NODE_ICONS[(node_type, model)]
+            return image
+        except KeyError:
+            image_stem = cls.get_image_file(gui_config, model)
+            if image_stem:
+                return Images.get_with_image_file(image_stem, ICON_SIZE)
 
     @classmethod
-    def node_image(cls, core_node: "core_pb2.Node") -> "ImageTk.PhotoImage":
-        image = cls.node_icon(core_node.type, core_node.model)
+    def node_image(
+        cls, core_node: "core_pb2.Node", gui_config: Dict[str, List[Dict[str, str]]]
+    ) -> "ImageTk.PhotoImage":
+        image = cls.node_icon(core_node.type, core_node.model, gui_config)
         if core_node.icon:
             try:
                 image = Images.create(core_node.icon, ICON_SIZE)
@@ -107,16 +120,26 @@ class NodeUtils:
         return image
 
     @classmethod
-    def is_custom(cls, model):
+    def is_custom(cls, model: str) -> bool:
         return model not in cls.NODE_MODELS
 
     @classmethod
-    def get_custom_node_services(cls, gui_config, name):
+    def get_custom_node_services(
+        cls, gui_config: Dict[str, List[Dict[str, str]]], name: str
+    ) -> List[str]:
         for m in gui_config["nodes"]:
             if m["name"] == name:
                 return m["services"]
         return []
 
+    @classmethod
+    def get_image_file(cls, gui_config, name: str) -> Union[str, None]:
+        if "nodes" in gui_config:
+            for m in gui_config["nodes"]:
+                if m["name"] == name:
+                    return m["image"]
+        return None
+
     @classmethod
     def setup(cls):
         nodes = [