def test_manager_methods(self):
"""
This test ensures that the correct set of methods from `QuerySet`
are copied onto `Manager`.
It's particularly useful to prevent accidentally leaking new methods
into `Manager`. New `QuerySet` methods that should also be copied onto
`Manager` will need to be added to `ManagerTest.QUERYSET_PROXY_METHODS`.
"""
self.assertEqual(
sorted(BaseManager._get_queryset_methods(QuerySet).keys()), sorted(self.QUERYSET_PROXY_METHODS)
)
def test_manager_methods(self):
"""
This test ensures that the correct set of methods from `QuerySet`
are copied onto `Manager`.
It's particularly useful to prevent accidentally leaking new methods
into `Manager`. New `QuerySet` methods that should also be copied onto
`Manager` will need to be added to `ManagerTest.QUERYSET_PROXY_METHODS`.
"""
self.assertEqual(
sorted(BaseManager._get_queryset_methods(QuerySet).keys()),
sorted(self.QUERYSET_PROXY_METHODS),
)
class ConversationMessageManager(
BaseManager.from_queryset(ConversationMessageQuerySet)):
def create(self, **kwargs):
if 'thread' not in kwargs:
# make sure author is participant (to receive notifications)
conversation = kwargs.get('conversation')
author = kwargs.get('author')
conversation.conversationparticipant_set.get_or_create(user=author)
obj = super().create(**kwargs)
# clear cached value
if obj.thread and hasattr(obj.thread, '_replies_count'):
del obj.thread._replies_count
return obj
class NotificationManager(BaseManager.from_queryset(NotificationQuerySet)):
def create_for_pickup_collectors(self, collectors, type):
for collector in collectors:
pickup = collector.pickupdate
super().create(
user=collector.user,
type=type,
context={
'group': pickup.group.id,
'store': pickup.store.id,
'pickup': pickup.id,
'pickup_collector': collector.id,
}
)
class TransactionManager(BaseManager.from_queryset(TransactionQuerySet)):
""" Custom manager """
def get_for_token(self, tidb64, token):
transaction = None
try:
tid = urlsafe_base64_decode(tidb64).decode()
transaction = Transaction.objects.get(pk=tid)
if not default_token_generator.check_token(transaction, token):
transaction = None
except (TypeError, ValueError, OverflowError, ObjectDoesNotExist, ValidationError):
pass
return transaction
class ExchangeRate(models.Model):
source = models.PositiveSmallIntegerField(
choices=currency_exchange_choices,
help_text=_('Where the exchange data comes from'),
)
base_currency = models.CharField(max_length=3, db_index=True)
rates = models.JSONField(default=dict)
last_updated_at = models.DateTimeField(auto_now=True, db_index=True)
objects = BaseManager.from_queryset(ExchangeRateQuerySet)()
class Meta:
ordering = ['base_currency', '-last_updated_at']
def __repr__(self):
return 'Exchange rate for {0}'.format(self.base_currency)
class EntityManager(BaseManager.from_queryset(EavQuerySet)):
"""
Our custom manager, overrides ``models.Manager``.
"""
_queryset_class = EavQuerySet
def create(self, **kwargs):
"""
Parse eav attributes out of *kwargs*, then try to create and save
the object, then assign and save it's eav attributes.
"""
config_cls = getattr(self.model, '_eav_config_cls', None)
if not config_cls or config_cls.manager_only:
return super(EntityManager, self).create(**kwargs)
prefix = '%s__' % config_cls.eav_attr
new_kwargs = {}
eav_kwargs = {}
for key, value in kwargs.items():
if key.startswith(prefix):
eav_kwargs.update({key[len(prefix):]: value})
else:
new_kwargs.update({key: value})
obj = self.model(**new_kwargs)
obj_eav = getattr(obj, config_cls.eav_attr)
for key, value in eav_kwargs.items():
setattr(obj_eav, key, value)
obj.save()
return obj
def get_or_create(self, **kwargs):
"""
Reproduces the behavior of get_or_create, eav friendly.
"""
try:
return self.get(**kwargs), False
except self.model.DoesNotExist:
return self.create(**kwargs), True
class SerializableManager(BaseManager.from_queryset(SerializableQuerySet)):
def __init__(self,
*,
select_related=None,
prefetch_related=None,
default_serializer=BaseSerializer):
super().__init__()
self.default_serializer = default_serializer(
select_related=select_related,
prefetch_related=prefetch_related,
)
def to_serialize(self, serializer=None):
if serializer is None:
serializer = self.default_serializer
return self.get_queryset().to_serialize(serializer)
@classmethod
def from_queryset(cls, queryset_class, class_name=None):
if not hasattr(queryset_class, 'to_serialize'):
queryset_class = type(
f'SerializableQuerySet__{queryset_class.__name__}',
(SerializableQuerySet, queryset_class), {})
return super().from_queryset(queryset_class, class_name=class_name)
class LocationManager(BaseManager.from_queryset(LocationQuerySet)):
"""
class TaskManager(BaseManager.from_queryset(TaskQueryset)):
""" Custom manager from queryset """
pass
class BulkUpdateManager(BaseManager.from_queryset(BulkUpdateQuerySet), models.Manager):
pass
class PreparedStatementManager(
BaseManager.from_queryset(PreparedQuerySqlBuilder)):
_queryset_class = PreparedQuerySqlBuilder
class AccountManager(BaseManager.from_queryset(AccountQuerySet)):
use_for_related_fields = True
class PIDManager(BaseManager.from_queryset(PIDQuerySet)):
pass
class CitusComparisonManager(BaseManager.from_queryset(CitusComparisonQuerySet)
):
pass
class TreeManager(BaseManager.from_queryset(TreeQuerySet)):
def get_queryset(self):
return super().get_queryset().order_by("path")
class TestManager(BaseManager.from_queryset(TestQuerySet)):
pass
class NoDeleteManager(BaseManager.from_queryset(NoDeleteQuerySet)):
def get_queryset(self):
return super().get_queryset().filter(date_deleted__isnull=True)
class IdMapManager(BaseManager.from_queryset(IdMapQuerySet), Manager):
pass
class PreparedManager(BaseManager.from_queryset(PreparedQuerySet)):
pass
class VoteManager(BaseManager.from_queryset(VoteQuerySet)):
use_for_related_fields = True
class JobManager(BaseManager.from_queryset(JobQuerySet)):
use_for_related_fields = True
class GroupManager(BaseManager.from_queryset(GroupQuerySet)):
use_for_related_fields = True
class RiverManager(BaseManager.from_queryset(RiverQuerySet)):
pass
class TaskManager(BaseManager.from_queryset(TaskQuerySet)):
use_for_related_fields = True
class LegalPersonManager(BaseManager.from_queryset(LegalPersonQuerySet)):
pass
class JobManager(BaseManager.from_queryset(JobQueryset)):
""" Custom manager from queryset """
pass
name = models.CharField(max_length=100)
slug = models.SlugField()
class Meta:
ordering = ['name']
def __str__(self):
return self.name
class DoesNotExistQuerySet(QuerySet):
def get(self, *args, **kwargs):
raise Author.DoesNotExist
DoesNotExistBookManager = BaseManager.from_queryset(DoesNotExistQuerySet)
class Book(models.Model):
name = models.CharField(max_length=300)
slug = models.SlugField()
pages = models.IntegerField()
authors = models.ManyToManyField(Author)
pubdate = models.DateField()
objects = models.Manager()
does_not_exist = DoesNotExistBookManager()
class Meta:
ordering = ['-pubdate']
class SmartManager(BaseManager.from_queryset(SmartQuerySet)):
pass
class VirtualDeletionBaseManager(
BaseManager.from_queryset(VirtualDeletionQuerySet), Manager):
pass
class EntryManager(BaseManager.from_queryset(EntryQuerySet)):
use_for_related_fields = True
def get_manager(cls):
return type('DictManager', (BaseManager.from_queryset(cls)), {})
