如何在django-filter中用choice field 的 value 值过滤对象

如果我们有这样一个model：

class IPInfoModel(models.Model):
    TYPE_INTRANET = 1
    TYPE_INTERNET = 2
    IP_TYPES = (
        (TYPE_INTRANET, u'内网'),
        (TYPE_INTERNET, u'外网'),
    )
    ip = models.GenericIPAddressField("IP", unique=True)
    ip_type = models.SmallIntegerField(choices=IP_TYPES)

然后我使用 rest_frame_work + django_filter 做API

from django_filters import rest_framework as django_filters 

class IPInfoFilter(django_filters.FilterSet):
    ip_type = django_filters.ChoiceFilter(choices=IPInfoModel.IP_TYPES)

    class Meta:
        model = IPInfoModel
        fields = ["ip_type",]


class IPInfoViewSet(mixins.ListModelMixin, viewsets.GenericViewSet):
    queryset = IPInfoModel.objects.all()
    serializer_class = IPInfoSerializer
    filter_class = IPInfoFilter

但是这样在过滤 ip_type 时，只能使用choice field 的 key 值 “1”， “2” 来进行过滤。

这样很不直观，那么如何才能使用choice field 的 value 值 “内网”， “外网” 来过滤对象呢？

---

方法一

我的做法是通过自定义一个 ChoiceValueFilter，并覆盖 Django 的 forms.ChoiceField 中的验证函数 valid_value 来实现。

代码如下：

from django import forms

from django_filters import rest_framework as django_filters
from django_filters.conf import settings


class ChoiceValueField(forms.ChoiceField):
    def valid_value(self, value):
        text_value = str(value)

        for k, v in self.choices:
            if value == v or text_value == str(v):
                return True

        return False


class ChoiceValueFilter(django_filters.Filter):
    field_class = ChoiceValueField

    def __init__(self, *args, **kwargs):
        self.null_value = kwargs.get('null_value', settings.NULL_CHOICE_VALUE)
        self.choices = kwargs.get('choices')
        super().__init__(*args, **kwargs)

    def filter(self, qs, value):
        value_map = {v: k for k, v in self.choices}

        return qs.filter(ip_type=value_map[value]) if value else qs

这样做的好处是

1. 比较通用，只要传入对应choice field 的元祖就可以使用了。
2. 仍然会对传入的参数进行验证，保证只能使用 choice field 的 value 值中已有的值进行过滤。

然后改写原有的IPInfoFilter, 改为使用 ChoiceValueFilter 就可以了。

class IPInfoFilter(django_filters.FilterSet):
    ip_type = ChoiceValueFilter(choices=IPInfoModel.IP_TYPES)

    class Meta:
        model = IPInfoModel
        fields = ["ip_type", "is_vip"]

---

方法二

也可以通过自定义filter的方法来实现

参考: https://django-filter.readthedocs.io/en/master/ref/filters.html#method

class IPInfoFilter(django_filters.FilterSet):
    ip_type = django_filters.CharFilter(method='filter_ip_type')


    def filter_ip_type(self, queryset, name, value):
        # create a dictionary string -> integer
        value_map = {v: k for k, v in IPInfoModel.IP_TYPES.items()}
        # get the integer value for the input string
        value = value_map[value]
        return queryset.filter(ip_type=value)