From 0fd81593c606d1a2ccef35d7d298241f185f542d Mon Sep 17 00:00:00 2001
From: Dov Benyomin Sohacheski <b@kloud.email>
Date: Tue, 8 Nov 2022 13:19:35 +0200
Subject: [PATCH] load icon image based on class name

---
 diagrams/__init__.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index b46fa468..627b40ee 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -1,5 +1,6 @@
 import contextvars
 import os
+import re
 import uuid
 from pathlib import Path
 from typing import List, Union, Dict
@@ -279,6 +280,7 @@ class Node:
 
     _icon_dir = None
     _icon = None
+    _icon_ext = ".png"
 
     _height = 1.9
 
@@ -394,7 +396,15 @@ class Node:
             else:
                 o.connect(self, Edge(self, reverse=True))
         return self
-
+    
+    def _get_icon(self):
+        if self._icon:
+            return self._icon
+        
+        icon = re.sub(r'(?<!^)(?=[A-Z])', '-', self.__class__.__name__).lower()
+        
+        return f'{icon}.{self._icon_ext}'
+    
     @property
     def nodeid(self):
         return self._id
@@ -421,7 +431,7 @@ class Node:
 
     def _load_icon(self):
         basedir = Path(os.path.abspath(os.path.dirname(__file__)))
-        return os.path.join(basedir.parent, self._icon_dir, self._icon)
+        return os.path.join(basedir.parent, self._icon_dir, self._get_icon())
 
 
 class Edge: