Exports refactoring

1 files changed, 3 insertions, 75 deletions

diff --git a/chimere/admin.py b/chimere/admin.py
index b6febd1..cfe3bf7 100644
--- a/chimere/admin.py
+++ b/chimere/admin.py
@@ -21,17 +21,10 @@
 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 _
@@ -42,7 +35,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.utils import unicode_normalize, ShapefileManager, KMLManager
 from chimere.widgets import TextareaWidget
 
 def get_areas_for_user(user):
@@ -68,15 +61,9 @@ def export_to_kml(modeladmin, request, queryset):
     u"""
     Export data to KML
     """
-    dct = {'name':settings.PROJECT_NAME,
-           'description':unicode(datetime.date.today()),
-           'locations':queryset.all()
-    }
-    result = render_to_response('chimere/export.kml', dct)
+    filename, result = KMLManager.export(queryset)
     response = HttpResponse(result,
                         mimetype='application/vnd.google-earth.kml+xml')
-    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")
@@ -85,66 +72,7 @@ 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()
-
+    filename, zip_stream = ShapefileManager.export(queryset)
     # Stick it all in a django HttpResponse
     response = HttpResponse()
     response['Content-Disposition'] = 'attachment; filename=%s.zip' % filename