Merge pull request #76 from nautobot/u/snaselj-fix-viewset

Fix UIViewSet sub class name

pylint_nautobot/utils.py

@@ -116,7 +116,7 @@ def get_model_name(node: ClassDef) -> str:
     """Get the model name from the class definition."""
     queryset = find_attr(node, "queryset")
     if queryset:
-        return find_model_name_from_queryset(queryset.value)
+        return get_model_name_from_queryset(queryset.value)
 
     model_attr = find_attr(node, "model")
     if not model_attr:
@@ -130,14 +130,19 @@ def get_model_name(node: ClassDef) -> str:
     return get_model_name_from_attr(model_attr)
 
 
-def find_model_name_from_queryset(node: NodeNG) -> str:
+def get_model_name_from_queryset(node: NodeNG) -> str:
     """Get the model name from the queryset assignment value."""
     while node:
         if isinstance(node, Call):
             node = node.func
         elif isinstance(node, Attribute):
-            if node.attrname == "objects" and isinstance(node.expr, Name):
-                return node.expr.name
+            if node.attrname == "objects":
+                if isinstance(node.expr, Name):
+                    # Covers `queryset = AddressObject.objects.all()`
+                    return node.expr.name
+                if isinstance(node.expr, Attribute):
+                    # Covers `queryset = models.AddressObject.objects.all()`
+                    return node.expr.attrname
             node = node.expr
         else:
             break

tests/inputs/sub-class-name/error_table_model.py

@@ -0,0 +1,12 @@
+from nautobot.apps.tables import BaseTable
+
+from . import models
+
+
+class MyAddressObjectTable(BaseTable):
+    """Filter for AddressObject."""
+
+    class Meta:
+        """Meta attributes for filter."""
+
+        model = models.AddressObject

tests/inputs/sub-class-name/error_viewset_model.py

@@ -0,0 +1,9 @@
+from nautobot.apps.views import NautobotUIViewSet
+
+from . import models
+
+
+class MyAddressObjectUIViewSet(NautobotUIViewSet):
+    """Filter for AddressObject."""
+
+    queryset = models.AddressObject.objects.all()

tests/inputs/sub-class-name/good_table_model.py

@@ -0,0 +1,12 @@
+from nautobot.apps.tables import BaseTable
+
+from . import models
+
+
+class AddressObjectTable(BaseTable):
+    """Filter for AddressObject."""
+
+    class Meta:
+        """Meta attributes for filter."""
+
+        model = models.AddressObject

tests/inputs/sub-class-name/good_viewset_model.py

@@ -0,0 +1,9 @@
+from nautobot.apps.views import NautobotUIViewSet
+
+from . import models
+
+
+class AddressObjectUIViewSet(NautobotUIViewSet):
+    """Filter for AddressObject."""
+
+    queryset = models.AddressObject.objects.all()

tests/inputs/sub-class-name/models.py

@@ -0,0 +1,5 @@
+from nautobot.core.models.generics import PrimaryModel
+
+
+class AddressObject(PrimaryModel):
+    pass

tests/test_sub_class_name.py

@@ -14,6 +14,10 @@ def _find_failing_node(module_node):
     return module_node.body[3]
 
 
+def _find_failing_node_model(module_node):
+    return module_node.body[2]
+
+
 _EXPECTED_ERRORS = {
     "filter_form": {
         "versions": ">=2",
@@ -55,6 +59,14 @@ def _find_failing_node(module_node):
         "args": ("AddressObjectUIViewSet",),
         "node": _find_failing_node,
     },
+    "viewset_model": {
+        "versions": ">=2",
+        "msg_id": "nb-sub-class-name",
+        "line": 6,
+        "col_offset": 0,
+        "args": ("AddressObjectUIViewSet",),
+        "node": _find_failing_node_model,
+    },
     "table": {
         "versions": ">=2",
         "msg_id": "nb-sub-class-name",
@@ -63,6 +75,14 @@ def _find_failing_node(module_node):
         "args": ("AddressObjectTable",),
         "node": _find_failing_node,
     },
+    "table_model": {
+        "versions": ">=2",
+        "msg_id": "nb-sub-class-name",
+        "line": 6,
+        "col_offset": 0,
+        "args": ("AddressObjectTable",),
+        "node": _find_failing_node_model,
+    },
 }