Refactor base Graphene object types

netbox/extras/graphql/mixins.py

@@ -1,12 +1,29 @@
 import graphene
+from graphene.types.generic import GenericScalar
 
 __all__ = (
+    'CustomFieldsMixin',
     'ImageAttachmentsMixin',
+    'TagsMixin',
 )
 
 
+class CustomFieldsMixin:
+    custom_fields = GenericScalar()
+
+    def resolve_custom_fields(self, info):
+        return self.custom_field_data
+
+
 class ImageAttachmentsMixin:
     image_attachments = graphene.List('extras.graphql.types.ImageAttachmentType')
 
     def resolve_image_attachments(self, info):
         return self.images.restrict(info.context.user, 'view')
+
+
+class TagsMixin:
+    tags = graphene.List(graphene.String)
+
+    def resolve_tags(self, info):
+        return self.tags.all()

netbox/netbox/graphql/types.py

@@ -1,8 +1,8 @@
-import graphene
 from django.contrib.contenttypes.models import ContentType
-from graphene.types.generic import GenericScalar
 from graphene_django import DjangoObjectType
 
+from extras.graphql.mixins import CustomFieldsMixin, TagsMixin
+
 __all__ = (
     'BaseObjectType',
     'ObjectType',
@@ -27,31 +27,21 @@ class BaseObjectType(DjangoObjectType):
         return queryset.restrict(info.context.user, 'view')
 
 
-class ObjectType(BaseObjectType):
+class ObjectType(CustomFieldsMixin, BaseObjectType):
     """
-    Extends BaseObjectType with support for custom field data.
+    Extends BaseObjectType with support for custom fields.
     """
-    custom_fields = GenericScalar()
-
     class Meta:
         abstract = True
 
-    def resolve_custom_fields(self, info):
-        return self.custom_field_data
 
-
-class TaggedObjectType(ObjectType):
+class TaggedObjectType(CustomFieldsMixin, TagsMixin, BaseObjectType):
     """
-    Extends ObjectType with support for Tags
+    Extends BaseObjectType with support for custom fields and tags
     """
-    tags = graphene.List(graphene.String)
-
     class Meta:
         abstract = True
 
-    def resolve_tags(self, info):
-        return self.tags.all()
-
 
 #
 # Miscellaneous types