from django.db import models


class LocalityQuerySet(models.QuerySet):
    get_by_name_or_code_fields = ()

    def get_by_name_or_code(self, name_or_code, **kwargs):
        for field in self.get_by_name_or_code_fields:
            try:
                return self.get(**{field: name_or_code}, **kwargs)
            except (self.model.DoesNotExist, self.model.MultipleObjectsReturned):
                continue
        return None


class CountryQuerySet(LocalityQuerySet):
    get_by_name_or_code_fields = ('name', 'alpha_2_code', 'alpha_3_code')


class RegionQuerySet(LocalityQuerySet):
    get_by_name_or_code_fields = ('name', 'code')


class MunicipalityQuerySet(LocalityQuerySet):
    get_by_name_or_code_fields = ('name', 'short_name')


__all__ = (
    'CountryQuerySet',
    'RegionQuerySet',
    'MunicipalityQuerySet',
)