Closes #3545: Add MultiObjectVar for custom scripts

CHANGELOG.md

@@ -12,6 +12,7 @@ v2.6.6 (FUTURE)
 * [#1941](https://github.com/netbox-community/netbox/issues/1941) - Add InfiniBand interface types
 * [#3259](https://github.com/netbox-community/netbox/issues/3259) - Add `rack` and `site` filters for cables
 * [#3471](https://github.com/netbox-community/netbox/issues/3471) - Disallow raw HTML in Markdown-rendered fields
+* [#3545](https://github.com/netbox-community/netbox/issues/3545) - Add `MultiObjectVar` for custom scripts
 * [#3563](https://github.com/netbox-community/netbox/issues/3563) - Enable editing of individual DeviceType components
 * [#3580](https://github.com/netbox-community/netbox/issues/3580) - Render text and URL fields as textareas in the custom link form
 

netbox/extras/scripts.py

@@ -11,7 +11,7 @@ from django import forms
 from django.conf import settings
 from django.core.validators import RegexValidator
 from django.db import transaction
-from mptt.forms import TreeNodeChoiceField
+from mptt.forms import TreeNodeChoiceField, TreeNodeMultipleChoiceField
 from mptt.models import MPTTModel
 
 from ipam.formfields import IPFormField
@@ -27,6 +27,7 @@ __all__ = [
     'FileVar',
     'IntegerVar',
     'IPNetworkVar',
+    'MultiObjectVar',
     'ObjectVar',
     'Script',
     'StringVar',
@@ -149,6 +150,23 @@ class ObjectVar(ScriptVariable):
             self.form_field = TreeNodeChoiceField
 
 
+class MultiObjectVar(ScriptVariable):
+    """
+    Like ObjectVar, but can represent one or more objects.
+    """
+    form_field = forms.ModelMultipleChoiceField
+
+    def __init__(self, queryset, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Queryset for field choices
+        self.field_attrs['queryset'] = queryset
+
+        # Update form field for MPTT (nested) objects
+        if issubclass(queryset.model, MPTTModel):
+            self.form_field = TreeNodeMultipleChoiceField
+
+
 class FileVar(ScriptVariable):
     """
     An uploaded file.

netbox/extras/tests/test_scripts.py

@@ -120,6 +120,29 @@ class ScriptVariablesTest(TestCase):
         self.assertTrue(form.is_valid())
         self.assertEqual(form.cleaned_data['var1'].pk, data['var1'])
 
+    def test_multiobjectvar(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 data
+        data = {'var1': [role.pk for role in DeviceRole.objects.all()[:3]]}
+        form = TestScript().as_form(data, None)
+        self.assertTrue(form.is_valid())
+        self.assertEqual(form.cleaned_data['var1'][0].pk, data['var1'][0])
+        self.assertEqual(form.cleaned_data['var1'][1].pk, data['var1'][1])
+        self.assertEqual(form.cleaned_data['var1'][2].pk, data['var1'][2])
+
     def test_filevar(self):
 
         class TestScript(Script):