diff --git a/pylint_nautobot/utils.py b/pylint_nautobot/utils.py index be5e816..655d060 100644 --- a/pylint_nautobot/utils.py +++ b/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,17 @@ 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): + return node.expr.name + if isinstance(node.expr, Attribute): + return node.expr.attrname node = node.expr else: break diff --git a/tests/inputs/sub-class-name/__init__.py b/tests/inputs/sub-class-name/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/inputs/sub-class-name/error_table_model.py b/tests/inputs/sub-class-name/error_table_model.py new file mode 100644 index 0000000..2e190d5 --- /dev/null +++ b/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 diff --git a/tests/inputs/sub-class-name/error_viewset_model.py b/tests/inputs/sub-class-name/error_viewset_model.py new file mode 100644 index 0000000..8fb25d4 --- /dev/null +++ b/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() diff --git a/tests/inputs/sub-class-name/good_table_model.py b/tests/inputs/sub-class-name/good_table_model.py new file mode 100644 index 0000000..0d86809 --- /dev/null +++ b/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 diff --git a/tests/inputs/sub-class-name/good_viewset_model.py b/tests/inputs/sub-class-name/good_viewset_model.py new file mode 100644 index 0000000..2028f5b --- /dev/null +++ b/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() diff --git a/tests/inputs/sub-class-name/models.py b/tests/inputs/sub-class-name/models.py new file mode 100644 index 0000000..5837b35 --- /dev/null +++ b/tests/inputs/sub-class-name/models.py @@ -0,0 +1,5 @@ +from nautobot.core.models.generics import PrimaryModel + + +class AddressObject(PrimaryModel): + pass diff --git a/tests/test_sub_class_name.py b/tests/test_sub_class_name.py index 3de290b..d067002 100644 --- a/tests/test_sub_class_name.py +++ b/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": 5, + "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": 5, + "col_offset": 0, + "args": ("AddressObjectTable",), + "node": _find_failing_node_model, + }, }