managers.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. # -*- coding: utf-8 -*-
  2. from __future__ import unicode_literals
  3. from django.conf import settings
  4. from django.core import checks
  5. from django.db import models
  6. from django.db.models.fields import FieldDoesNotExist
  7. class CurrentSiteManager(models.Manager):
  8. "Use this to limit objects to those associated with the current site."
  9. def __init__(self, field_name=None):
  10. super(CurrentSiteManager, self).__init__()
  11. self.__field_name = field_name
  12. def check(self, **kwargs):
  13. errors = super(CurrentSiteManager, self).check(**kwargs)
  14. errors.extend(self._check_field_name())
  15. return errors
  16. def _check_field_name(self):
  17. field_name = self._get_field_name()
  18. try:
  19. field = self.model._meta.get_field(field_name)
  20. except FieldDoesNotExist:
  21. return [
  22. checks.Error(
  23. "CurrentSiteManager could not find a field named '%s'." % field_name,
  24. hint=None,
  25. obj=self,
  26. id='sites.E001',
  27. )
  28. ]
  29. if not isinstance(field, (models.ForeignKey, models.ManyToManyField)):
  30. return [
  31. checks.Error(
  32. "CurrentSiteManager cannot use '%s.%s' as it is not a ForeignKey or ManyToManyField." % (
  33. self.model._meta.object_name, field_name
  34. ),
  35. hint=None,
  36. obj=self,
  37. id='sites.E002',
  38. )
  39. ]
  40. return []
  41. def _get_field_name(self):
  42. """ Return self.__field_name or 'site' or 'sites'. """
  43. if not self.__field_name:
  44. try:
  45. self.model._meta.get_field('site')
  46. except FieldDoesNotExist:
  47. self.__field_name = 'sites'
  48. else:
  49. self.__field_name = 'site'
  50. return self.__field_name
  51. def get_queryset(self):
  52. return super(CurrentSiteManager, self).get_queryset().filter(
  53. **{self._get_field_name() + '__id': settings.SITE_ID})