|
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465 |
- 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})
- )
|