Clean up RestrictedGenericForeignKey and RestrictedPrefetch

netbox/utilities/fields.py

@@ -78,10 +78,18 @@ class NaturalOrderingField(models.CharField):
 
 class RestrictedGenericForeignKey(GenericForeignKey):
 
-    # Replicated from GenericForeignKey
+    # Replicated largely from GenericForeignKey. Changes include:
+    #  1. Capture restrict_params from RestrictedPrefetch (hack)
+    #  2. If restrict_params is set, call restrict() on the queryset for
+    #     the related model
     def get_prefetch_queryset(self, instances, queryset=None):
+        restrict_params = {}
+
         # Compensate for the hack in RestrictedPrefetch
-        restrict_params = queryset if type(queryset) is dict else {}
+        if type(queryset) is dict:
+            restrict_params = queryset
+        elif queryset is not None:
+            raise ValueError("Custom queryset can't be used for this lookup.")
 
         # For efficiency, group the instances by content type and then do one
         # query per model

netbox/utilities/querysets.py

@@ -21,13 +21,12 @@ class RestrictedPrefetch(Prefetch):
             'action': self.restrict_action,
         }
 
-        qs = super().get_current_queryset(level)
-        if qs:
-            return qs.filter(**params)
+        if qs := super().get_current_queryset(level):
+            return qs.restrict(**params)
 
         # Bit of a hack. If no queryset is defined, pass through the dict of restrict()
-        # kwargs to be called later. This is necessary e.g. for GenericForeignKey fields,
-        # which do not permit setting a queryset on a Prefetch object.
+        # kwargs to be handled by the field. This is necessary e.g. for GenericForeignKey
+        # fields, which do not permit setting a queryset on a Prefetch object.
         return params