Add CustomObjectContentTypeChoiceField and object field filtering in edit form

Author: bctiemann

netbox_custom_objects/forms.py

@@ -13,6 +13,7 @@
 from extras.forms import CustomFieldForm
 from utilities.forms.fields import CommentField, ContentTypeChoiceField, DynamicModelChoiceField
 from utilities.forms.rendering import FieldSet
+from utilities.string import title
 
 __all__ = (
     'CustomObjectTypeForm',
@@ -43,6 +44,46 @@ class Meta:
 #         fields = ('name', 'label', 'custom_object_type', 'field_type',)
 
 
+class CustomObjectContentTypeChoiceField(ContentTypeChoiceField):
+    """
+    Selection field for a single content type.
+    """
+    custom_object_types = None
+
+    def __init__(self, queryset, *args, **kwargs):
+        # Order ContentTypes by app_label
+        queryset = queryset.order_by('app_label', 'model')
+        self.custom_object_types = CustomObjectType.objects.all()
+        super().__init__(queryset, *args, **kwargs)
+
+    def object_type_name(self, object_type, include_app=True):
+        """
+        Return a human-friendly ObjectType name (e.g. "DCIM > Site").
+        """
+        try:
+            meta = object_type.model_class()._meta
+            app_label = title(meta.app_config.verbose_name)
+            model_name = title(meta.verbose_name)
+            if include_app:
+                return f'{app_label} > {model_name}'
+            return model_name
+        except AttributeError:
+            if object_type.app_label == 'netbox_custom_objects':
+                try:
+                    model_name = self.custom_object_types.get(slug=object_type.model).name
+                    return f'Custom Objects > {model_name} (Custom)'
+                except CustomObjectType.DoesNotExist:
+                    pass
+            # Model does not exist
+            return f'{object_type.app_label} > {object_type.model}'
+
+    def label_from_instance(self, obj):
+        try:
+            return self.object_type_name(obj)
+        except AttributeError:
+            return super().label_from_instance(obj)
+
+
 class CustomObjectTypeFieldForm(CustomFieldForm):
     # This field should be removed or at least "required" should be defeated
     object_types = forms.CharField(
@@ -55,7 +96,7 @@ class CustomObjectTypeFieldForm(CustomFieldForm):
         required=True,
         label=_('Custom object type')
     )
-    related_object_type = ContentTypeChoiceField(
+    related_object_type = CustomObjectContentTypeChoiceField(
         label=_('Related object type'),
         queryset=CustomObjectObjectType.objects.public(),
         # choices=[("foo", "bar")],

netbox_custom_objects/models.py

@@ -412,6 +412,8 @@ def __str__(self):
 
     @property
     def model_class(self):
+        if self.related_object_type.app_label == 'netbox_custom_objects':
+            return apps.get_model(self.related_object_type.app_label, 'customobjecttype')
         return apps.get_model(self.related_object_type.app_label, self.related_object_type.model)
 
     @property
@@ -626,6 +628,26 @@ def deserialize(self, value):
             return model.objects.filter(pk__in=value)
         return value
 
+    @property
+    def related_custom_object_type(self):
+        if self.related_object_type.app_label != 'netbox_custom_objects':
+            return None
+        return CustomObjectType.objects.get(slug=self.related_object_type.model)
+
+    @property
+    def related_object_queryset(self):
+        if model_class := self.related_object_type.model_class():
+            return model_class.objects.all()
+        return CustomObject.objects.filter(custom_object_type=self.related_custom_object_type)
+
+    @property
+    def related_custom_object_type_filter(self):
+        if self.related_object_type.app_label != 'netbox_custom_objects':
+            return {}
+        return {
+            'custom_object_type': self.related_custom_object_type.id,
+        }
+
     def to_form_field(self, set_initial=True, enforce_required=True, enforce_visibility=True, for_csv_import=False):
         """
         Return a form field suitable for setting a CustomField's value for an object.
@@ -719,30 +741,32 @@ def to_form_field(self, set_initial=True, enforce_required=True, enforce_visibil
 
         # Object
         elif self.type == CustomFieldTypeChoices.TYPE_OBJECT:
-            model = self.related_object_type.model_class()
             field_class = CSVModelChoiceField if for_csv_import else DynamicModelChoiceField
             kwargs = {
-                'queryset': model.objects.all(),
+                'queryset': self.related_object_queryset,
                 'required': required,
                 'initial': initial,
             }
             if not for_csv_import:
-                kwargs['query_params'] = self.related_object_filter
+                query_params = self.related_object_filter or {}
+                query_params.update(self.related_custom_object_type_filter)
+                kwargs['query_params'] = query_params
                 kwargs['selector'] = True
 
             field = field_class(**kwargs)
 
         # Multiple objects
         elif self.type == CustomFieldTypeChoices.TYPE_MULTIOBJECT:
-            model = self.related_object_type.model_class()
             field_class = CSVModelMultipleChoiceField if for_csv_import else DynamicModelMultipleChoiceField
             kwargs = {
-                'queryset': model.objects.all(),
+                'queryset': self.related_object_queryset,
                 'required': required,
                 'initial': initial,
             }
             if not for_csv_import:
-                kwargs['query_params'] = self.related_object_filter
+                query_params = self.related_object_filter or {}
+                query_params.update(self.related_custom_object_type_filter)
+                kwargs['query_params'] = query_params
                 kwargs['selector'] = True
 
             field = field_class(**kwargs)