gis: LayerMapping: Added the unique keyword parameter and tests for the transform keyword and geometry collection conversion.

git-svn-id: http://code.djangoproject.com/svn/django/branches/gis@6980 bcc190cf-cafb-0310-a4f2-bffc1f526a37

django/contrib/gis/tests/layermap/models.py

@@ -1,5 +1,15 @@
 from django.contrib.gis.db import models
 
+class County(models.Model):
+    name = models.CharField(max_length=25)
+    mpoly = models.MultiPolygonField(srid=4269) # Multipolygon in NAD83
+    objects = models.GeoManager()
+
+class CountyFeat(models.Model):
+    name = models.CharField(max_length=25)
+    poly = models.PolygonField(srid=4269)
+    objects = models.GeoManager()
+
 class City(models.Model):
     name = models.CharField(max_length=25)
     population = models.IntegerField()
@@ -14,7 +24,15 @@ class Interstate(models.Model):
     path = models.LineStringField()
     objects = models.GeoManager()
 
-# Mapping dictionary for the City model.
+# Mapping dictionaries for the models above.
+co_mapping = {'name' : 'Name',
+              'mpoly' : 'MULTIPOLYGON', # Will convert POLYGON features into MULTIPOLYGONS.
+              }
+
+cofeat_mapping = {'name' : 'Name',
+                  'poly' : 'POLYGON',
+                  }
+
 city_mapping = {'name' : 'Name',
                 'population' : 'Population',
                 'density' : 'Density',
@@ -22,7 +40,6 @@ city_mapping = {'name' : 'Name',
                 'point' : 'POINT',
                 }
 
-# Mapping dictionary for the Interstate model.
 inter_mapping = {'name' : 'Name',
                  'length' : 'Length',
                  'path' : 'LINESTRING',

django/contrib/gis/tests/layermap/tests.py

@@ -2,12 +2,13 @@ import os, unittest
 from copy import copy
 from datetime import date
 from decimal import Decimal
-from models import City, Interstate, city_mapping, inter_mapping
+from models import City, County, CountyFeat, Interstate, city_mapping, co_mapping, cofeat_mapping, inter_mapping
 from django.contrib.gis.utils.layermapping import LayerMapping, LayerMapError, InvalidDecimal
 from django.contrib.gis.gdal import DataSource
 
 shp_path = os.path.dirname(__file__)
 city_shp = os.path.join(shp_path, 'cities/cities.shp')
+co_shp = os.path.join(shp_path, 'counties/counties.shp')
 inter_shp = os.path.join(shp_path, 'interstates/interstates.shp')
 
 class LayerMapTest(unittest.TestCase):
@@ -110,6 +111,61 @@ class LayerMapTest(unittest.TestCase):
                 self.assertAlmostEqual(p1[0], p2[0], 6)
                 self.assertAlmostEqual(p1[1], p2[1], 6)
 
+    def test04_layermap_unique_multigeometry(self):
+        "Testing the `unique`, and `transform` keywords and geometry collection conversion."
+        # All the following should work.
+        try:
+            # Telling LayerMapping that we want no transformations performed on the data.
+            lm = LayerMapping(County, co_shp, co_mapping, transform=False)
+        
+            # Specifying the source spatial reference system via the `source_srs` keyword.
+            lm = LayerMapping(County, co_shp, co_mapping, source_srs=4269)
+            lm = LayerMapping(County, co_shp, co_mapping, source_srs='NAD83')
+
+            # Unique may take tuple or string parameters.
+            for arg in ('name', ('name', 'mpoly')):
+                lm = LayerMapping(County, co_shp, co_mapping, transform=False, unique=arg)
+        except:
+            self.fail('No exception should be raised for proper use of keywords.')
+            
+        # Testing invalid params for the `unique` keyword.
+        for e, arg in ((TypeError, 5.0), (ValueError, 'foobar'), (ValueError, ('name', 'mpolygon'))):
+            self.assertRaises(e, LayerMapping, County, co_shp, co_mapping, transform=False, unique=arg)
+
+        # No source reference system defined in the shapefile, should raise an error.
+        self.assertRaises(LayerMapError, LayerMapping, County, co_shp, co_mapping)
+
+        # If a mapping is specified as a collection, all OGR fields that
+        # are not collections will be converted into them.  For example,
+        # a Point column would be converted to MultiPoint. Other things being done
+        # w/the keyword args:
+        #  `transform=False`: Specifies that no transform is to be done; this 
+        #    has the effect of ignoring the spatial reference check (because the
+        #    county shapefile does not have implicit spatial reference info).
+        # 
+        #  `unique='name'`: Creates models on the condition that they have 
+        #    unique county names; geometries from each feature however will be
+        #    appended to the geometry collection of the unique model.  Thus,
+        #    all of the various islands in Honolulu county will be in in one
+        #    database record with a MULTIPOLYGON type.
+        lm = LayerMapping(County, co_shp, co_mapping, transform=False, unique='name', silent=True)
+        lm.save()
+
+        # A reference that doesn't use the unique keyword; a new database record will
+        # created for each polygon.
+        lm = LayerMapping(CountyFeat, co_shp, cofeat_mapping, transform=False, silent=True)
+        lm.save()
+
+        # Dictionary to hold what's expected in the shapefile.
+        exp = {'names' : ('Bexar', 'Galveston', 'Harris', 'Honolulu', 'Pueblo'),
+               'num' : (1, 2, 2, 19, 1), # Number of polygons for each.
+               }
+        for name, n in zip(exp['names'], exp['num']):
+            c = County.objects.get(name=name) # Should only be one record.
+            self.assertEqual(n, len(c.mpoly))
+            qs = CountyFeat.objects.filter(name=name)
+            self.assertEqual(n, qs.count())
+            
 def suite():
     s = unittest.TestSuite()
     s.addTest(unittest.makeSuite(LayerMapTest))

django/contrib/gis/utils/layermapping.py

@@ -64,6 +64,13 @@
   transform:
    Setting this to False will disable all coordinate transformations.  
 
+  unique:
+   Setting this to the name, or a tuple of names, from the given
+   model will create models unique only to the given name(s).
+   Geometries will from each feature will be added into the collection
+   associated with the unique model.  Forces transaction mode to
+   be 'autocommit'.
+
 Example:
 
  1. You need a GDAL-supported data source, like a shapefile.
@@ -182,7 +189,8 @@ class LayerMapping(object):
     def __init__(self, model, data, mapping, layer=0, 
                  source_srs=None, encoding=None, check=True, 
                  progress=False, interval=1000, strict=False, silent=False,
-                 transaction_mode='commit_on_success', transform=True):
+                 transaction_mode='commit_on_success', transform=True,
+                 unique=False):
         "Takes the Django model, the data source, and the mapping (dictionary)"
 
         # Getting the field names and types from the model
@@ -208,8 +216,11 @@ class LayerMapping(object):
         # Checking the source spatial reference system, and getting
         # the coordinate transformation object (unless the `transform`
         # keyword is set to False)
+        if transform:
             self.source_srs = self.check_srs(source_srs)
-        self.transform = transform and self.coord_transform()
+            self.transform = self.coord_transform()
+        else:
+            self.transform = transform
 
         # Checking the layer -- intitialization of the object will fail if
         # things don't check out before hand.  This may be time-consuming,
@@ -232,6 +243,13 @@ class LayerMapping(object):
         else:
             self.encoding = None
 
+        if unique:
+            self.check_unique(unique)
+            transaction_mode = 'autocommit' # Has to be set to autocommit.
+            self.unique = unique
+        else:
+            self.unique = None
+
         # Setting the transaction decorator with the function in the 
         # transaction modes dictionary.
         if transaction_mode in self.TRANSACTION_MODES:
@@ -314,6 +332,26 @@ class LayerMapping(object):
         else:
             return sr
 
+    def check_unique(self, unique):
+        "Checks the `unique` keyword parameter -- may be a sequence or string."
+        # Getting the geometry field; only the first encountered GeometryField
+        # will be used.
+        self.geom_field = False
+        for model_field, ogr_fld in self.mapping.items():
+            if ogr_fld in self.OGC_TYPES:
+                self.geom_field = model_field
+                break
+
+        if isinstance(unique, (list, tuple)):
+            # List of fields to determine uniqueness with
+            for attr in unique: 
+                if not attr in self.mapping: raise ValueError
+        elif isinstance(unique, basestring):
+            # Only a single field passed in.
+            if unique not in self.mapping: raise ValueError
+        else:
+            raise TypeError('Unique keyword argument must be set with a tuple, list, or string.')
+
     def coord_transform(self):
         "Returns the coordinate transformation object."
         try:
@@ -371,6 +409,17 @@ class LayerMapping(object):
             
         return kwargs, all_prepped
 
+    def unique_kwargs(self, kwargs):
+        """
+        Given the feature keyword arguments (from `feature_kwargs`) this routine
+        will construct and return the uniqueness keyword arguments -- a subset
+        of the feature kwargs.
+        """
+        if isinstance(self.unique, basestring):
+            return {self.unique : kwargs[self.unique]}
+        else:
+            return dict((fld, kwargs[fld]) for fld in self.unique)
+
     def verify_field(self, fld, model_field):
         """
         Verifies if the OGR Field contents are acceptable to the Django
@@ -485,8 +534,31 @@ class LayerMapping(object):
                 else:
                     # Constructing the model using the keyword args
                     if all_prepped:
-                        m = self.model(**kwargs)
+                        if self.unique:
+                            # If we want unique models on a particular field, handle the
+                            # geometry appropriately.
                             try:
+                                # Getting the keyword arguments and retrieving
+                                # the unique model.
+                                u_kwargs = self.unique_kwargs(kwargs)
+                                m = self.model.objects.get(**u_kwargs)
+
+                                # Getting the geometry (in OGR form), creating 
+                                # one from the kwargs WKT, adding in additional 
+                                # geometries, and update the attribute with the 
+                                # just-updated geometry WKT.
+                                geom = getattr(m, self.geom_field).ogr
+                                new = OGRGeometry(kwargs[self.geom_field])
+                                for g in new: geom.add(g) 
+                                setattr(m, self.geom_field, geom.wkt)
+                            except ObjectDoesNotExist:
+                                # No unique model exists yet, create.
+                                m = self.model(**kwargs)
+                        else:
+                            m = self.model(**kwargs)
+
+                        try:
+                            # Attempting to save.
                             m.save()
                             num_saved += 1
                             if verbose: print 'Saved: %s' % m