Parcourir la source

Reimplemented FilterChoiceField

Jeremy Stretch il y a 8 ans
Parent
commit
e618bf40ec

+ 10 - 7
netbox/circuits/forms.py

@@ -1,4 +1,5 @@
 from django import forms
+from django.db.models import Count
 
 from dcim.models import Site, Device, Interface, Rack, IFACE_FF_VIRTUAL
 from extras.forms import CustomFieldForm, CustomFieldBulkEditForm, CustomFieldFilterForm
@@ -6,7 +7,7 @@ from tenancy.forms import bulkedit_tenant_choices
 from tenancy.models import Tenant
 from utilities.forms import (
     APISelect, BootstrapMixin, BulkImportForm, CommentField, CSVDataField, FilterChoiceField, Livesearch, SmallTextarea,
-    SlugField, get_filter_choices,
+    SlugField,
 )
 
 from .models import Circuit, CircuitType, Provider
@@ -59,7 +60,7 @@ class ProviderBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 
 class ProviderFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = Provider
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug'))
+    site = FilterChoiceField(queryset=Site.objects.all(), to_field_name='slug')
 
 
 #
@@ -185,8 +186,10 @@ class CircuitBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 
 class CircuitFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = Circuit
-    type = FilterChoiceField(choices=get_filter_choices(CircuitType, id_field='slug', count_field='circuits'))
-    provider = FilterChoiceField(choices=get_filter_choices(Provider, id_field='slug', count_field='circuits'))
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='circuits',
-                                                          null_option='None'))
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug', count_field='circuits'))
+    type = FilterChoiceField(queryset=CircuitType.objects.annotate(filter_count=Count('circuits')),
+                             to_field_name='slug')
+    provider = FilterChoiceField(queryset=Provider.objects.annotate(filter_count=Count('circuits')),
+                                 to_field_name='slug')
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('circuits')), to_field_name='slug',
+                               null_option=(0, 'None'))
+    site = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('circuits')), to_field_name='slug')

+ 22 - 26
netbox/dcim/forms.py

@@ -1,7 +1,7 @@
 import re
 
 from django import forms
-from django.db.models import Q
+from django.db.models import Count, Q
 
 from extras.forms import CustomFieldForm, CustomFieldBulkEditForm, CustomFieldFilterForm
 from ipam.models import IPAddress
@@ -10,7 +10,6 @@ from tenancy.models import Tenant
 from utilities.forms import (
     APISelect, add_blank_choice, BootstrapMixin, BulkImportForm, CommentField, CSVDataField, ExpandableNameField,
     FilterChoiceField, FlexibleModelChoiceField, Livesearch, SelectWithDisabled, SmallTextarea, SlugField,
-    get_filter_choices
 )
 
 from .models import (
@@ -120,8 +119,8 @@ class SiteBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 
 class SiteFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = Site
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='sites',
-                                                          null_option='None'))
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('sites')), to_field_name='slug',
+                               null_option=(0, 'None'))
 
 
 #
@@ -137,7 +136,7 @@ class RackGroupForm(forms.ModelForm, BootstrapMixin):
 
 
 class RackGroupFilterForm(forms.Form, BootstrapMixin):
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug', count_field='rack_groups'))
+    site = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('rack_groups')), to_field_name='slug')
 
 
 #
@@ -246,14 +245,13 @@ class RackBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 
 class RackFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = Rack
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug', count_field='racks'))
-    group_id = FilterChoiceField(choices=get_filter_choices(RackGroup, select_related=['site'], count_field='racks',
-                                                            null_option='None'),
-                                 label='Rack Group')
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='racks',
-                                                          null_option='None'))
-    role = FilterChoiceField(choices=get_filter_choices(RackRole, id_field='slug', count_field='racks',
-                                                        null_option='None'))
+    site = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('racks')), to_field_name='slug')
+    group_id = FilterChoiceField(queryset=RackGroup.objects.select_related('site')
+                                 .annotate(filter_count=Count('racks')), label='Rack group', null_option=(0, 'None'))
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('racks')), to_field_name='slug',
+                               null_option=(0, 'None'))
+    role = FilterChoiceField(queryset=RackRole.objects.annotate(filter_count=Count('racks')), to_field_name='slug',
+                             null_option=(0, 'None'))
 
 
 #
@@ -288,8 +286,8 @@ class DeviceTypeBulkEditForm(forms.Form, BootstrapMixin):
 
 
 class DeviceTypeFilterForm(forms.Form, BootstrapMixin):
-    manufacturer = FilterChoiceField(choices=get_filter_choices(Manufacturer, id_field='slug',
-                                                                count_field='device_types'))
+    manufacturer = FilterChoiceField(queryset=Manufacturer.objects.annotate(filter_count=Count('device_types')),
+                                     to_field_name='slug')
 
 
 #
@@ -594,18 +592,16 @@ class DeviceBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 
 class DeviceFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = Device
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug', count_field='racks__devices'))
-    rack_group_id = FilterChoiceField(choices=get_filter_choices(RackGroup, select_related=['site'],
-                                                                 count_field='racks__devices'),
+    site = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('racks__devices')), to_field_name='slug')
+    rack_group_id = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('racks__devices')),
                                       label='Rack Group')
-    role = FilterChoiceField(choices=get_filter_choices(DeviceRole, id_field='slug', count_field='devices'))
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='devices',
-                                                          null_option='None'))
-    device_type_id = FilterChoiceField(choices=get_filter_choices(DeviceType, select_related=['manufacturer'],
-                                                                  count_field='instances'),
-                                       label='Type')
-    platform = FilterChoiceField(choices=get_filter_choices(Platform, id_field='slug', count_field='devices',
-                                                            null_option='None'))
+    role = FilterChoiceField(queryset=DeviceRole.objects.annotate(filter_count=Count('devices')), to_field_name='slug')
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('devices')), to_field_name='slug',
+                               null_option=(0, 'None'))
+    device_type_id = FilterChoiceField(queryset=DeviceType.objects.select_related('manufacturer')
+                                       .annotate(filter_count=Count('instances')), label='Type')
+    platform = FilterChoiceField(queryset=Platform.objects.annotate(filter_count=Count('devices')),
+                                 to_field_name='slug', null_option=(0, 'None'))
     status = forms.NullBooleanField(required=False, widget=forms.Select(choices=FORM_STATUS_CHOICES))
 
 

+ 26 - 28
netbox/ipam/forms.py

@@ -7,7 +7,6 @@ from tenancy.forms import bulkedit_tenant_choices
 from tenancy.models import Tenant
 from utilities.forms import (
     APISelect, BootstrapMixin, CSVDataField, BulkImportForm, FilterChoiceField, Livesearch, SlugField,
-    get_filter_choices,
 )
 
 from .models import (
@@ -74,8 +73,8 @@ class VRFBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 
 class VRFFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = VRF
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='vrfs',
-                                                          null_option='None'))
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('vrfs')), to_field_name='slug',
+                               null_option=(0, None))
 
 
 #
@@ -129,7 +128,8 @@ class AggregateBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 class AggregateFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = Aggregate
     family = forms.ChoiceField(required=False, choices=IP_FAMILY_CHOICES, label='Address Family')
-    rir = FilterChoiceField(choices=get_filter_choices(RIR, id_field='slug', count_field='aggregates'), label='RIR')
+    rir = FilterChoiceField(queryset=RIR.objects.annotate(filter_count=Count('aggregates')), to_field_name='slug',
+                            label='RIR')
 
 
 #
@@ -273,16 +273,15 @@ class PrefixFilterForm(BootstrapMixin, CustomFieldFilterForm):
         'placeholder': 'Network',
     }))
     family = forms.ChoiceField(required=False, choices=IP_FAMILY_CHOICES, label='Address Family')
-    vrf = FilterChoiceField(choices=get_filter_choices(VRF, count_field='prefixes', null_option='Global'),
-                            label='VRF')
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='prefixes',
-                                                          null_option='None'),
-                               label='Tenant')
-    status = FilterChoiceField(choices=prefix_status_choices)
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug', count_field='prefixes',
-                                                        null_option='None'))
-    role = FilterChoiceField(choices=get_filter_choices(Role, id_field='slug', count_field='prefixes',
-                                                        null_option='None'))
+    vrf = FilterChoiceField(queryset=VRF.objects.annotate(filter_count=Count('prefixes')), to_field_name='slug',
+                            label='VRF', null_option=(0, 'Global'))
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('prefixes')), to_field_name='slug',
+                               null_option=(0, 'None'))
+    status = forms.MultipleChoiceField(choices=prefix_status_choices)
+    site = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('prefixes')), to_field_name='slug',
+                             null_option=(0, 'None'))
+    role = FilterChoiceField(queryset=Role.objects.annotate(filter_count=Count('prefixes')), to_field_name='slug',
+                             null_option=(0, 'None'))
     expand = forms.BooleanField(required=False, label='Expand prefix hierarchy')
 
 
@@ -419,11 +418,10 @@ class IPAddressFilterForm(BootstrapMixin, CustomFieldFilterForm):
         'placeholder': 'Prefix',
     }))
     family = forms.ChoiceField(required=False, choices=IP_FAMILY_CHOICES, label='Address Family')
-    vrf = FilterChoiceField(choices=get_filter_choices(VRF, count_field='ip_addresses', null_option='None'),
-                            label='VRF')
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='ip_addresses',
-                                                          null_option='None'),
-                               label='Tenant')
+    vrf = FilterChoiceField(queryset=VRF.objects.annotate(filter_count=Count('ip_addresses')), to_field_name='slug',
+                            label='VRF', null_option=(0, 'Global'))
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('ip_addresses')),
+                               to_field_name='slug', null_option=(0, 'None'))
 
 
 #
@@ -439,7 +437,7 @@ class VLANGroupForm(forms.ModelForm, BootstrapMixin):
 
 
 class VLANGroupFilterForm(forms.Form, BootstrapMixin):
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug', count_field='vlan_groups'))
+    site = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('vlan_groups')), to_field_name='slug')
 
 
 #
@@ -526,11 +524,11 @@ def vlan_status_choices():
 
 class VLANFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = VLAN
-    site = FilterChoiceField(choices=get_filter_choices(Site, id_field='slug', count_field='vlans'))
-    group_id = FilterChoiceField(choices=get_filter_choices(VLANGroup, select_related=['site'], count_field='vlans',
-                                                            null_option='None'),
-                                 label='VLAN Group')
-    tenant = FilterChoiceField(choices=get_filter_choices(Tenant, id_field='slug', count_field='vlans',
-                                                          null_option='None'))
-    status = FilterChoiceField(choices=vlan_status_choices)
-    role = FilterChoiceField(choices=get_filter_choices(Role, id_field='slug', count_field='vlans', null_option='None'))
+    site = FilterChoiceField(queryset=Site.objects.annotate(filter_count=Count('vlans')), to_field_name='slug')
+    group_id = FilterChoiceField(queryset=VLANGroup.objects.annotate(filter_count=Count('vlans')), label='VLAN group',
+                                 null_option=(0, 'None'))
+    tenant = FilterChoiceField(queryset=Tenant.objects.annotate(filter_count=Count('vlans')), to_field_name='slug',
+                               null_option=(0, 'None'))
+    status = forms.MultipleChoiceField(choices=vlan_status_choices)
+    role = FilterChoiceField(queryset=Role.objects.annotate(filter_count=Count('vlans')), to_field_name='slug',
+                             null_option=(0, 'None'))

+ 3 - 4
netbox/secrets/forms.py

@@ -2,11 +2,10 @@ from Crypto.Cipher import PKCS1_OAEP
 from Crypto.PublicKey import RSA
 
 from django import forms
+from django.db.models import Count
 
 from dcim.models import Device
-from utilities.forms import (
-    BootstrapMixin, BulkImportForm, CSVDataField, FilterChoiceField, SlugField, get_filter_choices,
-)
+from utilities.forms import BootstrapMixin, BulkImportForm, CSVDataField, FilterChoiceField, SlugField
 
 from .models import Secret, SecretRole, UserKey
 
@@ -97,7 +96,7 @@ class SecretBulkEditForm(forms.Form, BootstrapMixin):
 
 
 class SecretFilterForm(forms.Form, BootstrapMixin):
-    role = FilterChoiceField(choices=get_filter_choices(SecretRole, id_field='slug', count_field='secrets'))
+    role = FilterChoiceField(queryset=SecretRole.objects.annotate(filter_count=Count('secrets')), to_field_name='slug')
 
 
 #

+ 4 - 5
netbox/tenancy/forms.py

@@ -1,9 +1,8 @@
 from django import forms
+from django.db.models import Count
 
 from extras.forms import CustomFieldForm, CustomFieldBulkEditForm, CustomFieldFilterForm
-from utilities.forms import (
-    BootstrapMixin, BulkImportForm, CommentField, CSVDataField, FilterChoiceField, SlugField, get_filter_choices,
-)
+from utilities.forms import BootstrapMixin, BulkImportForm, CommentField, CSVDataField, FilterChoiceField, SlugField
 
 from .models import Tenant, TenantGroup
 
@@ -77,5 +76,5 @@ class TenantBulkEditForm(BootstrapMixin, CustomFieldBulkEditForm):
 
 class TenantFilterForm(BootstrapMixin, CustomFieldFilterForm):
     model = Tenant
-    group = FilterChoiceField(choices=get_filter_choices(TenantGroup, id_field='slug', count_field='tenants',
-                                                         null_option='None'))
+    group = FilterChoiceField(queryset=TenantGroup.objects.annotate(filter_count=Count('tenants')),
+                              to_field_name='slug', null_option=(0, 'None'))

+ 19 - 25
netbox/utilities/forms.py

@@ -1,4 +1,5 @@
 import csv
+import itertools
 import re
 
 from django import forms
@@ -35,29 +36,6 @@ def add_blank_choice(choices):
     return ((None, '---------'),) + choices
 
 
-def get_filter_choices(model, id_field='pk', select_related=[], count_field=None, null_option=None):
-    """
-    Return a list of choices suitable for a ChoiceField.
-
-    :param model: The base model to use for the queryset
-    :param id_field: Field to use as the object identifier
-    :param select_related: Any related tables to include
-    :param count_field: The field to use for a child COUNT() (optional)
-    :param null_option: A choice to include at the beginning of the list serving as "null"
-    """
-    queryset = model.objects.all()
-    if select_related:
-        queryset = queryset.select_related(*select_related)
-    if count_field:
-        queryset = queryset.annotate(child_count=Count(count_field))
-        choices = [(getattr(obj, id_field), u'{} ({})'.format(obj, obj.child_count)) for obj in queryset]
-    else:
-        choices = [(getattr(obj, id_field), u'{}'.format(obj)) for obj in queryset]
-    if null_option:
-        choices = [(0, null_option)] + choices
-    return choices
-
-
 #
 # Widgets
 #
@@ -250,15 +228,31 @@ class SlugField(forms.SlugField):
         self.widget.attrs['slug-source'] = slug_source
 
 
-class FilterChoiceField(forms.MultipleChoiceField):
+class FilterChoiceField(forms.ModelMultipleChoiceField):
+    iterator = forms.models.ModelChoiceIterator
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, null_option=None, *args, **kwargs):
+        self.null_option = null_option
         if 'required' not in kwargs:
             kwargs['required'] = False
         if 'widget' not in kwargs:
             kwargs['widget'] = forms.SelectMultiple(attrs={'size': 6})
         super(FilterChoiceField, self).__init__(*args, **kwargs)
 
+    def label_from_instance(self, obj):
+        if hasattr(obj, 'filter_count'):
+            return u'{} ({})'.format(obj, obj.filter_count)
+        return force_text(obj)
+
+    def _get_choices(self):
+        if hasattr(self, '_choices'):
+            return self._choices
+        if self.null_option is not None:
+            return itertools.chain([self.null_option], self.iterator(self))
+        return self.iterator(self)
+
+    choices = property(_get_choices, forms.ChoiceField._set_choices)
+
 
 #
 # Forms