diff --git a/netbox/extras/graphql/mixins.py b/netbox/extras/graphql/mixins.py index 5c857db8c..5004ab841 100644 --- a/netbox/extras/graphql/mixins.py +++ b/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() diff --git a/netbox/netbox/graphql/types.py b/netbox/netbox/graphql/types.py index 2958247b7..5e9d53338 100644 --- a/netbox/netbox/graphql/types.py +++ b/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