from django.conf import settings from django.core import checks from django.core.exceptions import FieldDoesNotExist from django.db import models class CurrentSiteManager(models.Manager): "Use this to limit objects to those associated with the current site." use_in_migrations = True def __init__(self, field_name=None): super().__init__() self.__field_name = field_name def check(self, **kwargs): errors = super().check(**kwargs) errors.extend(self._check_field_name()) return errors def _check_field_name(self): field_name = self._get_field_name() try: field = self.model._meta.get_field(field_name) except FieldDoesNotExist: return [ checks.Error( "CurrentSiteManager could not find a field named '%s'." % field_name, obj=self, id='sites.E001', ) ] if not field.many_to_many and not isinstance(field, (models.ForeignKey)): return [ checks.Error( "CurrentSiteManager cannot use '%s.%s' as it is not a foreign key or a many-to-many field." % ( self.model._meta.object_name, field_name ), obj=self, id='sites.E002', ) ] return [] def _get_field_name(self): """ Return self.__field_name or 'site' or 'sites'. """ if not self.__field_name: try: self.model._meta.get_field('site') except FieldDoesNotExist: self.__field_name = 'sites' else: self.__field_name = 'site' return self.__field_name def get_queryset(self): return super().get_queryset().filter(**{self._get_field_name() + '__id': settings.SITE_ID})