| from django.conf import settings |
| from django.db import models |
| from django.db.models.fields import FieldDoesNotExist |
| |
| class CurrentSiteManager(models.Manager): |
| "Use this to limit objects to those associated with the current site." |
| def __init__(self, field_name=None): |
| super(CurrentSiteManager, self).__init__() |
| self.__field_name = field_name |
| self.__is_validated = False |
| |
| def _validate_field_name(self): |
| field_names = self.model._meta.get_all_field_names() |
| |
| # If a custom name is provided, make sure the field exists on the model |
| if self.__field_name is not None and self.__field_name not in field_names: |
| raise ValueError("%s couldn't find a field named %s in %s." % \ |
| (self.__class__.__name__, self.__field_name, self.model._meta.object_name)) |
| |
| # Otherwise, see if there is a field called either 'site' or 'sites' |
| else: |
| for potential_name in ['site', 'sites']: |
| if potential_name in field_names: |
| self.__field_name = potential_name |
| self.__is_validated = True |
| break |
| |
| # Now do a type check on the field (FK or M2M only) |
| try: |
| field = self.model._meta.get_field(self.__field_name) |
| if not isinstance(field, (models.ForeignKey, models.ManyToManyField)): |
| raise TypeError("%s must be a ForeignKey or ManyToManyField." %self.__field_name) |
| except FieldDoesNotExist: |
| raise ValueError("%s couldn't find a field named %s in %s." % \ |
| (self.__class__.__name__, self.__field_name, self.model._meta.object_name)) |
| self.__is_validated = True |
| |
| def get_query_set(self): |
| if not self.__is_validated: |
| self._validate_field_name() |
| return super(CurrentSiteManager, self).get_query_set().filter(**{self.__field_name + '__id__exact': settings.SITE_ID}) |
