Blog

The Digital Agency for International Development

Extending django-countries for use in many to many fields

By Daniel Levy on 25 January 2013

Recently, I wanted to associate models in one of our applications with countries. The obvious starting point was to install django-countries and use its CountryField. However, this only allows you to associate one country with a model and I wanted to associate many countries.

My initial solution was to create the model below and use that.

class Country(models.Model):
        country = CountryField(unique=True)

However, this model didn't really do what I need it to do and was awkward to use because of the primary keys being numbers instead of country codes, as well as having country names hard-coded in a file instead of readily available in the database. This makes it harder to do things like adding countries and debugging data integrity issues in the database.

I needed to replace it without losing existing country associations. The first thing to do was to start a new application, let's call it "apti_countries". In apti_countries I created a new model that stored the country name and had the two letter country code as a primary key.

class Country(models.Model):
    code = CharField(max_length=2, primary_key=True)
    name = CharField(max_length=255, unique=True)

The next part was to add the data from the countries. To do this, I used a data migration:

class Migration(DataMigration):
    def forwards(self, orm):
       for old_country in orm.ExisitingModelWithCountries.objects.all():
             # Country.name involves non-ascii characters and this can cause
             # south difficulties. The soultion is to encode it as UTF-8

             country_name = old_country.country.name.encode("UTF-8")
             new_country = orm["apti_countries.country"](name=country_name,
                                       code=old_country.country.code)
             new_country.save()

Now I need to add new many to many fields on exisitng models that pointed to this one and move the data across. Because both models are called Country, the new model was imported as NewCountry.

class ExisitingModelWithCountries(models.Model):
    countries = ManyToManyField(Country)
    new_countries = ManyToManyField(NewCountry)

I used the migration below to transfer the old country data to the new field.

class Migration(DataMigration):
    def forwards(self, orm):
        for item in orm.ExisitingModelWithCountries.objects.all():
            for country in item.countries.all():
                country_name = country.country.name.encode("UTF-8")
                new_country = orm["apti_countries.country"].objects. \
                    get(name=country_name)
                item.new_countries.add(new_country)

This works, but has the problem of depending on a migration in a different app. The solution to this was to add a "depends_on" tuple of tuples pointing to the apti_countries migration for adding the new countries.

class Migration(DataMigration):
    depends_on = (
        ("apti_countries", "0002_create_new_countries"),
    )

    def forwards(self, orm):
        for item in orm.ExisitingModelWithCountries.objects.all():
            for country in item.countries.all():
                country_name = country.country.name.encode("UTF-8")
                new_country = orm["apti_countries.country"].objects. \
                    get(name=country_name)
                item.new_countries.add(new_country)

Finally, I deleted the old country many to many field and renamed the new one, as well as removing the old model that was no longer used.

class ExisitingModelWithCountries(models.Model):
    countries = ManyToManyField(Country)

I now have an extensible Country model that can easily be used with generic views and can show the country code in the url when used with generic views.