# LayerMapping -- A Django Model/OGR Layer Mapping Utility

"""

 The LayerMapping class provides a way to map the contents of OGR

 vector files (e.g. SHP files) to Geographic-enabled Django models.

"""

import sys

from datetime import date, datetime

from decimal import Decimal

from django.core.exceptions import ObjectDoesNotExist

from django.contrib.gis.db.models import GeometryField

from django.contrib.gis.gdal import CoordTransform, DataSource, \

    OGRException, OGRGeometry, OGRGeomType, SpatialReference

from django.contrib.gis.gdal.field import \

    OFTDate, OFTDateTime, OFTInteger, OFTReal, OFTString, OFTTime

from django.db import models, transaction

from django.contrib.localflavor.us.models import USStateField

# LayerMapping exceptions.

class LayerMapError(Exception): pass

class InvalidString(LayerMapError): pass

class InvalidDecimal(LayerMapError): pass

class InvalidInteger(LayerMapError): pass

class MissingForeignKey(LayerMapError): pass

class LayerMapping(object):

    "A class that maps OGR Layers to GeoDjango Models."

    # Acceptable 'base' types for a multi-geometry type.

    MULTI_TYPES = {1 : OGRGeomType('MultiPoint'),

                   2 : OGRGeomType('MultiLineString'),

                   3 : OGRGeomType('MultiPolygon'),

                   }

    # Acceptable Django field types and corresponding acceptable OGR

    # counterparts.

    FIELD_TYPES = {

        models.AutoField : OFTInteger,

        models.IntegerField : (OFTInteger, OFTReal, OFTString),

        models.FloatField : (OFTInteger, OFTReal),

        models.DateField : OFTDate,

        models.DateTimeField : OFTDateTime,

        models.EmailField : OFTString,

        models.TimeField : OFTTime,

        models.DecimalField : (OFTInteger, OFTReal),

        models.CharField : OFTString,

        models.SlugField : OFTString,

        models.TextField : OFTString,

        models.URLField : OFTString,

        USStateField : OFTString,

        models.XMLField : OFTString,

        models.SmallIntegerField : (OFTInteger, OFTReal, OFTString),

        models.PositiveSmallIntegerField : (OFTInteger, OFTReal, OFTString),

        }

    # The acceptable transaction modes.

    TRANSACTION_MODES = {'autocommit' : transaction.autocommit,

                         'commit_on_success' : transaction.commit_on_success,

                         }

                 source_srs=None, encoding=None,

        """

        A LayerMapping object is initialized using the given Model (not an instance),

        a DataSource (or string path to an OGR-supported data file), and a mapping

        dictionary.  See the module level docstring for more details and keyword

        argument usage.

        """

        # Getting the DataSource and the associated Layer.

        if isinstance(data, basestring):

            self.ds = DataSource(data)

        else:

            self.ds = data

        self.layer = self.ds[layer]

        # Setting the mapping & model attributes.

        self.mapping = mapping

        self.model = model

        # Checking the layer -- intitialization of the object will fail if

        # things don't check out before hand.

        self.check_layer()

        # exception will be raised if there is no geometry column).

            transform = False

        else:

        # 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 = self.coord_transform()

        else:

            self.transform = transform

        # Setting the encoding for OFTString fields, if specified.

        if encoding:

            # Making sure the encoding exists, if not a LookupError

            # exception will be thrown.

            from codecs import lookup

            lookup(encoding)

            self.encoding = encoding

        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

        # transaction modes dictionary.

        if transaction_mode in self.TRANSACTION_MODES:

            self.transaction_decorator = self.TRANSACTION_MODES[transaction_mode]

            self.transaction_mode = transaction_mode

        else:

            raise LayerMapError('Unrecognized transaction mode: %s' % transaction_mode)

    #### Checking routines used during initialization ####

    def check_fid_range(self, fid_range):

        "This checks the `fid_range` keyword."

        if fid_range:

            if isinstance(fid_range, (tuple, list)):

                return slice(*fid_range)

            elif isinstance(fid_range, slice):

                return fid_range

            else:

                raise TypeError

        else:

            return None

    def check_layer(self):

        """

        This checks the Layer metadata, and ensures that it is compatible

        with the mapping information and model.  Unlike previous revisions,

        there is no need to increment through each feature in the Layer.

        """

        # The geometry field of the model is set here.

        # TODO: Support more than one geometry field / model.  However, this

        # depends on the GDAL Driver in use.

        self.geom_field = False

        self.fields = {}

        # Getting lists of the field names and the field types available in

        # the OGR Layer.

        ogr_fields = self.layer.fields

        ogr_field_types = self.layer.field_types

        # Function for determining if the OGR mapping field is in the Layer.

        def check_ogr_fld(ogr_map_fld):

            try:

                idx = ogr_fields.index(ogr_map_fld)

            except ValueError:

                raise LayerMapError('Given mapping OGR field "%s" not found in OGR Layer.' % ogr_map_fld)

            return idx

        # No need to increment through each feature in the model, simply check

        # the Layer metadata against what was given in the mapping dictionary.

        for field_name, ogr_name in self.mapping.items():

            # Ensuring that a corresponding field exists in the model

            # for the given field name in the mapping.

            try:

                model_field = self.model._meta.get_field(field_name)

            except models.fields.FieldDoesNotExist:

                raise LayerMapError('Given mapping field "%s" not in given Model fields.' % field_name)

            # Getting the string name for the Django field class (e.g., 'PointField').

            fld_name = model_field.__class__.__name__

            if isinstance(model_field, GeometryField):

                if self.geom_field:

                    raise LayerMapError('LayerMapping does not support more than one GeometryField per model.')

                try:

                except OGRException:

                    raise LayerMapError('Invalid mapping for GeometryField "%s".' % field_name)

                # Making sure that the OGR Layer's Geometry is compatible.

                ltype = self.layer.geom_type

                # Setting the `geom_field` attribute w/the name of the model field

                self.geom_field = field_name

                fields_val = model_field

            elif isinstance(model_field, models.ForeignKey):

                if isinstance(ogr_name, dict):

                    # Is every given related model mapping field in the Layer?

                    rel_model = model_field.rel.to

                        idx = check_ogr_fld(ogr_field)

                        try:

                            rel_field = rel_model._meta.get_field(rel_name)

                        except models.fields.FieldDoesNotExist:

                                                (rel_name, rel_model.__class__.__name__))

                    fields_val = rel_model

                else:

                    raise TypeError('ForeignKey mapping must be of dictionary type.')

            else:

                # Is the model field type supported by LayerMapping?

                if not model_field.__class__ in self.FIELD_TYPES:

                    raise LayerMapError('Django field type "%s" has no OGR mapping (yet).' % fld_name)

                # Is the OGR field in the Layer?

                idx = check_ogr_fld(ogr_name)

                ogr_field = ogr_field_types[idx]

                # Can the OGR field type be mapped to the Django field type?

                if not issubclass(ogr_field, self.FIELD_TYPES[model_field.__class__]):

                                        (ogr_field, ogr_field.__name__, fld_name))

                fields_val = model_field

            self.fields[field_name] = fields_val

    def check_srs(self, source_srs):

        "Checks the compatibility of the given spatial reference object."

        if isinstance(source_srs, SpatialReference):

            sr = source_srs

            sr = source_srs.srs

        elif isinstance(source_srs, (int, basestring)):

            sr = SpatialReference(source_srs)

        else:

            # Otherwise just pulling the SpatialReference from the layer

            sr = self.layer.srs

        if not sr:

            raise LayerMapError('No source reference system defined.')

        else:

            return sr

    def check_unique(self, unique):

        "Checks the `unique` keyword parameter -- may be a sequence or string."

        if isinstance(unique, (list, tuple)):

            # List of fields to determine uniqueness with

                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.')

    #### Keyword argument retrieval routines ####

    def feature_kwargs(self, feat):

        """

        Given an OGR Feature, this will return a dictionary of keyword arguments

        for constructing the mapped model.

        """

        # The keyword arguments for model construction.

        kwargs = {}

        # Incrementing through each model field and OGR field in the

        # dictionary mapping.

        for field_name, ogr_name in self.mapping.items():

            model_field = self.fields[field_name]

            if isinstance(model_field, GeometryField):

                # Verify OGR geometry.

                val = self.verify_geom(feat.geom, model_field)

            elif isinstance(model_field, models.base.ModelBase):

                # The related _model_, not a field was passed in -- indicating

                # another mapping for the related Model.

                val = self.verify_fk(feat, model_field, ogr_name)

            else:

                # Otherwise, verify OGR Field type.

                val = self.verify_ogr_field(feat[ogr_name], model_field)

            # Setting the keyword arguments for the field name with the

            # value obtained above.

            kwargs[field_name] = val

        return kwargs

    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)

    #### Verification routines used in constructing model keyword arguments. ####

    def verify_ogr_field(self, ogr_field, model_field):

        """

        Verifies if the OGR Field contents are acceptable to the Django

        otherwise the proper exception is raised.

        """

            if self.encoding:

                # The encoding for OGR data sources may be specified here

                # (e.g., 'cp437' for Census Bureau boundary files).

                val = unicode(ogr_field.value, self.encoding)

            else:

                val = ogr_field.value

                if len(val) > model_field.max_length:

                    raise InvalidString('%s model field maximum string length is %s, given %s characters.' %

                                        (model_field.name, model_field.max_length, len(val)))

        elif isinstance(ogr_field, OFTReal) and isinstance(model_field, models.DecimalField):

            try:

                # Creating an instance of the Decimal value to use.

                d = Decimal(str(ogr_field.value))

            except:

                raise InvalidDecimal('Could not construct decimal from: %s' % ogr_field.value)

            # Getting the decimal value as a tuple.

            dtup = d.as_tuple()

            digits = dtup[1]

            d_idx = dtup[2] # index where the decimal is

            # Maximum amount of precision, or digits to the left of the decimal.

            max_prec = model_field.max_digits - model_field.decimal_places

            # given decimal.

            if d_idx < 0:

                n_prec = len(digits[:d_idx])

            else:

                n_prec = len(digits) + d_idx

            # InvalidDecimal exception.

            if n_prec > max_prec:

                raise InvalidDecimal('A DecimalField with max_digits %d, decimal_places %d must round to an absolute value less than 10^%d.' %

                                     (model_field.max_digits, model_field.decimal_places, max_prec))

            val = d

        elif isinstance(ogr_field, (OFTReal, OFTString)) and isinstance(model_field, models.IntegerField):

            # Attempt to convert any OFTReal and OFTString value to an OFTInteger.

            try:

                val = int(ogr_field.value)

            except:

                raise InvalidInteger('Could not construct integer from: %s' % ogr_field.value)

        else:

            val = ogr_field.value

        return val

    def verify_fk(self, feat, rel_model, rel_mapping):

        """

        Given an OGR Feature, the related model and its dictionary mapping,

        this routine will retrieve the related model for the ForeignKey

        mapping.

        """

        # TODO: It is expensive to retrieve a model for every record --

        #  ForeignKey models.

        # Constructing and verifying the related model keyword arguments.

        fk_kwargs = {}

        for field_name, ogr_name in rel_mapping.items():

            fk_kwargs[field_name] = self.verify_ogr_field(feat[ogr_name], rel_model._meta.get_field(field_name))

        # Attempting to retrieve and return the related model.

        try:

            return rel_model.objects.get(**fk_kwargs)

        except ObjectDoesNotExist:

            raise MissingForeignKey('No ForeignKey %s model found with keyword arguments: %s' % (rel_model.__name__, fk_kwargs))

    def verify_geom(self, geom, model_field):

        """

        Verifies the geometry -- will construct and return a GeometryCollection

        if necessary (for example if the model field is MultiPolygonField while

        the mapped shapefile only contains Polygons).

        """