Django dynamic fields model serializer

dynamic_fields_model_serializer.py

from rest_framework import serializers
from typing import Iterable, Optional, Union

from rest_framework import serializers


class DynamicFieldsModelSerializer(serializers.ModelSerializer):
    """
    A ModelSerializer that takes additional `include` and `exclude` arguments.
    These arguments control which fields we include or exclude.

    E.g.,
    If we have a model with a field called `plane` and we want to exclude
    it, we can do:
    `DynamicFieldsModelSerializer(model, exclude=['plane'])`

    If we want to include it, we can do:
    `DynamicFieldsModelSerializer(model, include=['plane'])`
    """
    include: Optional[Iterable[str]]
    exclude: Optional[Iterable[str]]

    def __init__(self, *args, **kwargs) -> None:
        self.include = self.specify_fields('include', kwargs)
        self.exclude = self.specify_fields('exclude', kwargs)
        self.check_only_include_or_exclude()
        self.check_fields_valid()
        super().__init__(*args, **kwargs)
        self.set_fields()

    def specify_fields(self, include_or_exclude: str, kwargs: dict) -> Optional[Iterable[str]]:
        """
        Returns the fields to include or exclude from the serializer.
        """
        arg = kwargs.pop(include_or_exclude, None)
        return self.to_iterable(arg) if arg else None

    def check_only_include_or_exclude(self) -> None:
        """
        Raises an error if both `include` and `exclude` are specified.
        """
        if self.include and self.exclude:
            raise ValueError("Specify only fields to include or exclude, not both")

    def check_fields_valid(self) -> None:
        """
        Check that the model contains the specified fields.
        """
        target_fields = self.include or self.exclude
        if target_fields and not set(target_fields).issubset(set(self.fields)):
            raise ValueError(
                "Specified fields {} are not in model fields {}.".
                    format(target_fields, list(self.fields.keys()))
            )

    def set_fields(self) -> None:
        """
        Sets the fields to include or exclude on the serializer.
        """
        if self.include:
            existing = set(self.fields)
            for field_name in existing - set(self.include):
                self.fields.pop(field_name)
        elif self.exclude:
            for field_name in set(self.exclude):
                self.fields.pop(field_name)

    @staticmethod
    def to_iterable(arg: Union[Iterable[str], str]) -> Iterable[str]:
        """
        Returns the argument as an iterable.
        """
        return arg if type(arg) in [list, tuple] else [arg]