Added better support for relation fields

Co-authored-by: Baptiste Mispelon <[email protected]>

src/modelsubquery/functions.py

@@ -13,7 +13,7 @@ def _model_fields(model, fields):
     then return all the declared fields.
     """
     # TODO: the pk/id field should always be returned
-    declared = {f.name for f in model._meta.get_fields()}
+    declared = {f.column for f in model._meta.local_concrete_fields}
     if fields is None:
         return declared
 

testproject/testapp/models.py

@@ -11,6 +11,9 @@ class Book(models.Model):
     rating = models.IntegerField(blank=True, null=True)
     published = models.DateField(default=date.today)
     has_cover = models.BooleanField(default=True)
+    author = models.ForeignKey(
+        "Person", on_delete=models.CASCADE, related_name="books", null=True
+    )
 
 
 class PersonQuerySet(models.QuerySet):
@@ -29,3 +32,8 @@ class Person(models.Model):
     birth = models.DateField(default=date.today)
 
     objects = PersonQuerySet.as_manager()
+
+
+class Shelve(models.Model):
+    title = models.CharField(max_length=100)
+    books = models.ManyToManyField("Book", related_name="shelves")

testproject/testapp/tests.py

@@ -4,7 +4,21 @@
 from django.test import TestCase
 from model_bakery import baker
 
-from .models import Book, Person
+from modelsubquery.functions import _model_fields
+
+from .models import Book, Person, Shelve
+
+
+class ModelFieldTest(TestCase):
+    def test_model_fields_book(self):
+        fields = _model_fields(Book, None)
+        self.assertEqual(
+            fields, {"id", "title", "author_id", "published", "rating", "has_cover"}
+        )
+
+    def test_model_fields_shelve(self):
+        fields = _model_fields(Shelve, None)
+        self.assertEqual(fields, {"id", "title"})
 
 
 @time_machine.travel("2000-01-01")
@@ -79,3 +93,16 @@ def test_deffered_fields(self):
             self.assertEqual(person.book_of_year.title, "test")
         with self.assertNumQueries(1):
             self.assertEqual(person.book_of_year.rating, 5)
+
+    def test_with_related_field(self):
+        author = baker.make(Person)
+        book = baker.make(Book, author=author)
+
+        # Yes he wrote a book in his birthyear! ;-)
+
+        person = Person.objects.with_book_of_the_year().get()
+        self.assertEqual(person.book_of_year, book)
+
+        # The relation is there, but it's lazy!
+        with self.assertNumQueries(1):
+            self.assertEqual(person.book_of_year.author, author)