Commits

Anonymous committed 155b327

Code cleanup

  • Participants
  • Parent commits be4b13d

Comments (0)

Files changed (2)

File shapes/forms.py

+import zipfile
+import tempfile
 from django import forms
+from django.forms.util import ValidationError
 from django.utils.translation import ugettext_lazy as _
 from django.conf import settings
 from django.contrib.gis.gdal import DataSource
-from django.forms.util import ValidationError
-import zipfile
-import tempfile
+
 
 #http://docs.djangoproject.com/en/dev/topics/http/file-uploads/
 #http://www.neverfriday.com/sweetfriday/2008/09/-a-long-time-ago.html 

File shapes/views/export.py

 import zipfile
 import tempfile
 import datetime
-
-from django.db.models.query import QuerySet, ValuesQuerySet
+import cStringIO
+from django.http import HttpResponse
 from django.contrib.gis.db.models.fields import GeometryField
-from django.utils import simplejson
-from django.http import HttpResponse
-from django.utils.encoding import smart_str, smart_unicode
-from django.utils.translation import ugettext as _
-
-# ctypes stuff
-from cStringIO import StringIO
-from ctypes import string_at, c_double
-from django.contrib.gis.gdal.libgdal import lgdal
-
-#from django.contrib.gis.gdal import *
-from django.contrib.gis.gdal import OGRGeometry, OGRGeomType, SpatialReference, check_err
+from django.contrib.gis.gdal.libgdal import lgdal as ogr
+from django.contrib.gis.gdal import OGRGeometry, OGRGeomType, SpatialReference, check_err 
 
 # todo use: qs.query._geo_field()
-# todo abstract lgdal stuff
 # todo: support multiple querysets == multiple shapefiles
 
 class ShpResponder(object):
-    def __init__(self, queryset, proj_transform=None, *args, **kwargs):
-        # Available data
+    def __init__(self, queryset, *args, **kwargs):
+        """
+        """
         self.queryset = queryset
+        self.proj_transform = None
         self.mimetype = 'application/zip'
-        self.proj_transform = proj_transform
+        self.file_name = 'shp_download'
+        self.file_name = self.file_name.rstrip('.shp')
     
     def __call__(self, request, *args, **kwargs):
-        
-        lgdal.OGRRegisterAll()
-        
-        tmp = tempfile.NamedTemporaryFile(suffix='.shp', mode = 'w')
-        path = os.path.dirname(tmp.name)
-        name = tmp.name.split('/')[-1]
-        fullname = tmp.name
-        basename = fullname.strip('.shp')
-        tmp.close()
-        
-        # Todo: contruct download name dynamically from queryset
-        download_name = 'shp_download'
-        
+        """
+        """
         fields = self.queryset.model._meta.fields
         geo_fields = [f for f in fields if isinstance(f, GeometryField)]
-        other_fields = [f for f in fields if not isinstance(f, GeometryField)]
+        attributes = [f for f in fields if not isinstance(f, GeometryField)]
         
         if len(geo_fields) > 1:
             geo_field = geo_fields[0] # no support yet for multiple geometry fields
         else:
             geo_field = geo_fields[0]
         
-        dr = lgdal.OGRGetDriverByName('ESRI Shapefile')
-    
+        ogr.OGRRegisterAll()
+        # Get the shapefile driver
+        dr = ogr.OGRGetDriverByName('ESRI Shapefile')
+        
+        # create a temporary file to write the shapefile to
+        # since we are ultimately going to zip it up
+        tmp = tempfile.NamedTemporaryFile(suffix='.shp', mode = 'w')
+        tmp_shp_name = tmp.name.split('/')[-1]
+        basename = tmp.name.strip('.shp')
+        tmp.close()
+        
         # Creating the datasource
-        ds = lgdal.OGR_Dr_CreateDataSource(dr,fullname, None)
+        ds = ogr.OGR_Dr_CreateDataSource(dr,tmp.name, None)
         
         # Get the right geometry type number for ogr
         ogr_type = OGRGeomType(geo_field._geom).num
 
-        # Set up the spatial reference with epsg cod
+        # Set up the spatial reference with epsg code
         srs = SpatialReference(geo_field._srid)
-        # we're going to reproject later on 
+        
+        # If true we're going to reproject later on 
         if self.proj_transform:
           srs = SpatialReference(self.proj_transform)
         
         # Creating the layer
-        layer = lgdal.OGR_DS_CreateLayer(ds,name, srs._ptr, ogr_type, None)
-        #import pdb;pdb.set_trace()
+        layer = ogr.OGR_DS_CreateLayer(ds,tmp_shp_name, srs._ptr, ogr_type, None)
         
-        for field in other_fields:
-          fld = lgdal.OGR_Fld_Create(str(field.name), 4)
-          added = lgdal.OGR_L_CreateField(layer, fld, 0)
+        # Create the fields
+        # Todo: control field order as param
+        for field in attributes:
+          fld = ogr.OGR_Fld_Create(str(field.name), 4)
+          added = ogr.OGR_L_CreateField(layer, fld, 0)
           check_err(added) 
         
         # Getting the Layer feature definition.
-        fdefn = lgdal.OGR_L_GetLayerDefn(layer) 
+        feature_def = ogr.OGR_L_GetLayerDefn(layer) 
         
+        # Loop through queryset creating features
         for item in self.queryset:
-        
-            feat = lgdal.OGR_F_Create(fdefn)
+            feat = ogr.OGR_F_Create(feature_def)
             
-            # Setting the fields (for now all as strings)
+            # For now, set all fields as strings
             # TODO: catch model types and convert to ogr fields
             #OFTReal => FloatField DecimalField
             #OFTInteger => IntegerField
             #OFTDate => TimeField
             
             idx = 0
-            for field in other_fields:
+            for field in attributes:
               value = getattr(item,field.name)
-              lgdal.OGR_F_SetFieldString(feat, idx, str(value))
+              ogr.OGR_F_SetFieldString(feat, idx, str(value))
               idx += 1
               
             # Transforming & setting the geometry
             geom = getattr(item,geo_field.name)
             
-            # reproject is requested so we transform the input
-            # geometry to match the shapefiles projection 'to-be'
+            # if requested we transform the input geometry
+            # to match the shapefiles projection 'to-be'
             if self.proj_transform:
               geom.transform(self.proj_transform)
             ogr_geom = OGRGeometry(geom.wkt,srs)
-            #if self.proj_transform:
-              #ogr_geom.transform(self.proj_transform)
-            check_err(lgdal.OGR_F_SetGeometry(feat, ogr_geom._ptr))
-            # Creating the feature in the layer.
-            check_err(lgdal.OGR_L_SetFeature(layer, feat))
+            
+            # create the geometry
+            check_err(ogr.OGR_F_SetGeometry(feat, ogr_geom._ptr))
+            
+            # creat the feature in the layer.
+            check_err(ogr.OGR_L_SetFeature(layer, feat))
         
         # Cleaning up
-        check_err(lgdal.OGR_L_SyncToDisk(layer))
-        lgdal.OGR_DS_Destroy(ds)
-    
-        buffer = StringIO()
+        check_err(ogr.OGR_L_SyncToDisk(layer))
+        ogr.OGR_DS_Destroy(ds)
+        ogr.OGRCleanupAll()
+        
+        #Read resulting shapefile into a zipfile buffer
+        buffer = cStringIO.StringIO()
         zip = zipfile.ZipFile(buffer, 'w', zipfile.ZIP_DEFLATED)
         files = ['shp','shx','prj','dbf']
         for item in files:
           filename= '%s.%s' % (basename, item)
-          zip.write(filename, '%s.%s' % (download_name, item))
+          zip.write(filename, '%s.%s' % (self.file_name, item))
         zip.close()
         buffer.flush()
-        ret_zip = buffer.getvalue()
-        buffer.close()     
+        zip_stream = buffer.getvalue()
+        buffer.close()
+        
+        # Stick it all in a django HttpResponse
         response = HttpResponse()
-        response['Content-Disposition'] = 'filename=%s.zip' % download_name
-        response['Content-length'] = str(len(ret_zip))
+        response['Content-Disposition'] = 'filename=%s.zip' % self.file_name
+        response['Content-length'] = str(len(zip_stream))
         response['Content-Type'] = self.mimetype
-    
-        response.write(ret_zip)
-        lgdal.OGRCleanupAll()
-        return response
+        response.write(zip_stream)
+        return response