添加 choices 获取 label方法

apps/assets/const.py

@@ -1,6 +1,6 @@
 from django.db import models
 from django.utils.translation import gettext_lazy as _
-from common.db.models import IncludesTextChoicesMeta
+from common.db.models import IncludesTextChoicesMeta, ChoicesMixin
 from common.tree import TreeNode
 
 
@@ -24,7 +24,7 @@ class PlatformMixin:
         }
 
 
-class Category(PlatformMixin, models.TextChoices):
+class Category(PlatformMixin, ChoicesMixin, models.TextChoices):
     HOST = 'host', _('Host')
     NETWORK = 'network', _("NetworkDevice")
     DATABASE = 'database', _("Database")
@@ -60,7 +60,7 @@ class Category(PlatformMixin, models.TextChoices):
         }
 
 
-class HostTypes(PlatformMixin, models.TextChoices):
+class HostTypes(PlatformMixin, ChoicesMixin, models.TextChoices):
     LINUX = 'linux', 'Linux'
     WINDOWS = 'windows', 'Windows'
     UNIX = 'unix', 'Unix'
@@ -84,14 +84,14 @@ class HostTypes(PlatformMixin, models.TextChoices):
         }
 
 
-class NetworkTypes(PlatformMixin, models.TextChoices):
+class NetworkTypes(PlatformMixin, ChoicesMixin, models.TextChoices):
     SWITCH = 'switch', _("Switch")
     ROUTER = 'router', _("Router")
     FIREWALL = 'firewall', _("Firewall")
     OTHER_NETWORK = 'other_network', _("Other device")
 
 
-class DatabaseTypes(PlatformMixin, models.TextChoices):
+class DatabaseTypes(PlatformMixin, ChoicesMixin, models.TextChoices):
     MYSQL = 'mysql', 'MySQL'
     MARIADB = 'mariadb', 'MariaDB'
     POSTGRESQL = 'postgresql', 'PostgreSQL'
@@ -110,15 +110,15 @@ class DatabaseTypes(PlatformMixin, models.TextChoices):
         return meta
 
 
-class WebTypes(PlatformMixin, models.TextChoices):
+class WebTypes(PlatformMixin, ChoicesMixin, models.TextChoices):
     General = 'general', 'General'
 
 
-class CloudTypes(PlatformMixin, models.TextChoices):
+class CloudTypes(PlatformMixin, ChoicesMixin, models.TextChoices):
     K8S = 'k8s', 'Kubernetes'
 
 
-class AllTypes(metaclass=IncludesTextChoicesMeta):
+class AllTypes(ChoicesMixin, metaclass=IncludesTextChoicesMeta):
     choices: list
     includes = [
         HostTypes, NetworkTypes, DatabaseTypes,
@@ -202,7 +202,7 @@ class AllTypes(metaclass=IncludesTextChoicesMeta):
         return nodes
 
 
-class Protocol(models.TextChoices):
+class Protocol(ChoicesMixin, models.TextChoices):
     ssh = 'ssh', 'SSH'
     rdp = 'rdp', 'RDP'
     telnet = 'telnet', 'Telnet'

apps/assets/models/asset/common.py

@@ -10,6 +10,7 @@ from django.db import models
 from django.utils.translation import ugettext_lazy as _
 
 from orgs.mixins.models import OrgManager, JMSOrgBaseModel
+from ...const import Category
 from ..platform import Platform
 from ..base import AbsConnectivity
 
@@ -121,7 +122,8 @@ class Asset(AbsConnectivity, NodesRelationMixin, JMSOrgBaseModel):
 
     @property
     def type_display(self):
-        return self.platform.type
+        value = self.platform.type
+        return Category.get_label(value)
 
     @property
     def category(self):
@@ -129,7 +131,8 @@ class Asset(AbsConnectivity, NodesRelationMixin, JMSOrgBaseModel):
 
     @property
     def category_display(self):
-        return self.platform.category
+        value = self.platform.category
+        return Category.get_label(value)
 
     def as_node(self):
         from assets.models import Node

apps/common/db/models.py

@@ -85,6 +85,14 @@ class BitOperationChoice:
         return [(cls.NAME_MAP[i], j) for i, j in cls.DB_CHOICES]
 
 
+class ChoicesMixin:
+    _value2label_map_: dict
+
+    @classmethod
+    def get_label(cls, value: (str, int)):
+        return cls._value2label_map_[value]
+
+
 class BaseCreateUpdateModel(models.Model):
     created_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by'))
     updated_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Updated by'))