Exports refactoring

2 files changed, 83 insertions, 81 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
diff --git a/chimere/utils.py b/chimere/utils.py
index 3266548..09255d0 100644
--- a/chimere/utils.py
+++ b/chimere/utils.py
@@ -21,20 +21,24 @@
 Utilitaries
 """
 
+import datetime
 import os
-import tempfile
 import re
+import StringIO
+import tempfile
 import urllib2
 import unicodedata
 import zipfile
-import StringIO
 
 from lxml import etree
 from osgeo import osr
 
 from django.conf import settings
-from django.contrib.gis.gdal import DataSource
+from django.contrib.gis.gdal import DataSource, Driver, OGRGeometry, \
+                                    OGRGeomType, SpatialReference, check_err
+from django.contrib.gis.gdal.libgdal import lgdal
 from django.core.exceptions import ObjectDoesNotExist
+from django.shortcuts import render_to_response
 from django.utils.translation import ugettext_lazy as _
 
 from chimere import get_version
@@ -192,13 +196,21 @@ class KMLManager(ImportManager):
                     m.categories.add(cat)
         return (new_item, updated_item, msg)
 
+    @classmethod
+    def export(cls, queryset):
+        dct = {'name':settings.PROJECT_NAME,
+               'description':unicode(datetime.date.today()),
+               'locations':queryset.all()
+        }
+        filename = unicode_normalize(settings.PROJECT_NAME + dct['description']\
+                                     + '.kml')
+        result = render_to_response('chimere/export.kml', dct)
+        return filename, result
+
 class ShapefileManager(ImportManager):
     u"""
     Shapefile importer
     """
-    def __init__(self, importer_instance):
-        self.importer_instance = importer_instance
-
     def get(self, source=None):
         u"""
         Get data from the source
@@ -306,6 +318,68 @@ class ShapefileManager(ImportManager):
             os.removedirs(dr)
         return (new_item, updated_item, msg)
 
+    @classmethod
+    def export(cls, queryset):
+        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.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()
+        return filename, zip_stream
 
 RE_NODE = re.compile('node\[([^\]]*)\]')