4525: better support for ScriptVariable fields with multiple values

netbox/extras/forms.py

@@ -1,6 +1,7 @@
 from django import forms
 from django.contrib.auth.models import User
 from django.contrib.contenttypes.models import ContentType
+from django.utils.datastructures import MultiValueDict
 from mptt.forms import TreeNodeMultipleChoiceField
 from taggit.forms import TagField as TagField_
 
@@ -431,6 +432,23 @@ class ScriptForm(BootstrapMixin, forms.Form):
     )
 
     def __init__(self, vars, *args, commit_default=True, **kwargs):
+        from .scripts import MultiObjectVar
+
+        if 'initial' in kwargs:
+            orig_initial = kwargs['initial']
+            if isinstance(orig_initial, MultiValueDict):
+                # Convert MultiValueDict to a normal dict with single or multiple values based on the field type
+                new_initial = {}
+                for name in orig_initial:
+                    var = vars.get(name)
+                    if var and var.multiple_values:
+                        # Force MultiValueDict to give us the list of values
+                        new_initial[name] = orig_initial.getlist(name)
+                    else:
+                        # By default MultiValueDict gives us the last value
+                        new_initial[name] = orig_initial.get(name)
+
+                kwargs['initial'] = new_initial
 
         super().__init__(*args, **kwargs)
 

netbox/extras/scripts.py

@@ -50,6 +50,9 @@ class ScriptVariable:
     """
     form_field = forms.CharField
 
+    # Accept multiple values (e.g. MultiObjectVar)
+    multiple_values = False
+
     def __init__(self, label='', description='', default=None, required=True, widget=None):
 
         # Initialize field attributes
@@ -187,6 +190,7 @@ class MultiObjectVar(ScriptVariable):
     Like ObjectVar, but can represent one or more objects.
     """
     form_field = DynamicModelMultipleChoiceField
+    multiple_values = True
 
     def __init__(self, queryset, *args, **kwargs):
         super().__init__(*args, **kwargs)

netbox/extras/tests/test_scripts.py

@@ -1,5 +1,7 @@
 from django.core.files.uploadedfile import SimpleUploadedFile
+from django.http import QueryDict
 from django.test import TestCase
+from django.utils.datastructures import MultiValueDict
 from netaddr import IPAddress, IPNetwork
 
 from dcim.models import DeviceRole
@@ -145,6 +147,33 @@ class TestScript(Script):
         self.assertTrue(form.is_valid())
         self.assertEqual(form.cleaned_data['var1'].pk, data['var1'])
 
+    def test_objectvar_initial(self):
+
+        class TestScript(Script):
+
+            var1 = ObjectVar(
+                queryset=DeviceRole.objects.all()
+            )
+
+        # Populate some objects
+        for i in range(1, 6):
+            DeviceRole(
+                name='Device Role {}'.format(i),
+                slug='device-role-{}'.format(i)
+            ).save()
+
+        # Validate valid initial data
+        keys = list(DeviceRole.objects.values_list('pk', flat=True)[:2])
+
+        initial = QueryDict(f"var1={keys[0]}&var1={keys[1]}")
+        form = TestScript().as_form(initial=initial)
+
+        # For a single-value field the initial value must be the LAST value in the query-string
+        self.assertIn(
+            form.get_initial_for_field(form.fields['var1'], 'var1'),
+            [keys[1], str(keys[1])]  # Accept both numeric and string values
+        )
+
     def test_multiobjectvar(self):
 
         class TestScript(Script):
@@ -168,6 +197,37 @@ class TestScript(Script):
         self.assertEqual(form.cleaned_data['var1'][1].pk, data['var1'][1])
         self.assertEqual(form.cleaned_data['var1'][2].pk, data['var1'][2])
 
+    def test_multiobjectvar_initial(self):
+
+        class TestScript(Script):
+
+            var1 = MultiObjectVar(
+                queryset=DeviceRole.objects.all()
+            )
+
+        # Populate some objects
+        for i in range(1, 6):
+            DeviceRole(
+                name='Device Role {}'.format(i),
+                slug='device-role-{}'.format(i)
+            ).save()
+
+        # Validate valid initial data
+        keys = list(DeviceRole.objects.values_list('pk', flat=True)[:2])
+
+        initial = QueryDict(f"var1={keys[0]}&var1={keys[1]}")
+        form = TestScript().as_form(initial=initial)
+
+        self.assertEqual(len(form.get_initial_for_field(form.fields['var1'], 'var1')), 2)
+        self.assertIn(
+            form.get_initial_for_field(form.fields['var1'], 'var1')[0],
+            [keys[0], str(keys[0])]  # Accept both numeric and string values
+        )
+        self.assertIn(
+            form.get_initial_for_field(form.fields['var1'], 'var1')[1],
+            [keys[1], str(keys[1])]  # Accept both numeric and string values
+        )
+
     def test_filevar(self):
 
         class TestScript(Script):