Custom iterable object for PluginMenuItem groups (#97)

Author: bctiemann

netbox_custom_objects/navigation.py

@@ -1,6 +1,6 @@
 from django.apps import apps
 from django.conf import settings
-from django.urls import reverse
+from django.urls import reverse_lazy
 from django.utils.translation import gettext_lazy as _
 from netbox.plugins import PluginMenu, PluginMenuButton, PluginMenuItem
 from packaging import version
@@ -21,50 +21,43 @@
 )
 
 
-def get_menu():
-    CustomObjectType = apps.get_model(APP_LABEL, "CustomObjectType")
-    menu_items = []
-    for custom_object_type in CustomObjectType.objects.all():
-        model = custom_object_type.get_model()
-        menu_items.append(
-            PluginMenuItem(
+class CustomObjectTypeMenuItems:
+
+    def __iter__(self):
+        CustomObjectType = apps.get_model(APP_LABEL, "CustomObjectType")
+        for custom_object_type in CustomObjectType.objects.all():
+            model = custom_object_type.get_model()
+            add_button = PluginMenuButton(
+                None,
+                _("Add"),
+                "mdi mdi-plus-thick",
+            )
+            add_button.url = reverse_lazy(
+                f"plugins:{APP_LABEL}:customobject_add",
+                kwargs={
+                    "custom_object_type": custom_object_type.name.lower()
+                },
+            )
+            menu_item = PluginMenuItem(
                 link=None,
-                url=reverse(
-                    f"plugins:{APP_LABEL}:customobject_list",
-                    kwargs={"custom_object_type": custom_object_type.name.lower()},
-                ),
                 link_text=_(title(model._meta.verbose_name_plural)),
-                buttons=(
-                    PluginMenuButton(
-                        None,
-                        _("Add"),
-                        "mdi mdi-plus-thick",
-                        url=reverse(
-                            f"plugins:{APP_LABEL}:customobject_add",
-                            kwargs={
-                                "custom_object_type": custom_object_type.name.lower()
-                            },
-                        ),
-                    ),
-                ),
+                buttons=(add_button,),
             )
-        )
-    return PluginMenu(
-        label="Custom Objects",
-        groups=(
-            (_("Object Types"), (custom_object_type_plugin_menu_item,)),
-            (_("Objects"), tuple(menu_items)),
-        ),
-        icon_class="mdi mdi-toy-brick-outline",
-    )
+            menu_item.url = reverse_lazy(
+                f"plugins:{APP_LABEL}:customobject_list",
+                kwargs={"custom_object_type": custom_object_type.name.lower()},
+            )
+            yield menu_item
 
 
 current_version = version.parse(settings.RELEASE.version)
-if current_version < version.parse("4.4.0"):
-    menu = PluginMenu(
-        label="Custom Objects",
-        groups=((_("Object Types"), (custom_object_type_plugin_menu_item,)),),
-        icon_class="mdi mdi-toy-brick-outline",
-    )
-else:
-    menu = get_menu
+
+groups = [(_("Object Types"), (custom_object_type_plugin_menu_item,))]
+if current_version >= version.parse("4.3.4"):
+    groups.append((_("Objects"), CustomObjectTypeMenuItems()))
+
+menu = PluginMenu(
+    label=_("Custom Objects"),
+    groups=tuple(groups),
+    icon_class="mdi mdi-toy-brick-outline",
+)