diff --git a/django_pandas/tests/models.py b/django_pandas/tests/models.py index 6a3496a..3bd96f3 100644 --- a/django_pandas/tests/models.py +++ b/django_pandas/tests/models.py @@ -1,6 +1,6 @@ from django.db import models from django.utils.encoding import python_2_unicode_compatible -from django_pandas.managers import DataFrameManager +from django_pandas.managers import DataFrameManager, PassThroughManager @python_2_unicode_compatible @@ -158,3 +158,71 @@ class Portfolio(models.Model): def __str__(self): return self.name + + +class DudeQuerySet(models.query.QuerySet): + def abiding(self): + return self.filter(abides=True) + + def rug_positive(self): + return self.filter(has_rug=True) + + def rug_negative(self): + return self.filter(has_rug=False) + + def by_name(self, name): + return self.filter(name__iexact=name) + + +class AbidingManager(PassThroughManager): + def get_queryset(self): + return DudeQuerySet(self.model).abiding() + + get_query_set = get_queryset + + def get_stats(self): + return { + "abiding_count": self.count(), + "rug_count": self.rug_positive().count(), + } + + +class Dude(models.Model): + abides = models.BooleanField(default=True) + name = models.CharField(max_length=20) + has_rug = models.BooleanField(default=False) + + objects = PassThroughManager(DudeQuerySet) + abiders = AbidingManager() + + +class Car(models.Model): + name = models.CharField(max_length=20) + owner = models.ForeignKey(Dude, related_name='cars_owned') + + objects = PassThroughManager(DudeQuerySet) + + +class SpotManager(PassThroughManager): + def get_queryset(self): + return super(SpotManager, self).get_queryset().filter(secret=False) + + get_query_set = get_queryset + + +class SpotQuerySet(models.query.QuerySet): + def closed(self): + return self.filter(closed=True) + + def secured(self): + return self.filter(secure=True) + + +class Spot(models.Model): + name = models.CharField(max_length=20) + secure = models.BooleanField(default=True) + closed = models.BooleanField(default=False) + secret = models.BooleanField(default=False) + owner = models.ForeignKey(Dude, related_name='spots_owned') + + objects = SpotManager.for_queryset_class(SpotQuerySet)() diff --git a/django_pandas/tests/test_manager.py b/django_pandas/tests/test_manager.py index af76069..a39d117 100644 --- a/django_pandas/tests/test_manager.py +++ b/django_pandas/tests/test_manager.py @@ -1,9 +1,11 @@ from django.test import TestCase import pandas as pd import numpy as np +import pickle +import django from .models import ( DataFrame, WideTimeSeries, - LongTimeSeries, PivotData, MyModelChoice + LongTimeSeries, PivotData, MyModelChoice, Dude, Car, Spot ) import pandas.util.testing as tm @@ -183,3 +185,114 @@ def test_pivot(self): self.assertEqual(pt.index.names, rows) self.assertEqual(pt.columns.names, cols) + +if django.VERSION < (1, 9): + + class PassThroughManagerTests(TestCase): + + def setUp(self): + Dude.objects.create(name='The Dude', abides=True, has_rug=False) + Dude.objects.create(name='His Dudeness', abides=False, has_rug=True) + Dude.objects.create(name='Duder', abides=False, has_rug=False) + Dude.objects.create(name='El Duderino', abides=True, has_rug=True) + + + def test_chaining(self): + self.assertEqual(Dude.objects.by_name('Duder').count(), 1) + self.assertEqual(Dude.objects.all().by_name('Duder').count(), 1) + self.assertEqual(Dude.abiders.rug_positive().count(), 1) + self.assertEqual(Dude.abiders.all().rug_positive().count(), 1) + + + def test_manager_only_methods(self): + stats = Dude.abiders.get_stats() + self.assertEqual(stats['rug_count'], 1) + with self.assertRaises(AttributeError): + Dude.abiders.all().get_stats() + + + def test_queryset_pickling(self): + qs = Dude.objects.all() + saltyqs = pickle.dumps(qs) + unqs = pickle.loads(saltyqs) + self.assertEqual(unqs.by_name('The Dude').count(), 1) + + + def test_queryset_not_available_on_related_manager(self): + dude = Dude.objects.by_name('Duder').get() + Car.objects.create(name='Ford', owner=dude) + self.assertFalse(hasattr(dude.cars_owned, 'by_name')) + + + def test_using_dir(self): + # make sure introspecing via dir() doesn't actually cause queries, + # just as a sanity check. + with self.assertNumQueries(0): + querysets_to_dir = ( + Dude.objects, + Dude.objects.by_name('Duder'), + Dude.objects.all().by_name('Duder'), + Dude.abiders, + Dude.abiders.rug_positive(), + Dude.abiders.all().rug_positive() + ) + for qs in querysets_to_dir: + self.assertTrue('by_name' in dir(qs)) + self.assertTrue('abiding' in dir(qs)) + self.assertTrue('rug_positive' in dir(qs)) + self.assertTrue('rug_negative' in dir(qs)) + # some standard qs methods + self.assertTrue('count' in dir(qs)) + self.assertTrue('order_by' in dir(qs)) + self.assertTrue('select_related' in dir(qs)) + # make sure it's been de-duplicated + self.assertEqual(1, dir(qs).count('distinct')) + + # manager only method. + self.assertTrue('get_stats' in dir(Dude.abiders)) + # manager only method shouldn't appear on the non AbidingManager + self.assertFalse('get_stats' in dir(Dude.objects)) + # standard manager methods + self.assertTrue('get_query_set' in dir(Dude.abiders)) + self.assertTrue('contribute_to_class' in dir(Dude.abiders)) + + + + class CreatePassThroughManagerTests(TestCase): + + def setUp(self): + self.dude = Dude.objects.create(name='El Duderino') + self.other_dude = Dude.objects.create(name='Das Dude') + + def test_reverse_manager(self): + Spot.objects.create( + name='The Crib', owner=self.dude, closed=True, secure=True, + secret=False) + self.assertEqual(self.dude.spots_owned.closed().count(), 1) + Spot.objects.create( + name='The Crux', owner=self.other_dude, closed=True, secure=True, + secret=False + ) + self.assertEqual(self.dude.spots_owned.closed().all().count(), 1) + self.assertEqual(self.dude.spots_owned.closed().count(), 1) + + def test_related_queryset_pickling(self): + Spot.objects.create( + name='The Crib', owner=self.dude, closed=True, secure=True, + secret=False) + qs = self.dude.spots_owned.closed() + pickled_qs = pickle.dumps(qs) + unpickled_qs = pickle.loads(pickled_qs) + self.assertEqual(unpickled_qs.secured().count(), 1) + + def test_related_queryset_superclass_method(self): + Spot.objects.create( + name='The Crib', owner=self.dude, closed=True, secure=True, + secret=False) + Spot.objects.create( + name='The Secret Crib', owner=self.dude, closed=False, secure=True, + secret=True) + self.assertEqual(self.dude.spots_owned.count(), 1) + + def test_related_manager_create(self): + self.dude.spots_owned.create(name='The Crib', closed=True, secure=True)