Restricting DRF serialized fields using only and defer.

restricted_fields.py

from collections import OrderedDict

from rest_framework.fields import SkipField
from rest_framework.relations import PKOnlyObject


class RestrictedFieldsSerializerMixin(object):
    """
    API Serializer mixin which provides support for restricting serialized data to only a subset of fields.

    This requires using the ``only`` and ``defer`` query parameters.

    ---
    only: Restricted to only a subset of fields
    defer: All other fields except the listed fields.
    ---

    Examples:
        GET https://.../api/users/?only=id&only=name

        # This returns the serialized data with only the `id` and `name` fields.

        [
            {
              "id":257,
              "name": "Test 1"
            },
            {
              "id":365,
              "name": "Test 2"
            },
            {
              "id":378,
              "name": "Test 3"
            },
            {
               "id":377,
               "name": "Test 4"
            }
        ]
        
        GET https://.../api/users/?defer=id

        # This returns the serialized data deferring the `id` field (i.e Return all fields except the `id`).

        [
            {
              "name": "Test 1",
              "age": 2
            },
            {
              "name": "Test 2",
              "age": 2
            },
            {
              "name": "Test 3",
              "age": 2
            },
            {
               "name": "Test 4",
               "age": 2
            }
        ]

    """
    RESTRICTED_FIELDS_PARAM = 'only'
    DEFERRED_FIELDS_PARAM = 'defer'

    def to_representation(self, instance):
        """
        Convert Model Object instance -> Dict of primitive datatypes.
        """
        request = self.context['request']
        ret = OrderedDict()
        restricted_fields = request.query_params.getlist(self.RESTRICTED_FIELDS_PARAM)
        deferred_fields = request.query_params.getlist(self.DEFERRED_FIELDS_PARAM)
        fields = self._readable_fields

        if restricted_fields:
            fields = [f for f in fields if f.field_name in restricted_fields]

        if deferred_fields:
            fields = [f for f in fields if f.field_name not in deferred_fields]

        for field in fields:
            try:
                attribute = field.get_attribute(instance)
            except SkipField:
                continue

            # We skip `to_representation` for `None` values so that fields do
            # not have to explicitly deal with that case.
            #
            # For related fields with `use_pk_only_optimization` we need to
            # resolve the pk value.
            check_for_none = attribute.pk if isinstance(attribute, PKOnlyObject) else attribute
            if check_for_none is None:
                ret[field.field_name] = None
            else:
                ret[field.field_name] = field.to_representation(attribute)

        return ret