Source

django-countries / django_countries / fields.py

Full commit
from django.db.models.fields import CharField
from django.utils.encoding import force_unicode, StrAndUnicode
from django_countries import settings, mobile_phone_numbers

from widgets import CountrySelectWidget, DynamicProvincesWidget

class Country(StrAndUnicode):
    def __init__(self, code=None, name=None):
        if not code and name:
            from django_countries.countries import COUNTRIES
            for code, cname in COUNTRIES:
                if cname == name:
                    break
                code = ''
        else:
            code = code.upper()
        self.code = code
    
    def __unicode__(self):
        return force_unicode(self.code or u'')

    def __eq__(self, other):
        return unicode(self) == force_unicode(other)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __cmp__(self, other):
        return cmp(unicode(self), force_unicode(other))

    def __hash__(self):
        return hash(unicode(self))

    def __repr__(self):
        return "%s(code=%r)" % (self.__class__.__name__, unicode(self))

    def __nonzero__(self):
        return bool(self.code)

    def __len__(self):
        return len(unicode(self))
    
    @property
    def name(self):
        # Local import so the countries aren't loaded unless they are needed. 
        from django_countries.countries import COUNTRIES
        for code, name in COUNTRIES:
            if self.code == code:
                return name
        return ''
    
    @property
    def flag(self):
        if not self.code:
            return ''
        return settings.FLAG_URL % {'code_upper': self.code,
                                    'code': self.code.lower()}
    
    @property
    def net(self):
        """
        Some of the ISO country codes differ from the 2-letter IANA network
        code. This is a mapping to allow these to still be used to get the
        province information for these countries.
        """
        from django_countries.countries import COUNTRY_CODES
        return COUNTRY_CODES.get(self.code)[1]
    
    @property
    def icc(self):
        from django_countries.countries import COUNTRY_CODES
        return COUNTRY_CODES.get(self.code)[0]
    
    @property
    def localflavor(self):
        """
        Helper for province-based stuff.
        """
        if not hasattr(self, '_localflavor'):
            try:
                # The following would be useful, except the Indian list of
                # states is a flat list, not a list of tuples.
                # if self.net in ["in", "is"]:
                #     net = self.net + "_"
                # else:
                #     net = self.net
                self._localflavor = __import__('django.contrib.localflavor.%s.forms' % self.net, fromlist=["django.contrib.localflavor"])
            except ImportError:
                try:
                    self._localflavor = __import__('django_countries.localflavor.%s.forms' % self.net, fromlist=["django_countries.localflavor"])
                except ImportError:
                    self._localflavor = None
        return self._localflavor
    
    @property
    def provinces_filtered_select_multiple(self):
        """
        This will currently only work for people who can view admin pages,
        since it uses an admin widget (and admin media).
        
        There is also a provinces_select_multiple that lacks the filtering,
        but is available anywhere.
        """
        if self.provinces_name:
            provinces_name = self.provinces_name
            provinces = self.provinces_select.choices
            from django.contrib.admin.widgets import FilteredSelectMultiple
            class SelectMultipleProvinces(FilteredSelectMultiple):
                def __init__(self, attrs=None):
                    super(SelectMultipleProvinces, self).__init__(provinces_name, False)
                def render(self, name, value, attrs=None):
                    return super(SelectMultipleProvinces, self).render(name, value, attrs, choices=provinces)
            return SelectMultipleProvinces()
        else:
            from django.forms import TextInput
            return TextInput()
    
    @property
    def provinces_select_multiple(self):
        """
        Multiple select widget for selecting provinces.
        """
        if self.provinces_name:
            provinces = self.provinces_select.choices
            from django.forms import SelectMultiple
            class SelectMultipleProvinces(SelectMultiple):
                def render(self, name, value, attrs=None):
                    return super(SelectMultipleProvinces, self).render(name, value, attrs, choices=provinces)
            return SelectMultipleProvinces()
                
    @property
    def provinces_select(self):
        """
        A Select Widget that provides choices based on the provinces of
        the country. Uses django.contrib.localflavor.
        Adds in a first empty entry. Note that self.provinces removes this,
        which if you want to use it in any other way, you may want to do.
        Or just use self.provinces.
        """
        if self.provinces_name:
            widget = getattr(self.localflavor, '%s%sSelect' % (self.net.upper(), self.provinces_name.replace(' ', '')))()
            widget.choices.insert(0, ('', 'Please select a province...'))
            return widget
    
    @property
    def provinces(self):
        if self.provinces_select:
            return [
                {'code': a, 'name': b} for a,b in self.provinces_select.choices[1:]
            ]
        
    @property
    def provinces_name(self):
        """
        What is the name of provinces in this country?
        
        .. note: 
            If there is more than one of these for a country, the `choices`
            list determines the order they will be checked.  For example,
            UK/GB will find Nation before County.
        """
        
        choices = [
            "Nation",
            "State", 
            "Province", 
            "Region", 
            "Municipality", 
            "Department", 
            "County",
            "Prefecture",
            "Provincial District"
        ]
        
        if self.localflavor:
            for choice in choices:
                if hasattr(self.localflavor, '%s%sSelect' % (self.net.upper(), choice.replace(' ',''))):
                    return choice
        return None
    
    def validate_mobile(self, number):
        """
        Validate the passed in string into number is a valid mobile
        phone number for this country.
        """
        return mobile_phone_numbers.validate(self.code, number)

class CountryDescriptor(object):
    """
    A descriptor for country fields on a model instance. Returns a Country when
    accessed so you can do stuff like::

        >>> instance.country.name
        u'New Zealand'
        
        >>> instance.country.flag
        '/static/flags/nz.gif'
    """
    def __init__(self, field):
        self.field = field

    def __get__(self, instance=None, owner=None):
        if instance is None:
            raise AttributeError(
                "The '%s' attribute can only be accessed from %s instances."
                % (self.field.name, owner.__name__))
        return Country(code=instance.__dict__[self.field.name])

    def __set__(self, instance, value):
        if value is not None:
            value = force_unicode(value)
        instance.__dict__[self.field.name] = value


class CountryField(CharField):
    """
    A country field for Django models that provides all ISO 3166-1 countries as
    choices.
    
    """
    descriptor_class = CountryDescriptor
 
    def __init__(self, *args, **kwargs):
        # Local import so the countries aren't loaded unless they are needed. 
        from django_countries.countries import COUNTRIES 

        kwargs.setdefault('max_length', 2) 
        kwargs.setdefault('choices', COUNTRIES) 
        
        super(CharField, self).__init__(*args, **kwargs) 

    def get_internal_type(self): 
        return "CharField"
    
    def formfield(self, **kwargs):
        defaults = {
            'widget': CountrySelectWidget
        }
        defaults.update(kwargs)
        return super(CountryField, self).formfield(**defaults)

    def contribute_to_class(self, cls, name):
        super(CountryField, self).contribute_to_class(cls, name)
        setattr(cls, self.name, self.descriptor_class(self))

    def get_prep_lookup(self, lookup_type, value):
        if hasattr(value, 'code'):
            value = value.code
        return super(CountryField, self).get_prep_lookup(lookup_type, value)

    def pre_save(self, *args, **kwargs):
        "Returns field's value just before saving."
        value = super(CharField, self).pre_save(*args, **kwargs)
        return self.get_prep_value(value)

    def get_prep_value(self, value):
        "Returns field's value prepared for saving into a database."
        # Convert the Country to unicode for database insertion.
        if value is None:
            return None
        return unicode(value)


# If south is installed, ensure that CountryField will be introspected just
# like a normal CharField.
try:
    from south.modelsinspector import add_introspection_rules
    add_introspection_rules([], ['^django_countries\.fields\.CountryField'])
except ImportError:
    pass