From 431e5e2447b5632581edc5b6ee758aae779beea4 Mon Sep 17 00:00:00 2001
From: Jonas_Jones <me@jonasjones.dev>
Date: Tue, 4 Mar 2025 19:51:03 +0100
Subject: [PATCH] added enum and interface class type support

---
 puml_generator.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/puml_generator.py b/puml_generator.py
index f3c445a..2318c1e 100644
--- a/puml_generator.py
+++ b/puml_generator.py
@@ -129,12 +129,14 @@ class JavaClass:
     def __init__(self,
                  class_name,
                  class_package,
+                 class_type,
                  puml_content,
                  uml_relations,
                  class_fields,
                  class_methods):
         self.class_name = class_name
         self.class_package = class_package
+        self.class_type = class_type
         self.puml_content = puml_content
         self.uml_relations = uml_relations
         self.class_fields = class_fields
@@ -149,6 +151,9 @@ class JavaClass:
     def get_class_package(self):
         return self.class_package
 
+    def get_class_type(self):
+        return self.class_type
+
     def get_puml_content(self):
         return self.puml_content
 
@@ -173,6 +178,7 @@ def gen_puml_code_from_class(java_class:JavaClass, no_pkgs:bool=False) -> str:
     puml_code = ""
     class_name = java_class.get_class_name()
     class_package = java_class.get_class_package()
+    class_type = java_class.get_class_type()
     class_fields = java_class.get_class_fields()
     class_methods = java_class.get_class_methods()
     uml_relations = java_class.get_uml_relations()
@@ -180,7 +186,7 @@ def gen_puml_code_from_class(java_class:JavaClass, no_pkgs:bool=False) -> str:
     if not no_pkgs:
         puml_code += f"package {class_package + " {"}\n"
 
-    puml_code += f"class \"{class_name + "\" as " + class_package + "." + class_name + " {"}\n"
+    puml_code += f"{class_type} \"{class_name + "\" as " + class_package + "." + class_name + " {"}\n"
 
     for field, visibility in class_fields.items():
         puml_code += f"  {visibility} {field}\n"
@@ -231,6 +237,8 @@ def class_to_puml(filename:str, base_package_slug:str) -> JavaClass:
     class_name = filename.split("/")[-1].replace(".java", "")
     class_package = get_package_slug(filename)
 
+    class_type = "class"
+
     reached_class = False
     class_getter = False
     class_setter = False
@@ -286,6 +294,10 @@ def class_to_puml(filename:str, base_package_slug:str) -> JavaClass:
                         class_setter = True
                     else:
                         next_line_setter = True
+            elif " enum " in javaline:
+                class_type = "enum"
+            elif " interface " in javaline:
+                class_type = "interface"
             elif " class " in javaline:
                 reached_class = True
             elif is_line_field_declaration(javaline):
@@ -311,6 +323,7 @@ def class_to_puml(filename:str, base_package_slug:str) -> JavaClass:
     java_class = JavaClass(
         class_name,
         class_package,
+        class_type,
         puml_content,
         uml_relations,
         class_fields,