From b0c0ad7c8277fa6e55abad584084f87a6655f096 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Wed, 5 Feb 2020 15:29:35 -0500
Subject: [PATCH] Adapt component creation forms to infer parent device/VM from
 initial or bound data

---
 netbox/dcim/forms.py           | 28 ++++++++++++++++------------
 netbox/virtualization/forms.py |  8 +++++---
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/netbox/dcim/forms.py b/netbox/dcim/forms.py
index 2a6aaf451..b01fbe8f0 100644
--- a/netbox/dcim/forms.py
+++ b/netbox/dcim/forms.py
@@ -2456,9 +2456,10 @@ class PowerOutletCreateForm(ComponentForm):
         super().__init__(*args, **kwargs)
 
         # Limit power_port queryset to PowerPorts which belong to the parent Device
-        if 'device' in self.initial:
-            device = Device.objects.filter(pk=self.initial['device']).first()
-            self.fields['power_port'].queryset = PowerPort.objects.filter(device=device)
+        device = Device.objects.get(
+            pk=self.initial.get('device') or self.data.get('device')
+        )
+        self.fields['power_port'].queryset = PowerPort.objects.filter(device=device)
 
 
 class PowerOutletCSVForm(forms.ModelForm):
@@ -2707,12 +2708,13 @@ class InterfaceCreateForm(InterfaceCommonForm, ComponentForm, forms.Form):
         super().__init__(*args, **kwargs)
 
         # Limit LAG choices to interfaces which belong to the parent device (or VC master)
-        if 'device' in self.initial:
-            device = Device.objects.filter(pk=self.initial['device']).first()
-            self.fields['lag'].queryset = Interface.objects.filter(
-                device__in=[device, device.get_vc_master()],
-                type=InterfaceTypeChoices.TYPE_LAG
-            )
+        device = Device.objects.get(
+            pk=self.initial.get('device') or self.data.get('device')
+        )
+        self.fields['lag'].queryset = Interface.objects.filter(
+            device__in=[device, device.get_vc_master()],
+            type=InterfaceTypeChoices.TYPE_LAG
+        )
 
 
 class InterfaceCSVForm(forms.ModelForm):
@@ -2952,18 +2954,20 @@ class FrontPortCreateForm(ComponentForm):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        parent = Device.objects.get(pk=self.initial.get('device'))
+        device = Device.objects.get(
+            pk=self.initial.get('device') or self.data.get('device')
+        )
 
         # Determine which rear port positions are occupied. These will be excluded from the list of available
         # mappings.
         occupied_port_positions = [
             (front_port.rear_port_id, front_port.rear_port_position)
-            for front_port in parent.frontports.all()
+            for front_port in device.frontports.all()
         ]
 
         # Populate rear port choices
         choices = []
-        rear_ports = RearPort.objects.filter(device=parent)
+        rear_ports = RearPort.objects.filter(device=device)
         for rear_port in rear_ports:
             for i in range(1, rear_port.positions + 1):
                 if (rear_port.pk, i) not in occupied_port_positions:
diff --git a/netbox/virtualization/forms.py b/netbox/virtualization/forms.py
index 718784d3e..98c2dd752 100644
--- a/netbox/virtualization/forms.py
+++ b/netbox/virtualization/forms.py
@@ -797,9 +797,12 @@ class InterfaceCreateForm(ComponentForm):
     )
 
     def __init__(self, *args, **kwargs):
-
         super().__init__(*args, **kwargs)
 
+        virtual_machine = VirtualMachine.objects.get(
+            pk=self.initial.get('virtual_machine') or self.data.get('virtual_machine')
+        )
+
         # Limit VLAN choices to those in: global vlans, global groups, the current site's group, the current site
         vlan_choices = []
         global_vlans = VLAN.objects.filter(site=None, group=None)
@@ -812,8 +815,7 @@ class InterfaceCreateForm(ComponentForm):
                 (group.name, [(vlan.pk, vlan) for vlan in global_group_vlans])
             )
 
-        parent = VirtualMachine.objects.get(pk=self.initial.get('virtual_machine'))
-        site = getattr(parent.cluster, 'site', None)
+        site = getattr(virtual_machine.cluster, 'site', None)
         if site is not None:
 
             # Add non-grouped site VLANs