From 80d6a97e1f08c05bbd3dc54c738c37499f2e4b1c Mon Sep 17 00:00:00 2001
From: Oliver Walters <oliver.henry.walters@gmail.com>
Date: Thu, 21 Nov 2024 12:10:38 +0000
Subject: [PATCH] Refactor get_custom_classes

- Add StatusCode.custom_values method
---
 .../InvenTree/generic/states/custom.py        | 43 ++++++-------------
 .../InvenTree/generic/states/states.py        | 15 ++++++-
 2 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/src/backend/InvenTree/generic/states/custom.py b/src/backend/InvenTree/generic/states/custom.py
index 8fc8948c28..66c94202b8 100644
--- a/src/backend/InvenTree/generic/states/custom.py
+++ b/src/backend/InvenTree/generic/states/custom.py
@@ -53,37 +53,22 @@ def get_custom_classes(
     if not include_custom:
         return discovered_classes
 
-    # Gather DB settings
-    from common.models import InvenTreeCustomUserStateModel
-
-    custom_db_states = {}
-    custom_db_mdls = {}
-    for item in list(InvenTreeCustomUserStateModel.objects.all()):
-        if not custom_db_states.get(item.reference_status):
-            custom_db_states[item.reference_status] = []
-        custom_db_states[item.reference_status].append(item)
-        custom_db_mdls[item.model.app_label] = item.reference_status
-    custom_db_mdls_keys = custom_db_mdls.keys()
-
     states = {}
+
     for cls in discovered_classes:
-        tag = cls.tag()
-        states[tag] = cls
-        if custom_db_mdls and tag in custom_db_mdls_keys:
-            data = [(str(m.name), (m.value, m.label, m.color)) for m in states[tag]]
-            data_keys = [i[0] for i in data]
+        name = cls.__name__
+        states[name] = cls
 
-            # Extent with non present tags
-            for entry in custom_db_states[custom_db_mdls[tag]]:
-                ref_name = str(entry.name.upper().replace(' ', ''))
-                if ref_name not in data_keys:
-                    data += [
-                        (
-                            str(entry.name.upper().replace(' ', '')),
-                            (entry.key, entry.label, entry.color),
-                        )
-                    ]
+        data = [(str(m.name), (m.value, m.label, m.color)) for m in cls]
+
+        labels = [str(item[0]) for item in data]
+
+        for item in cls.custom_values():
+            label = str(item.name)
+            if label not in labels and label not in cls.labels():
+                data += [(str(item.name), (item.key, item.label, item.color))]
+
+        # Re-assemble the enum
+        states[name] = base_class(name, data)
 
-            # Re-assemble the enum
-            states[tag] = base_class(f'{tag.capitalize()}Status', data)
     return states.values()
diff --git a/src/backend/InvenTree/generic/states/states.py b/src/backend/InvenTree/generic/states/states.py
index 966ed76210..7f447b7f67 100644
--- a/src/backend/InvenTree/generic/states/states.py
+++ b/src/backend/InvenTree/generic/states/states.py
@@ -103,9 +103,22 @@ class StatusCode(BaseEnum):
         return isinstance(value.value, int)
 
     @classmethod
-    def values(cls, key=None):
+    def custom_values(cls):
+        """Return all user-defined custom values for this status class."""
+        from common.models import InvenTreeCustomUserStateModel
+
+        try:
+            return InvenTreeCustomUserStateModel.objects.filter(
+                reference_status=cls.__name__
+            )
+        except Exception:
+            return []
+
+    @classmethod
+    def values(cls, key=None, custom=True):
         """Return a dict representation containing all required information."""
         elements = [itm for itm in cls if cls._is_element(itm.name)]
+
         if key is None:
             return elements