33"""
44
55import typing as t
6- from enum import Enum
6+ from enum import Enum , Flag
77
88from django .db .models import Field as ModelField
9- from django_filters import (
10- Filter ,
11- TypedChoiceFilter ,
12- TypedMultipleChoiceFilter ,
13- filterset ,
14- )
15-
16- from django_enum .fields import EnumField
17- from django_enum .forms import EnumChoiceField , EnumMultipleChoiceField
9+ from django .db .models import Q
10+ from django_filters import filterset
11+ from django_filters .filters import Filter , TypedChoiceFilter , TypedMultipleChoiceFilter
12+ from django_filters .utils import try_dbfield
13+
14+ from django_enum .fields import EnumField , FlagField
15+ from django_enum .forms import EnumChoiceField , EnumFlagField , EnumMultipleChoiceField
1816from django_enum .utils import choices
1917
2018__all__ = [
2119 "EnumFilter" ,
2220 "MultipleEnumFilter" ,
21+ "EnumFlagFilter" ,
2322 "FilterSet" ,
2423]
2524
2625
2726class EnumFilter (TypedChoiceFilter ):
2827 """
29- Use this filter class instead of :ref:` ChoiceFilter <django-filter:choice-filter> `
28+ Use this filter class instead of :class:`~django_filters.filters. ChoiceFilter`
3029 to get filters to accept :class:`~enum.Enum` labels and symmetric properties.
3130
3231 For example if we have an enumeration field defined with the following
@@ -45,7 +44,7 @@ class Color(TextChoices):
4544
4645 color = EnumField(Color)
4746
48- The default :ref:` ChoiceFilter <django-filter:choice-filter> ` will only work with
47+ The default :class:`~django_filters.filters. ChoiceFilter` will only work with
4948 the enumeration values: ?color=R, ?color=G, ?color=B. ``EnumFilter`` will accept
5049 query parameter values from any of the symmetric properties: ?color=Red,
5150 ?color=ff0000, etc...
@@ -54,66 +53,157 @@ class Color(TextChoices):
5453 filter on
5554 :param strict: If False (default), values not in the enumeration will
5655 be searchable.
57- :param kwargs: Any additional arguments for base classes
56+ :param kwargs: Any additional arguments from the base classes
57+ (:class:`~django_filters.filters.TypedChoiceFilter`)
5858 """
5959
6060 enum : t .Type [Enum ]
6161 field_class = EnumChoiceField
6262
63- def __init__ (self , * , enum : t .Type [Enum ], strict : bool = False , ** kwargs ):
63+ def __init__ (self , * , enum : t .Type [Enum ], ** kwargs ):
6464 self .enum = enum
6565 super ().__init__ (
6666 enum = enum ,
6767 choices = kwargs .pop ("choices" , choices (self .enum )),
68- strict = strict ,
6968 ** kwargs ,
7069 )
7170
7271
7372class MultipleEnumFilter (TypedMultipleChoiceFilter ):
7473 """
7574 Use this filter class instead of
76- :ref:`MultipleChoiceFilter <django-filter:multiple-choice-filter>`
77- to get filters to accept multiple :class:`~enum.Enum` labels and symmetric
78- properties.
75+ :class:`~django_filters.filters.MultipleChoiceFilter` to get filters to accept
76+ multiple :class:`~enum.Enum` labels and symmetric properties.
7977
8078 :param enum: The class of the enumeration containing the values to
8179 filter on
8280 :param strict: If False (default), values not in the enumeration will
8381 be searchable.
84- :param kwargs: Any additional arguments for base classes
82+ :param conjoined: If True require all values to be present, if False require any
83+ :param kwargs: Any additional arguments from base classes,
84+ (:class:`~django_filters.filters.TypedMultipleChoiceFilter`)
8585 """
8686
8787 enum : t .Type [Enum ]
8888 field_class = EnumMultipleChoiceField
8989
90- def __init__ (self , * , enum : t .Type [Enum ], strict : bool = False , ** kwargs ):
90+ def __init__ (
91+ self ,
92+ * ,
93+ enum : t .Type [Enum ],
94+ conjoined : bool = False ,
95+ ** kwargs ,
96+ ):
97+ self .enum = enum
98+ super ().__init__ (
99+ enum = enum ,
100+ choices = kwargs .pop ("choices" , choices (self .enum )),
101+ conjoined = conjoined ,
102+ ** kwargs ,
103+ )
104+
105+
106+ class EnumFlagFilter (TypedMultipleChoiceFilter ):
107+ """
108+ Use this filter class instead of
109+ :class:`~django_filters.filters.MultipleChoiceFilter` to get filters to accept
110+ multiple :class:`~enum.Enum` labels and symmetric properties.
111+
112+ :param enum: The class of the enumeration containing the values to
113+ filter on
114+ :param strict: If False (default), values not in the enumeration will
115+ be searchable.
116+ :param conjoined: If True use :ref:`has_all` lookup, otherwise use :ref:`has_any`
117+ (default)
118+ :param kwargs: Any additional arguments from base classes,
119+ (:class:`~django_filters.filters.TypedMultipleChoiceFilter`)
120+ """
121+
122+ enum : t .Type [Flag ]
123+ field_class = EnumFlagField
124+
125+ def __init__ (
126+ self ,
127+ * ,
128+ enum : t .Type [Flag ],
129+ conjoined : bool = False ,
130+ strict : bool = False ,
131+ ** kwargs ,
132+ ):
91133 self .enum = enum
134+ self .lookup_expr = "has_all" if conjoined else "has_any"
92135 super ().__init__ (
93136 enum = enum ,
94137 choices = kwargs .pop ("choices" , choices (self .enum )),
95138 strict = strict ,
139+ conjoined = conjoined ,
96140 ** kwargs ,
97141 )
98142
143+ def filter (self , qs , value ):
144+ if value == self .null_value :
145+ value = None
146+
147+ if not value :
148+ return qs
149+
150+ if self .is_noop (qs , value ):
151+ return qs
152+
153+ qs = self .get_method (qs )(Q (** self .get_filter_predicate (value )))
154+ return qs .distinct () if self .distinct else qs
155+
99156
100157class FilterSet (filterset .FilterSet ):
101158 """
102159 Use this class instead of the :doc:`django-filter <django-filter:index>`
103- :doc:`FilterSet <django-filter:ref/ filterset>` class to automatically set all
160+ :class:`~django_filters. filterset.FilterSet` to automatically set all
104161 :class:`~django_enum.fields.EnumField` filters to
105162 :class:`~django_enum.filters.EnumFilter` by default instead of
106- :ref:` ChoiceFilter <django-filter:choice-filter> `.
163+ :class:`~django_filters.filters. ChoiceFilter`.
107164 """
108165
166+ @staticmethod
167+ def enum_extra (f : EnumField ) -> t .Dict [str , t .Any ]:
168+ return {"enum" : f .enum , "choices" : f .choices }
169+
170+ FILTER_DEFAULTS = {
171+ ** {
172+ FlagField : {
173+ "filter_class" : EnumFlagFilter ,
174+ "extra" : enum_extra ,
175+ },
176+ EnumField : {
177+ "filter_class" : EnumFilter ,
178+ "extra" : enum_extra ,
179+ },
180+ },
181+ ** filterset .FilterSet .FILTER_DEFAULTS ,
182+ }
183+
109184 @classmethod
110185 def filter_for_lookup (
111186 cls , field : ModelField , lookup_type : str
112187 ) -> t .Tuple [t .Optional [t .Type [Filter ]], t .Dict [str , t .Any ]]:
113188 """For EnumFields use the EnumFilter class by default"""
114189 if isinstance (field , EnumField ):
115- return EnumFilter , {
116- "enum" : field .enum ,
117- "strict" : getattr (field , "strict" , False ),
118- }
190+ data = (
191+ try_dbfield (
192+ {
193+ ** cls .FILTER_DEFAULTS ,
194+ ** (
195+ getattr (getattr (cls , "_meta" , None ), "filter_overrides" , {})
196+ ),
197+ }.get ,
198+ field .__class__ ,
199+ )
200+ or {}
201+ )
202+ return (
203+ data ["filter_class" ],
204+ {
205+ ** cls .enum_extra (field ),
206+ ** data .get ("extra" , lambda f : {})(field ),
207+ },
208+ )
119209 return super ().filter_for_lookup (field , lookup_type )
0 commit comments