Commits

Anonymous committed 483b043

Add support for multiple geometry fields, add option to generate a README file, handle null geometries

Comments (0)

Files changed (1)

shapes/views/export.py

 # todo: support multiple querysets == multiple shapefiles
 
 class ShpResponder(object):
-    def __init__(self, queryset, *args, **kwargs):
+    def __init__(self, queryset, readme=None, geo_field=None, proj_transform=None, mimetype='application/zip',file_name='shp_download'):
         """
         """
         self.queryset = queryset
-        self.proj_transform = None
-        self.mimetype = 'application/zip'
-        self.file_name = 'shp_download'
-        self.file_name = self.file_name.rstrip('.shp')
+        self.readme = readme
+        self.geo_field = geo_field
+        self.proj_transform = proj_transform
+        self.mimetype = mimetype
+        self.file_name = file_name
     
-    def __call__(self, *args, **kwargs):
+    def __call__(self):
         """
         """
         fields = self.queryset.model._meta.fields
         geo_fields = [f for f in fields if isinstance(f, GeometryField)]
+        geo_fields_names = ', '.join([f.name for f in geo_fields])
         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
+            if not self.geo_field:
+              raise ValueError("More than one geodjango geometry field found, please specify which to use by name. Available fields are: '%s'" % geo_fields_names)
+            else:
+              geo_field_by_name = [fld for fld in geo_fields if fld.name == self.geo_field]
+              if not geo_field_by_name:
+                raise ValueError("Geodjango geometry field not found with the name '%s', fields available are: '%s'" % (self.geo_field,geo_fields_names))
+              else:
+                geo_field = geo_field_by_name[0]
+        elif geo_fields:
+            geo_field = geo_fields[0]
         else:
-            geo_field = geo_fields[0]
+            raise ValueError('No geodjango geometry fields found in this model queryset')
         
         ogr.OGRRegisterAll()
         # Get the shapefile driver
         
         # 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 = tempfile.NamedTemporaryFile(suffix='.shp', mode = 'w+b')
+        # we must close the file for GDAL to be able to open and write to it
         tmp.close()
         
         # Creating the datasource
           srs = SpatialReference(self.proj_transform)
         
         # Creating the layer
-        layer = ogr.OGR_DS_CreateLayer(ds,tmp_shp_name, srs._ptr, ogr_type, None)
+        layer = ogr.OGR_DS_CreateLayer(ds,os.path.basename(tmp.name), srs._ptr, ogr_type, None)
         
         # Create the fields
         # Todo: control field order as param
             
             # For now, set all fields as strings
             # TODO: catch model types and convert to ogr fields
+            # http://www.gdal.org/ogr/classOGRFeature.html
+            
+            # OGR_F_SetFieldDouble
             #OFTReal => FloatField DecimalField
+            
+            # OGR_F_SetFieldInteger
             #OFTInteger => IntegerField
+            
+            #OGR_F_SetFieldStrin
             #OFTString => CharField
-            #OFTDate => DateField
+            
+            
+            # OGR_F_SetFieldDateTime()
             #OFTDateTime => DateTimeField
             #OFTDate => TimeField
+            #OFTDate => DateField
             
             idx = 0
             for field in attributes:
               
             # Transforming & setting the geometry
             geom = getattr(item,geo_field.name)
-            
+            #import pdb;pdb.set_trace()
             # 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)
             
-            # create the geometry
-            check_err(ogr.OGR_F_SetGeometry(feat, ogr_geom._ptr))
+            if geom:
+              if self.proj_transform:
+                geom.transform(self.proj_transform)
+              ogr_geom = OGRGeometry(geom.wkt,srs)
+              # create the geometry
+              check_err(ogr.OGR_F_SetGeometry(feat, ogr_geom._ptr))
+            else:
+              # Case where geometry object is not found because of null value for field
+              # effectively looses whole record in shapefile if geometry does not exist
+              pass
+            
             
             # creat the feature in the layer.
             check_err(ogr.OGR_L_SetFeature(layer, feat))
         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' % (self.file_name, item))
+          filename= '%s.%s' % (tmp.name.strip('.shp'), item)
+          zip.write(filename, arcname='%s.%s' % (self.file_name.rstrip('.shp'), item))
+        if self.readme:
+          zip.writestr('README.txt',self.readme)
         zip.close()
         buffer.flush()
         zip_stream = buffer.getvalue()
         
         # Stick it all in a django HttpResponse
         response = HttpResponse()
-        response['Content-Disposition'] = 'filename=%s.zip' % self.file_name
+        response['Content-Disposition'] = 'filename=%s.zip' % self.file_name.rstrip('.shp')
         response['Content-length'] = str(len(zip_stream))
         response['Content-Type'] = self.mimetype
         response.write(zip_stream)