gis LayerMapping now supports loading while mapping ForeignKey to other tables.

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

django/contrib/gis/utils/LayerMapping.py

@@ -20,7 +20,8 @@ Usage:
 
   model -- GeoDjango model (not an instance)
 
-  source_file -- OGR-supported data source file (e.g. a shapefile)
+  data -- OGR-supported data source file (e.g. a shapefile) or
+          gdal.DataSource instance
 
   mapping -- A python dictionary, keys are strings corresponding
              to the GeoDjango model field, and values correspond to
@@ -91,6 +92,9 @@ from django.contrib.gis.gdal import \
 from django.contrib.gis.gdal.Field import Field, OFTInteger, OFTReal, OFTString, OFTDateTime
 from django.contrib.gis.models import GeometryColumns, SpatialRefSys
 
+from django.db import connection, transaction
+from django.core.exceptions import ObjectDoesNotExist
+
 # A mapping of given geometry types to their OGR integer type.
 ogc_types = {'POINT' : OGRGeomType('Point'),
              'LINESTRING' : OGRGeomType('LineString'),
@@ -116,8 +120,21 @@ multi_types = {'POINT' : OGRGeomType('MultiPoint'),
                'POLYGON' : OGRGeomType('MultiPolygon'),
                }
 
+def map_foreign_key(django_field):
+    from django.db.models.fields.related import ForeignKey
+
+    if not django_field.__class__ is ForeignKey:
+        return django_field.__class__.__name__
+
+     
+    rf=django_field.rel.get_related_field()
+
+    return rf.get_internal_type()
+                
 # The acceptable Django field types that map to OGR fields.
-field_types = {'IntegerField' : OFTInteger,
+field_types = {
+               'AutoField' : OFTInteger,
+               'IntegerField' : OFTInteger,
                'FloatField' : OFTReal,
                'DateTimeField' : OFTDateTime,
                'DecimalField' : OFTReal,
@@ -142,10 +159,12 @@ def check_feature(feat, model_fields, mapping):
     for model_field, ogr_field in mapping.items():
 
         # Making sure the given mapping model field is in the given model fields.
-        if not model_field in model_fields:
-            raise Exception, 'Given mapping field "%s" not in given Model fields!' % model_field
-        else:
+        if model_field in model_fields:
             model_type = model_fields[model_field]
+        elif model_field[:-3] in model_fields: #foreign key
+            model_type = model_fields[model_field[:-3]]
+        else:
+            raise Exception, 'Given mapping field "%s" not in given Model fields!' % model_field
 
         ## Handling if we get a geometry in the Field ###
         if ogr_field in ogc_types:
@@ -208,14 +227,16 @@ def check_srs(layer, source_srs):
 class LayerMapping:
     "A class that maps OGR Layers to Django Models."
 
-    def __init__(self, model, ogr_file, mapping, layer=0, source_srs=None):
+    def __init__(self, model, data, mapping, layer=0, source_srs=None):
         "Takes the Django model, the mapping (dictionary), and the SHP file."
 
         # Getting the field names and types from the model
-        fields = dict((f.name, f.__class__.__name__) for f in model._meta.fields)
-
+        fields = dict((f.name, map_foreign_key(f)) for f in model._meta.fields)
         # Getting the DataSource and its Layer
-        self.ds = DataSource(ogr_file)
+        if isinstance(data, basestring):
+            self.ds = DataSource(data)
+        else:
+            self.ds = data
         self.layer = self.ds[layer]
 
         # Checking the layer -- intitialization of the object will fail if
@@ -228,6 +249,7 @@ class LayerMapping:
         self.model = model
         self.source_srs = check_srs(self.layer, source_srs)
 
+    @transaction.commit_on_success
     def save(self, verbose=False):
         "Runs the layer mapping on the given SHP file, and saves to the database."
 
@@ -252,8 +274,15 @@ class LayerMapping:
             kwargs = {}
 
             # Incrementing through each model field and the OGR field in the mapping
+            all_prepped = True
+
             for model_field, ogr_field in self.mapping.items():
+                is_fk = False
+                try:
                     model_type = self.fields[model_field]
+                except KeyError: #foreign key
+                    model_type = self.fields[model_field[:-3]]
+                    is_fk = True
 
                 if ogr_field in ogc_types:
                     ## Getting the OGR geometry from the field
@@ -271,17 +300,40 @@ class LayerMapping:
                     g.transform(ct)
 
                     # Updating the keyword args with the WKT of the transformed model.
-                    kwargs[model_field] = g.wkt
+                    val = g.wkt
                 else:
                     ## Otherwise, this is an OGR field type
                     fi = feat.index(ogr_field)
                     val = feat[fi].value
+
+                if is_fk:
+                    rel_obj = None
+                    field_name = model_field[:-3]
+                    try:
+                        #FIXME: refactor to efficiently fetch FKs.
+                        #  Requires significant re-work. :-/
+                        rel = self.model._meta.get_field(field_name).rel
+                        rel_obj = rel.to._default_manager.get(**{('%s__exact' % rel.field_name):val})
+                    except ObjectDoesNotExist:
+                        all_prepped = False
+                    
+                    kwargs[model_field[:-3]] = rel_obj
+                else:
                     kwargs[model_field] = val
 
             # Constructing the model using the constructed keyword args
+            if all_prepped:
                 m = self.model(**kwargs)
 
                 # Saving the model
+                try:
+                    if all_prepped:
                         m.save()
                         if verbose: print 'Saved: %s' % str(m)                        
+                    else:
+                        print "Skipping %s due to missing relation." % kwargs
+                except SystemExit:
+                    raise
+                except Exception, e:
+                    print "Failed to save %s\n  Continuing" % kwargs