Exports refactoring

1 files changed, 80 insertions, 6 deletions

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\[([^\]]*)\]')