如何在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)