Shapefile export

2 files changed, 90 insertions, 5 deletions

diff --git a/chimere/admin.py b/chimere/admin.py
index 0ef34f3..b6febd1 100644
--- a/chimere/admin.py
+++ b/chimere/admin.py
@@ -21,11 +21,17 @@
 Settings for administration pages
 """
 import datetime
+import tempfile
 import unicodedata
+import zipfile
+from StringIO import StringIO
 
 from django import forms
 from django.conf import settings
 from django.contrib import admin
+from django.contrib.gis.gdal.libgdal import lgdal
+from django.contrib.gis.gdal import Driver, OGRGeometry, OGRGeomType, \
+                                    SpatialReference, check_err
 from django.http import HttpResponse
 from django.shortcuts import render_to_response
 from django.utils.translation import ugettext_lazy as _
@@ -36,6 +42,7 @@ from chimere.forms import MarkerAdminForm, RouteAdminForm, AreaAdminForm,\
 from chimere.models import Category, Icon, SubCategory, Marker, \
            PropertyModel, News, Route, Area, ColorTheme, Color, RouteFile,\
            MultimediaType, MultimediaFile, PictureFile, Importer
+from chimere.utils import unicode_normalize
 from chimere.widgets import TextareaWidget
 
 def get_areas_for_user(user):
@@ -68,13 +75,85 @@ def export_to_kml(modeladmin, request, queryset):
     result = render_to_response('chimere/export.kml', dct)
     response = HttpResponse(result,
                         mimetype='application/vnd.google-earth.kml+xml')
-    filename = ''.join(
-        (c for c in unicodedata.normalize('NFD', settings.PROJECT_NAME)
-        if unicodedata.category(c) != 'Mn')) + '-' + dct['description'] + '.kml'
+    filename = unicode_normalize(settings.PROJECT_NAME + dct['description'] + \
+                                 '.kml')
     response['Content-Disposition'] = 'attachment; filename=%s' % filename
     return response
 export_to_kml.short_description = _(u"Export to KML")
 
+def export_to_shapefile(modeladmin, request, queryset):
+    u"""
+    Export data to Shapefile
+    """
+    date = unicode(datetime.date.today())
+
+    tmp = tempfile.NamedTemporaryFile(suffix='.shp', mode='w+b')
+    tmp.close()
+
+    tmp_name = tmp.name
+    field_names = [field.name for field in queryset.model._meta.fields]
+    geo_field = getattr(queryset.model,
+                        'point' if 'point' in field_names else 'route')._field
+
+    dr = Driver('ESRI Shapefile')
+    ds = lgdal.OGR_Dr_CreateDataSource(dr._ptr, tmp_name, None)
+    if ds is None:
+        raise Exception(_(u'Could not create file!'))
+    ogr_type = OGRGeomType(geo_field.geom_type).num
+    srs = SpatialReference(geo_field.srid)
+
+    layer = lgdal.OGR_DS_CreateLayer(ds, 'lyr', srs._ptr, ogr_type, None)
+
+    for field_name in ('name', 'category'):
+        fld = lgdal.OGR_Fld_Create(field_name, 4)
+        added = lgdal.OGR_L_CreateField(layer, fld, 0)
+        check_err(added)
+
+    feature_def = lgdal.OGR_L_GetLayerDefn(layer)
+
+    for item in queryset:
+        # duplicate items when in several categories
+        for category in item.categories.all():
+            feat = lgdal.OGR_F_Create(feature_def)
+            lgdal.OGR_F_SetFieldString(feat, 0,
+                                   str(unicode_normalize(item.name)[:80]))
+            lgdal.OGR_F_SetFieldString(feat, 1,
+                                   str(unicode_normalize(category.name)[:80]))
+
+            geom = getattr(item, geo_field.name)
+            if not geom:
+                continue
+            ogr_geom = OGRGeometry(geom.wkt, srs)
+            check_err(lgdal.OGR_F_SetGeometry(feat, ogr_geom._ptr))
+            check_err(lgdal.OGR_L_SetFeature(layer, feat))
+    # Cleaning up
+    check_err(lgdal.OGR_L_SyncToDisk(layer))
+    lgdal.OGR_DS_Destroy(ds)
+    lgdal.OGRCleanupAll()
+
+    # writing to a zip file
+    filename = unicode_normalize(settings.PROJECT_NAME) + '-' + date
+    buff = StringIO()
+    zip_file = zipfile.ZipFile(buff, 'w', zipfile.ZIP_DEFLATED)
+    suffixes = ['shp', 'shx', 'prj', 'dbf']
+    for suffix in suffixes:
+        name = tmp_name.replace('.shp', '.' + suffix)
+        arcname = '.'.join((filename, suffix))
+        zip_file.write(name, arcname=arcname)
+    zip_file.close()
+    buff.flush()
+    zip_stream = buff.getvalue()
+    buff.close()
+
+    # Stick it all in a django HttpResponse
+    response = HttpResponse()
+    response['Content-Disposition'] = 'attachment; filename=%s.zip' % filename
+    response['Content-length'] = str(len(zip_stream))
+    response['Content-Type'] = 'application/zip'
+    response.write(zip_stream)
+    return response
+export_to_shapefile.short_description = _(u"Export to Shapefile")
+
 class PictureInline(admin.TabularInline):
     model = PictureFile
     extra = 1
@@ -90,7 +169,7 @@ class MarkerAdmin(admin.ModelAdmin):
     search_fields = ("name",)
     list_display = ('name', 'status')
     list_filter = ('status', 'categories')
-    actions = [validate, export_to_kml]
+    actions = [validate, export_to_kml, export_to_shapefile]
     exclude = ['submiter_session_key', 'import_key', 'import_version']
     readonly_fields = ['submiter_email', 'submiter_comment', 'import_source']
     if 'chimere_rss' in settings.INSTALLED_APPS:
diff --git a/chimere/utils.py b/chimere/utils.py
index b2c8ccc..97f710f 100644
--- a/chimere/utils.py
+++ b/chimere/utils.py
@@ -22,6 +22,7 @@ Utilitaries
 """
 
 import urllib2, re
+import unicodedata
 from external_utils import OsmApi
 from lxml import etree
 from chimere import get_version
@@ -30,6 +31,11 @@ from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
 from django.utils.translation import ugettext_lazy as _
 
+def unicode_normalize(string):
+    return ''.join(
+        (c for c in unicodedata.normalize('NFD', string)
+        if unicodedata.category(c) != 'Mn'))
+
 class ImportManager:
     u"""
     Generic class for specific importers
@@ -130,7 +136,7 @@ RE_NODE = re.compile('node\[([^\]]*)\]')
 
 class OSMManager(ImportManager):
     u"""
-    OSM importer
+    OSM importer/exporter
     The source url is a path to an OSM file or a XAPI url
     The filtr argument is XAPI args or empty if it is an OSM file.
     """