Finish django countries implementation

* Add a migration to drop the old countries field.
* Update all templates/views/utility methods to point at the new country
  field and dereference it as necessary.
* Add the flags images to a few views where it makes sense.
* Cleanup the download page layout quite a bit.
* Bump the mirror status JSON version to 3; add country_code attribute.

Signed-off-by: Dan McGee <dan@archlinux.org>

4 files changed, 130 insertions, 34 deletions

diff --git a/mirrors/migrations/0016_auto__del_field_mirror_country_old__del_field_mirrorurl_country_old.py b/mirrors/migrations/0016_auto__del_field_mirror_country_old__del_field_mirrorurl_country_old.py
new file mode 100644
index 0000000..ae917ec
--- /dev/null
+++ b/mirrors/migrations/0016_auto__del_field_mirror_country_old__del_field_mirrorurl_country_old.py
@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+from south.db import db
+from south.v2 import SchemaMigration
+from django.db import models
+
+
+class Migration(SchemaMigration):
+
+    def forwards(self, orm):
+        db.delete_column('mirrors_mirror', 'country_old')
+        db.delete_column('mirrors_mirrorurl', 'country_old')
+        db.create_index('mirrors_mirror', ['country'])
+        db.create_index('mirrors_mirrorurl', ['country'])
+
+    def backwards(self, orm):
+        db.delete_index('mirrors_mirrorurl', ['country'])
+        db.delete_index('mirrors_mirror', ['country'])
+        db.add_column('mirrors_mirror', 'country_old',
+                      self.gf('django.db.models.fields.CharField')(default='Any', max_length=255, db_index=True),
+                      keep_default=False)
+        db.add_column('mirrors_mirrorurl', 'country_old',
+                      self.gf('mirrors.models.NullCharField')(blank=True, max_length=255, null=True, db_index=True),
+                      keep_default=False)
+
+    models = {
+        'mirrors.mirror': {
+            'Meta': {'ordering': "('country', 'name')", 'object_name': 'Mirror'},
+            'active': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
+            'admin_email': ('django.db.models.fields.EmailField', [], {'max_length': '255', 'blank': 'True'}),
+            'country': ('django_countries.fields.CountryField', [], {'db_index': 'True', 'max_length': '2', 'blank': 'True'}),
+            'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
+            'isos': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
+            'name': ('django.db.models.fields.CharField', [], {'unique': 'True', 'max_length': '255'}),
+            'notes': ('django.db.models.fields.TextField', [], {'blank': 'True'}),
+            'public': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
+            'rsync_password': ('django.db.models.fields.CharField', [], {'default': "''", 'max_length': '50', 'blank': 'True'}),
+            'rsync_user': ('django.db.models.fields.CharField', [], {'default': "''", 'max_length': '50', 'blank': 'True'}),
+            'tier': ('django.db.models.fields.SmallIntegerField', [], {'default': '2'}),
+            'upstream': ('django.db.models.fields.related.ForeignKey', [], {'to': "orm['mirrors.Mirror']", 'null': 'True', 'on_delete': 'models.SET_NULL'})
+        },
+        'mirrors.mirrorlog': {
+            'Meta': {'object_name': 'MirrorLog'},
+            'check_time': ('django.db.models.fields.DateTimeField', [], {'db_index': 'True'}),
+            'duration': ('django.db.models.fields.FloatField', [], {'null': 'True'}),
+            'error': ('django.db.models.fields.CharField', [], {'default': "''", 'max_length': '255', 'blank': 'True'}),
+            'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
+            'is_success': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
+            'last_sync': ('django.db.models.fields.DateTimeField', [], {'null': 'True'}),
+            'url': ('django.db.models.fields.related.ForeignKey', [], {'related_name': "'logs'", 'to': "orm['mirrors.MirrorUrl']"})
+        },
+        'mirrors.mirrorprotocol': {
+            'Meta': {'ordering': "('protocol',)", 'object_name': 'MirrorProtocol'},
+            'default': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
+            'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
+            'is_download': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
+            'protocol': ('django.db.models.fields.CharField', [], {'unique': 'True', 'max_length': '10'})
+        },
+        'mirrors.mirrorrsync': {
+            'Meta': {'object_name': 'MirrorRsync'},
+            'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
+            'ip': ('django.db.models.fields.CharField', [], {'max_length': '24'}),
+            'mirror': ('django.db.models.fields.related.ForeignKey', [], {'related_name': "'rsync_ips'", 'to': "orm['mirrors.Mirror']"})
+        },
+        'mirrors.mirrorurl': {
+            'Meta': {'object_name': 'MirrorUrl'},
+            'country': ('django_countries.fields.CountryField', [], {'db_index': 'True', 'max_length': '2', 'blank': 'True'}),
+            'has_ipv4': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
+            'has_ipv6': ('django.db.models.fields.BooleanField', [], {'default': 'False'}),
+            'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
+            'mirror': ('django.db.models.fields.related.ForeignKey', [], {'related_name': "'urls'", 'to': "orm['mirrors.Mirror']"}),
+            'protocol': ('django.db.models.fields.related.ForeignKey', [], {'related_name': "'urls'", 'on_delete': 'models.PROTECT', 'to': "orm['mirrors.MirrorProtocol']"}),
+            'url': ('django.db.models.fields.CharField', [], {'unique': 'True', 'max_length': '255'})
+        }
+    }
+
+    complete_apps = ['mirrors']
diff --git a/mirrors/models.py b/mirrors/models.py
index 7996841..c23622f 100644
--- a/mirrors/models.py
+++ b/mirrors/models.py
@@ -25,8 +25,7 @@ class Mirror(models.Model):
     name = models.CharField(max_length=255, unique=True)
     tier = models.SmallIntegerField(default=2, choices=TIER_CHOICES)
     upstream = models.ForeignKey('self', null=True, on_delete=models.SET_NULL)
-    country_old = models.CharField(max_length=255, db_index=True)
-    country = CountryField(blank=True)
+    country = CountryField(blank=True, db_index=True)
     admin_email = models.EmailField(max_length=255, blank=True)
     public = models.BooleanField(default=True)
     active = models.BooleanField(default=True)
@@ -70,9 +69,7 @@ class MirrorUrl(models.Model):
     protocol = models.ForeignKey(MirrorProtocol, related_name="urls",
             editable=False, on_delete=models.PROTECT)
     mirror = models.ForeignKey(Mirror, related_name="urls")
-    country_old = NullCharField(max_length=255, null=True, blank=True,
-            db_index=True)
-    country = CountryField(blank=True)
+    country = CountryField(blank=True, db_index=True)
     has_ipv4 = models.BooleanField("IPv4 capable", default=True,
             editable=False)
     has_ipv6 = models.BooleanField("IPv6 capable", default=False,
@@ -90,7 +87,7 @@ class MirrorUrl(models.Model):
 
     @property
     def real_country(self):
-        return self.country_old or self.mirror.country_old
+        return self.country or self.mirror.country
 
     def clean(self):
         try:
diff --git a/mirrors/utils.py b/mirrors/utils.py
index aa1e9f7..728c304 100644
--- a/mirrors/utils.py
+++ b/mirrors/utils.py
@@ -1,6 +1,7 @@
 from datetime import timedelta
 
 from django.db.models import Avg, Count, Max, Min, StdDev
+from django_countries.fields import Country
 
 from main.utils import cache_function, utc_now
 from .models import MirrorLog, MirrorProtocol, MirrorUrl
@@ -89,13 +90,14 @@ def get_mirror_errors(cutoff=default_cutoff):
     errors = MirrorLog.objects.filter(
             is_success=False, check_time__gte=cutoff_time,
             url__mirror__active=True, url__mirror__public=True).values(
-            'url__url', 'url__country_old', 'url__protocol__protocol',
-            'url__mirror__country_old', 'error').annotate(
+            'url__url', 'url__country', 'url__protocol__protocol',
+            'url__mirror__country', 'error').annotate(
             error_count=Count('error'), last_occurred=Max('check_time')
             ).order_by('-last_occurred', '-error_count')
     errors = list(errors)
     for err in errors:
-        err['country'] = err['url__country_old'] or err['url__mirror__country_old']
+        ctry_code = err['url__country'] or err['url__mirror__country']
+        err['country'] = Country(ctry_code)
     return errors
 
 
@@ -114,16 +116,15 @@ def get_mirror_url_for_download(cutoff=default_cutoff):
         best_logs = MirrorLog.objects.filter(is_success=True,
                 check_time__gte=min_check_time, last_sync__gte=min_sync_time,
                 url__mirror__public=True, url__mirror__active=True,
-                url__protocol__protocol__iexact='HTTP').order_by(
+                url__protocol__default=True).order_by(
                 'duration')[:1]
         if best_logs:
             return MirrorUrl.objects.get(id=best_logs[0].url_id)
 
     mirror_urls = MirrorUrl.objects.filter(
-            mirror__public=True, mirror__active=True,
-            protocol__protocol__iexact='HTTP')
-    # look first for an 'Any' URL, then fall back to any HTTP URL
-    filtered_urls = mirror_urls.filter(mirror__country_old='Any')[:1]
+            mirror__public=True, mirror__active=True, protocol__default=True)
+    # look first for a country-agnostic URL, then fall back to any HTTP URL
+    filtered_urls = mirror_urls.filter(mirror__country='')[:1]
     if not filtered_urls:
         filtered_urls = mirror_urls[:1]
     if not filtered_urls:
diff --git a/mirrors/views.py b/mirrors/views.py
index c5989d6..349c17d 100644
--- a/mirrors/views.py
+++ b/mirrors/views.py
@@ -1,3 +1,6 @@
+import datetime
+from operator import attrgetter, itemgetter
+
 from django import forms
 from django.core.serializers.json import DjangoJSONEncoder
 from django.db.models import Q
@@ -6,12 +9,13 @@ from django.shortcuts import get_object_or_404
 from django.views.decorators.csrf import csrf_exempt
 from django.views.generic.simple import direct_to_template
 from django.utils import simplejson
+from django_countries.countries import COUNTRIES
 
-from main.utils import make_choice
 from .models import Mirror, MirrorUrl, MirrorProtocol
 from .utils import get_mirror_statuses, get_mirror_errors
 
-import datetime
+COUNTRY_LOOKUP = dict(COUNTRIES)
+
 
 class MirrorlistForm(forms.Form):
     country = forms.MultipleChoiceField(required=False)
@@ -22,17 +26,27 @@ class MirrorlistForm(forms.Form):
 
     def __init__(self, *args, **kwargs):
         super(MirrorlistForm, self).__init__(*args, **kwargs)
-        countries = Mirror.objects.filter(active=True).values_list(
-                'country_old', flat=True).distinct().order_by('country_old')
-        self.fields['country'].choices = [('all','All')] + make_choice(
-                countries)
-        self.fields['country'].initial = ['all']
-        protos = make_choice(
-                MirrorProtocol.objects.filter(is_download=True))
+        fields = self.fields
+        fields['country'].choices = [('all','All')] + self.get_countries()
+        fields['country'].initial = ['all']
+        protos = [(p.protocol, p.protocol) for p in
+                MirrorProtocol.objects.filter(is_download=True)]
         initial = MirrorProtocol.objects.filter(is_download=True, default=True)
-        self.fields['protocol'].choices = protos
-        self.fields['protocol'].initial = [p.protocol for p in initial]
-        self.fields['ip_version'].initial = ['4']
+        fields['protocol'].choices = protos
+        fields['protocol'].initial = [p.protocol for p in initial]
+        fields['ip_version'].initial = ['4']
+
+    def get_countries(self):
+        country_codes = set()
+        country_codes.update(Mirror.objects.filter(active=True).exclude(
+                country='').values_list(
+                'country', flat=True).order_by().distinct())
+        country_codes.update(MirrorUrl.objects.filter(
+                mirror__active=True).exclude(country='').values_list(
+                'country', flat=True).order_by().distinct())
+        countries = [(code, COUNTRY_LOOKUP[code]) for code in country_codes]
+        return sorted(countries, key=itemgetter(1))
+
 
 @csrf_exempt
 def generate_mirrorlist(request):
@@ -52,6 +66,7 @@ def generate_mirrorlist(request):
     return direct_to_template(request, 'mirrors/index.html',
             {'mirrorlist_form': form})
 
+
 def find_mirrors(request, countries=None, protocols=None, use_status=False,
         ipv4_supported=True, ipv6_supported=True):
     if not protocols:
@@ -62,8 +77,8 @@ def find_mirrors(request, countries=None, protocols=None, use_status=False,
             mirror__public=True, mirror__active=True,
     )
     if countries and 'all' not in countries:
-        qset = qset.filter(Q(country_old__in=countries) |
-                Q(mirror__country_old__in=countries))
+        qset = qset.filter(Q(country__in=countries) |
+                Q(mirror__country__in=countries))
 
     ip_version = Q()
     if ipv4_supported:
@@ -74,7 +89,7 @@ def find_mirrors(request, countries=None, protocols=None, use_status=False,
 
     if not use_status:
         urls = qset.order_by('mirror__name', 'url')
-        urls = sorted(urls, key=lambda x: x.real_country)
+        urls = sorted(urls, key=attrgetter('real_country'))
         template = 'mirrors/mirrorlist.txt'
     else:
         status_info = get_mirror_statuses()
@@ -96,13 +111,15 @@ def find_mirrors(request, countries=None, protocols=None, use_status=False,
             },
             mimetype='text/plain')
 
+
 def mirrors(request):
-    mirror_list = Mirror.objects.select_related().order_by('tier', 'country_old')
+    mirror_list = Mirror.objects.select_related().order_by('tier', 'country')
     if not request.user.is_authenticated():
         mirror_list = mirror_list.filter(public=True, active=True)
     return direct_to_template(request, 'mirrors/mirrors.html',
             {'mirror_list': mirror_list})
 
+
 def mirror_details(request, name):
     mirror = get_object_or_404(Mirror, name=name)
     if not request.user.is_authenticated() and \
@@ -116,11 +133,12 @@ def mirror_details(request, name):
     # get each item from checked_urls and supplement with anything in all_urls
     # if it wasn't there
     all_urls = set(checked_urls).union(all_urls)
-    all_urls = sorted(all_urls, key=lambda x: x.url)
+    all_urls = sorted(all_urls, key=attrgetter('url'))
 
     return direct_to_template(request, 'mirrors/mirror_details.html',
             {'mirror': mirror, 'urls': all_urls})
 
+
 def status(request):
     bad_timedelta = datetime.timedelta(days=3)
     status_info = get_mirror_statuses()
@@ -143,6 +161,7 @@ def status(request):
     })
     return direct_to_template(request, 'mirrors/status.html', context)
 
+
 class MirrorStatusJSONEncoder(DjangoJSONEncoder):
     '''Base JSONEncoder extended to handle datetime.timedelta and MirrorUrl
     serialization. The base class takes care of datetime.datetime types.'''
@@ -159,17 +178,20 @@ class MirrorStatusJSONEncoder(DjangoJSONEncoder):
         if isinstance(obj, MirrorUrl):
             data = dict((attr, getattr(obj, attr))
                     for attr in self.url_attributes)
-            # separate because it isn't on the URL directly
-            data['country'] = obj.real_country
+            # get any override on the country attribute first
+            country = obj.real_country
+            data['country'] = unicode(country.name)
+            data['country_code'] = country.code
             return data
         if isinstance(obj, MirrorProtocol):
             return unicode(obj)
         return super(MirrorStatusJSONEncoder, self).default(obj)
 
+
 def status_json(request):
     status_info = get_mirror_statuses()
     data = status_info.copy()
-    data['version'] = 2
+    data['version'] = 3
     to_json = simplejson.dumps(data, ensure_ascii=False,
             cls=MirrorStatusJSONEncoder)
     response = HttpResponse(to_json, mimetype='application/json')