From 911184375250d3d8e6d1fb0afd3ec6333d608469 Mon Sep 17 00:00:00 2001 From: Travis Swicegood Date: Mon, 7 Oct 2013 18:15:41 -0500 Subject: [PATCH 1/2] Refactor to make sure get_subclass() is on QuerySet There's an edge case where you might want to call get_subclass() on a QuerySet the same way you can call get() after you have already called various filter/exclude methods. --- model_utils/managers.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/model_utils/managers.py b/model_utils/managers.py index abb18b4..3c72585 100644 --- a/model_utils/managers.py +++ b/model_utils/managers.py @@ -96,6 +96,8 @@ class InheritanceQuerySet(QuerySet): else: return node + def get_subclass(self, *args, **kwargs): + return self.select_subclasses().get(*args, **kwargs) class InheritanceManager(models.Manager): @@ -108,8 +110,7 @@ class InheritanceManager(models.Manager): return self.get_query_set().select_subclasses(*subclasses) def get_subclass(self, *args, **kwargs): - return self.get_query_set().select_subclasses().get(*args, **kwargs) - + return self.get_query_set().get_subclass(*args, **kwargs) class QueryManager(models.Manager): From 03c8a1929b8c33bf690c4ead13f658805420a837 Mon Sep 17 00:00:00 2001 From: Travis Swicegood Date: Tue, 8 Oct 2013 10:38:52 -0500 Subject: [PATCH 2/2] Add test to verify get_subclass() on QuerySet --- model_utils/tests/tests.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/model_utils/tests/tests.py b/model_utils/tests/tests.py index 8381815..e8df6a2 100644 --- a/model_utils/tests/tests.py +++ b/model_utils/tests/tests.py @@ -574,6 +574,12 @@ class InheritanceManagerTests(TestCase): self.child1) + def test_get_subclass_on_queryset(self): + self.assertEqual( + self.get_manager().all().get_subclass(pk=self.child1.pk), + self.child1) + + def test_prior_select_related(self): with self.assertNumQueries(1): obj = self.get_manager().select_related(