Separate method for _build_prefetch to allow overrides

vintasoftware · May 31, 2023 · fce78df · fce78df
1 parent bb7a526
commit fce78df
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/django_virtual_models/fields.py b/django_virtual_models/fields.py
@@ -289,6 +289,12 @@ def _hydrate_queryset_with_nested_declared_fields(
 
         return new_qs
 
+    def _build_prefetch(self, prefetch_queryset: QuerySet):
+        if self.lookup:
+            return Prefetch(self.lookup, queryset=prefetch_queryset, to_attr=self.to_attr)
+        else:
+            return Prefetch(self.field_name, queryset=prefetch_queryset)
+
     def hydrate_queryset(
         self,
         qs: QuerySet,
@@ -325,12 +331,8 @@ def hydrate_queryset(
             deferred_fields=self.deferred_fields,
         )
 
-        # build prefetch with lookups
-        if self.lookup:
-            prefetch = Prefetch(self.lookup, queryset=prefetch_queryset, to_attr=self.to_attr)
-        else:
-            prefetch = Prefetch(self.field_name, queryset=prefetch_queryset)
-
+        # build prefetch object, call prefetch_related
+        prefetch = self._build_prefetch(prefetch_queryset)
         new_qs = qs.prefetch_related(prefetch)
         return new_qs