Rename and refactor CustomFieldForm

netbox/extras/forms.py

@@ -21,59 +21,58 @@ from .utils import FeatureQuery
 # Custom fields
 #
 
-class CustomFieldForm(forms.Form):
+class CustomFieldsMixin:
     """
-    Extend Form to include custom field support.
+    Extend a Form to include custom field support.
     """
-    model = None
-
     def __init__(self, *args, **kwargs):
-        if self.model is None:
-            raise NotImplementedError("CustomFieldForm must specify a model class.")
-        self.custom_fields = []
-
-        super().__init__(*args, **kwargs)
-
-        # Append relevant custom fields to the form instance
-        obj_type = ContentType.objects.get_for_model(self.model)
-        for cf in CustomField.objects.filter(content_types=obj_type):
-            field_name = 'cf_{}'.format(cf.name)
-            self.fields[field_name] = cf.to_form_field()
-
-            # Annotate the field in the list of CustomField form fields
-            self.custom_fields.append(field_name)
-
-
-class CustomFieldModelForm(forms.ModelForm):
-    """
-    Extend ModelForm to include custom field support.
-    """
-
-    def __init__(self, *args, **kwargs):
-
-        self.obj_type = ContentType.objects.get_for_model(self._meta.model)
         self.custom_fields = []
 
         super().__init__(*args, **kwargs)
 
         self._append_customfield_fields()
 
+    def _get_content_type(self):
+        """
+        Return the ContentType of the form's model.
+        """
+        if not hasattr(self, 'model'):
+            raise NotImplementedError(f"{self.__class__.__name__} must specify a model class.")
+        return ContentType.objects.get_for_model(self.model)
+
+    def _get_form_field(self, customfield):
+        return customfield.to_form_field()
+
     def _append_customfield_fields(self):
         """
-        Append form fields for all CustomFields assigned to this model.
+        Append form fields for all CustomFields assigned to this object type.
         """
+        content_type = self._get_content_type()
+
         # Append form fields; assign initial values if modifying and existing object
-        for cf in CustomField.objects.filter(content_types=self.obj_type):
-            field_name = 'cf_{}'.format(cf.name)
-            if self.instance.pk:
-                self.fields[field_name] = cf.to_form_field(set_initial=False)
-                self.fields[field_name].initial = self.instance.custom_field_data.get(cf.name)
-            else:
-                self.fields[field_name] = cf.to_form_field()
+        for customfield in CustomField.objects.filter(content_types=content_type):
+            field_name = f'cf_{customfield.name}'
+            self.fields[field_name] = self._get_form_field(customfield)
 
             # Annotate the field in the list of CustomField form fields
             self.custom_fields.append(field_name)
 
+
+class CustomFieldModelForm(CustomFieldsMixin, forms.ModelForm):
+    """
+    Extend ModelForm to include custom field support.
+    """
+    def _get_content_type(self):
+        return ContentType.objects.get_for_model(self._meta.model)
+
+    def _get_form_field(self, customfield):
+        if self.instance.pk:
+            form_field = customfield.to_form_field(set_initial=False)
+            form_field.initial = self.instance.custom_field_data.get(customfield.name, None)
+            return form_field
+
+        return customfield.to_form_field()
+
     def clean(self):
 
         # Save custom field data on instance
@@ -85,15 +84,8 @@ class CustomFieldModelForm(forms.ModelForm):
 
 class CustomFieldModelCSVForm(CSVModelForm, CustomFieldModelForm):
 
-    def _append_customfield_fields(self):
-
-        # Append form fields
-        for cf in CustomField.objects.filter(content_types=self.obj_type):
-            field_name = 'cf_{}'.format(cf.name)
-            self.fields[field_name] = cf.to_form_field(for_csv_import=True)
-
-            # Annotate the field in the list of CustomField form fields
-            self.custom_fields.append(field_name)
+    def _get_form_field(self, customfield):
+        return customfield.to_form_field(for_csv_import=True)
 
 
 class CustomFieldBulkEditForm(BulkEditForm):