From 9abc67bbebc00eb078e34d5aa7a4453d3888dd51 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Wed, 10 Jun 2020 11:38:23 -0400
Subject: [PATCH] Fixes #4737: Introduce ColoredLabelColumn for consistent
 display of colored labels

---
 docs/release-notes/version-2.8.md |  1 +
 netbox/dcim/tables.py             | 23 +++--------------------
 netbox/utilities/tables.py        | 17 +++++++++++++++++
 netbox/virtualization/tables.py   | 10 ++--------
 4 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/docs/release-notes/version-2.8.md b/docs/release-notes/version-2.8.md
index 348a6fa36..219f94e23 100644
--- a/docs/release-notes/version-2.8.md
+++ b/docs/release-notes/version-2.8.md
@@ -13,6 +13,7 @@
 * [#4707](https://github.com/netbox-community/netbox/issues/4707) - Fix `prefix_count` population on VLAN API serializer
 * [#4725](https://github.com/netbox-community/netbox/issues/4725) - Fix "brief" rendering of various REST API endpoints
 * [#4736](https://github.com/netbox-community/netbox/issues/4736) - Add cable trace endpoints for pass-through ports
+* [#4737](https://github.com/netbox-community/netbox/issues/4737) - Fix display of role labels in virtual machines table
 
 ---
 
diff --git a/netbox/dcim/tables.py b/netbox/dcim/tables.py
index 9018625a0..d8cf41eaa 100644
--- a/netbox/dcim/tables.py
+++ b/netbox/dcim/tables.py
@@ -2,7 +2,7 @@ import django_tables2 as tables
 from django_tables2.utils import Accessor
 
 from tenancy.tables import COL_TENANT
-from utilities.tables import BaseTable, BooleanColumn, ColorColumn, TagColumn, ToggleColumn
+from utilities.tables import BaseTable, BooleanColumn, ColorColumn, ColoredLabelColumn, TagColumn, ToggleColumn
 from .models import (
     Cable, ConsolePort, ConsolePortTemplate, ConsoleServerPort, ConsoleServerPortTemplate, Device, DeviceBay,
     DeviceBayTemplate, DeviceRole, DeviceType, FrontPort, FrontPortTemplate, Interface, InterfaceTemplate,
@@ -72,15 +72,6 @@ RACKROLE_ACTIONS = """
 {% endif %}
 """
 
-RACK_ROLE = """
-{% if record.role %}
-    {% load helpers %}
-    <label class="label" style="color: {{ record.role.color|fgcolor }}; background-color: #{{ record.role.color }}">{{ value }}</label>
-{% else %}
-    &mdash;
-{% endif %}
-"""
-
 RACK_DEVICE_COUNT = """
 <a href="{% url 'dcim:device_list' %}?rack_id={{ record.pk }}">{{ value }}</a>
 """
@@ -137,11 +128,6 @@ PLATFORM_ACTIONS = """
 {% endif %}
 """
 
-DEVICE_ROLE = """
-{% load helpers %}
-<label class="label" style="color: {{ record.device_role.color|fgcolor }}; background-color: #{{ record.device_role.color }}">{{ value }}</label>
-"""
-
 STATUS_LABEL = """
 <span class="label label-{{ record.get_status_class }}">{{ record.get_status_display }}</span>
 """
@@ -325,9 +311,7 @@ class RackTable(BaseTable):
     status = tables.TemplateColumn(
         template_code=STATUS_LABEL
     )
-    role = tables.TemplateColumn(
-        template_code=RACK_ROLE
-    )
+    role = ColoredLabelColumn()
     u_height = tables.TemplateColumn(
         template_code="{{ record.u_height }}U",
         verbose_name='Height'
@@ -806,8 +790,7 @@ class DeviceTable(BaseTable):
         viewname='dcim:rack',
         args=[Accessor('rack.pk')]
     )
-    device_role = tables.TemplateColumn(
-        template_code=DEVICE_ROLE,
+    device_role = ColoredLabelColumn(
         verbose_name='Role'
     )
     device_type = tables.LinkColumn(
diff --git a/netbox/utilities/tables.py b/netbox/utilities/tables.py
index 97108b5b2..10e408b43 100644
--- a/netbox/utilities/tables.py
+++ b/netbox/utilities/tables.py
@@ -84,6 +84,10 @@ class BaseTable(tables.Table):
         return [name for name in self.sequence if self.columns[name].visible]
 
 
+#
+# Table columns
+#
+
 class ToggleColumn(tables.CheckBoxColumn):
     """
     Extend CheckBoxColumn to add a "toggle all" checkbox in the column header.
@@ -129,6 +133,19 @@ class ColorColumn(tables.Column):
         )
 
 
+class ColoredLabelColumn(tables.TemplateColumn):
+    """
+    Render a colored label (e.g. for DeviceRoles).
+    """
+    template_code = """
+    {% load helpers %}
+    {% if value %}<label class="label" style="color: {{ value.color|fgcolor }}; background-color: #{{ value.color }}">{{ value }}</label>{% else %}&mdash;{% endif %}
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(template_code=self.template_code, *args, **kwargs)
+
+
 class TagColumn(tables.TemplateColumn):
     """
     Display a list of tags assigned to the object.
diff --git a/netbox/virtualization/tables.py b/netbox/virtualization/tables.py
index 077add945..d957e0053 100644
--- a/netbox/virtualization/tables.py
+++ b/netbox/virtualization/tables.py
@@ -3,7 +3,7 @@ from django_tables2.utils import Accessor
 
 from dcim.models import Interface
 from tenancy.tables import COL_TENANT
-from utilities.tables import BaseTable, TagColumn, ToggleColumn
+from utilities.tables import BaseTable, ColoredLabelColumn, TagColumn, ToggleColumn
 from .models import Cluster, ClusterGroup, ClusterType, VirtualMachine
 
 CLUSTERTYPE_ACTIONS = """
@@ -28,10 +28,6 @@ VIRTUALMACHINE_STATUS = """
 <span class="label label-{{ record.get_status_class }}">{{ record.get_status_display }}</span>
 """
 
-VIRTUALMACHINE_ROLE = """
-{% if record.role %}<label class="label" style="background-color: #{{ record.role.color }}">{{ value }}</label>{% else %}&mdash;{% endif %}
-"""
-
 VIRTUALMACHINE_PRIMARY_IP = """
 {{ record.primary_ip6.address.ip|default:"" }}
 {% if record.primary_ip6 and record.primary_ip4 %}<br />{% endif %}
@@ -132,9 +128,7 @@ class VirtualMachineTable(BaseTable):
         viewname='virtualization:cluster',
         args=[Accessor('cluster.pk')]
     )
-    role = tables.TemplateColumn(
-        template_code=VIRTUALMACHINE_ROLE
-    )
+    role = ColoredLabelColumn()
     tenant = tables.TemplateColumn(
         template_code=COL_TENANT
     )