Custom class for ObjectType

Author: bctiemann

netbox_custom_objects/forms.py

@@ -3,7 +3,9 @@
 from django import forms
 from django.utils.translation import gettext_lazy as _
 
-from netbox_custom_objects.models import CustomObject, CustomObjectType, CustomObjectTypeField, CustomObjectRelation
+from netbox_custom_objects.models import (
+    CustomObject, CustomObjectType, CustomObjectTypeField, CustomObjectRelation, CustomObjectObjectType
+)
 
 from netbox.forms import NetBoxModelForm
 from core.models import ObjectType
@@ -53,11 +55,12 @@ class CustomObjectTypeFieldForm(CustomFieldForm):
         required=True,
         label=_('Custom object type')
     )
-    # related_object_type = ContentTypeChoiceField(
-    #     label=_('Related object type'),
-    #     queryset=ObjectType.objects.public(),
-    #     help_text=_("Type of the related object (for object/multi-object fields only)")
-    # )
+    related_object_type = ContentTypeChoiceField(
+        label=_('Related object type'),
+        queryset=CustomObjectObjectType.objects.public(),
+        # choices=[("foo", "bar")],
+        help_text=_("Type of the related object (for object/multi-object fields only)")
+    )
 
     fieldsets = (
         FieldSet(
@@ -84,6 +87,8 @@ def __init__(self, *args, **kwargs):
         if self.instance.pk:
             self.fields['custom_object_type'].disabled = True
 
+        # self.fields['related_object_type'].choices = [("foo", "bar")]
+
 
 class CustomObjectForm(NetBoxModelForm):
     fieldsets = (

netbox_custom_objects/models.py

@@ -3,12 +3,13 @@
 import json
 import re
 from datetime import datetime, date
+from copy import deepcopy
 
 import django_filters
 from django import forms
 from django.conf import settings
 from django.db import models
-from django.db.models import F, Func, Value
+from django.db.models import F, Func, Value, Q, QuerySet
 from django.db.models.expressions import RawSQL
 from django.apps import apps
 from django.contrib.contenttypes.models import ContentType
@@ -17,6 +18,7 @@
 from django.utils.html import escape
 from django.utils.safestring import mark_safe
 from django.utils.translation import gettext_lazy as _
+from core.models.contenttypes import ObjectTypeManager
 from netbox.models import NetBoxModel, ChangeLoggedModel
 from netbox.models.features import (
     BookmarksMixin,
@@ -31,6 +33,7 @@
     TagsMixin,
     EventRulesMixin,
 )
+from netbox.registry import registry
 from extras.choices import (
     CustomFieldTypeChoices, CustomFieldFilterLogicChoices, CustomFieldUIVisibleChoices, CustomFieldUIEditableChoices
 )
@@ -948,3 +951,28 @@ class CustomObjectRelation(models.Model):
     def instance(self):
         model_class = self.field.related_object_type.model_class()
         return model_class.objects.get(pk=self.object_id)
+
+
+class CustomObjectObjectTypeManager(ObjectTypeManager):
+
+    def public(self):
+        """
+        Filter the base queryset to return only ContentTypes corresponding to "public" models; those which are listed
+        in registry['models'] and intended for reference by other objects.
+        """
+        q = Q()
+        model_registry = deepcopy(registry['models'])
+        model_registry['netbox_custom_objects'] = self.get_queryset().values_list('model', flat=True)
+        for app_label, models in model_registry.items():
+            q |= Q(app_label=app_label, model__in=models)
+        return self.get_queryset().filter(q)
+
+
+class CustomObjectObjectType(ContentType):
+    """
+    Wrap Django's native ContentType model to use our custom manager.
+    """
+    objects = CustomObjectObjectTypeManager()
+
+    class Meta:
+        proxy = True