1 import os, unittest

     2 from copy import copy

     3 from datetime import date

     4 from decimal import Decimal

     5 from models import City, County, CountyFeat, Interstate, State, city_mapping, co_mapping, cofeat_mapping, inter_mapping

     6 from django.contrib.gis.utils.layermapping import LayerMapping, LayerMapError, InvalidDecimal, MissingForeignKey

     7 from django.contrib.gis.gdal import DataSource

     8 

     9 shp_path = os.path.dirname(__file__)

    10 city_shp = os.path.join(shp_path, 'cities/cities.shp')

    11 co_shp = os.path.join(shp_path, 'counties/counties.shp')

    12 inter_shp = os.path.join(shp_path, 'interstates/interstates.shp')

    13 

    14 # Dictionaries to hold what's expected in the county shapefile.  

    15 NAMES  = ['Bexar', 'Galveston', 'Harris', 'Honolulu', 'Pueblo']

    16 NUMS   = [1, 2, 1, 19, 1] # Number of polygons for each.                                                                                                                                                  

    17 STATES = ['Texas', 'Texas', 'Texas', 'Hawaii', 'Colorado']

    18 

    19 class LayerMapTest(unittest.TestCase):

    20 

    21     def test01_init(self):

    22         "Testing LayerMapping initialization."

    23 

    24         # Model field that does not exist.

    25         bad1 = copy(city_mapping)

    26         bad1['foobar'] = 'FooField'

    27 

    28         # Shapefile field that does not exist.

    29         bad2 = copy(city_mapping)

    30         bad2['name'] = 'Nombre'

    31 

    32         # Nonexistent geographic field type.

    33         bad3 = copy(city_mapping)

    34         bad3['point'] = 'CURVE'

    35 

    36         # Incrementing through the bad mapping dictionaries and

    37         # ensuring that a LayerMapError is raised.

    38         for bad_map in (bad1, bad2, bad3):

    39             try:

    40                 lm = LayerMapping(City, city_shp, bad_map)

    41             except LayerMapError:

    42                 pass

    43             else:

    44                 self.fail('Expected a LayerMapError.')

    45 

    46         # A LookupError should be thrown for bogus encodings.

    47         try:

    48             lm = LayerMapping(City, city_shp, city_mapping, encoding='foobar')

    49         except LookupError:

    50             pass

    51         else:

    52             self.fail('Expected a LookupError')

    53 

    54     def test02_simple_layermap(self):

    55         "Test LayerMapping import of a simple point shapefile."

    56 

    57         # Setting up for the LayerMapping.

    58         lm = LayerMapping(City, city_shp, city_mapping)

    59         lm.save()

    60 

    61         # There should be three cities in the shape file.

    62         self.assertEqual(3, City.objects.count())

    63 

    64         # Opening up the shapefile, and verifying the values in each

    65         # of the features made it to the model.

    66         ds = DataSource(city_shp)

    67         layer = ds[0]

    68         for feat in layer:

    69             city = City.objects.get(name=feat['Name'].value)

    70             self.assertEqual(feat['Population'].value, city.population)

    71             self.assertEqual(Decimal(str(feat['Density'])), city.density)

    72             self.assertEqual(feat['Created'].value, city.dt)

    73 

    74             # Comparing the geometries.

    75             pnt1, pnt2 = feat.geom, city.point

    76             self.assertAlmostEqual(pnt1.x, pnt2.x, 6)

    77             self.assertAlmostEqual(pnt1.y, pnt2.y, 6)

    78 

    79     def test03_layermap_strict(self):

    80         "Testing the `strict` keyword, and import of a LineString shapefile."

    81 

    82         # When the `strict` keyword is set an error encountered will force

    83         # the importation to stop.

    84         try:

    85             lm = LayerMapping(Interstate, inter_shp, inter_mapping)

    86             lm.save(silent=True, strict=True)

    87         except InvalidDecimal:

    88             pass

    89         else:

    90             self.fail('Should have failed on strict import with invalid decimal values.')

    91 

    92         # This LayerMapping should work b/c `strict` is not set.

    93         lm = LayerMapping(Interstate, inter_shp, inter_mapping)

    94         lm.save(silent=True)

    95 

    96         # Two interstate should have imported correctly.

    97         self.assertEqual(2, Interstate.objects.count())

    98         

    99         # Verifying the values in the layer w/the model.

   100         ds = DataSource(inter_shp)

   101 

   102         # Only the first two features of this shapefile are valid.

   103         valid_feats = ds[0][:2]

   104         for feat in valid_feats:

   105             istate = Interstate.objects.get(name=feat['Name'].value)

   106             

   107             if feat.fid == 0:

   108                 self.assertEqual(Decimal(str(feat['Length'])), istate.length)

   109             elif feat.fid == 1:

   110                 # Everything but the first two decimal digits were truncated,

   111                 # because the Interstate model's `length` field has decimal_places=2.

   112                 self.assertAlmostEqual(feat.get('Length'), float(istate.length), 2)

   113 

   114             for p1, p2 in zip(feat.geom, istate.path):

   115                 self.assertAlmostEqual(p1[0], p2[0], 6)

   116                 self.assertAlmostEqual(p1[1], p2[1], 6)

   117 

   118     def county_helper(self, county_feat=True):

   119         "Helper function for ensuring the integrity of the mapped County models."

   120         

   121         for name, n, st in zip(NAMES, NUMS, STATES):

   122             # Should only be one record b/c of `unique` keyword.

   123             c = County.objects.get(name=name)

   124             self.assertEqual(n, len(c.mpoly))

   125             self.assertEqual(st, c.state.name) # Checking ForeignKey mapping.

   126             

   127             # Multiple records because `unique` was not set.

   128             if county_feat:

   129                 qs = CountyFeat.objects.filter(name=name)

   130                 self.assertEqual(n, qs.count())

   131 

   132     def test04_layermap_unique_multigeometry_fk(self):

   133         "Testing the `unique`, and `transform`, geometry collection conversion, and ForeignKey mappings."

   134         # All the following should work.

   135         try:

   136             # Telling LayerMapping that we want no transformations performed on the data.

   137             lm = LayerMapping(County, co_shp, co_mapping, transform=False)

   138         

   139             # Specifying the source spatial reference system via the `source_srs` keyword.

   140             lm = LayerMapping(County, co_shp, co_mapping, source_srs=4269)

   141             lm = LayerMapping(County, co_shp, co_mapping, source_srs='NAD83')

   142 

   143             # Unique may take tuple or string parameters.

   144             for arg in ('name', ('name', 'mpoly')):

   145                 lm = LayerMapping(County, co_shp, co_mapping, transform=False, unique=arg)

   146         except:

   147             self.fail('No exception should be raised for proper use of keywords.')

   148             

   149         # Testing invalid params for the `unique` keyword.

   150         for e, arg in ((TypeError, 5.0), (ValueError, 'foobar'), (ValueError, ('name', 'mpolygon'))):

   151             self.assertRaises(e, LayerMapping, County, co_shp, co_mapping, transform=False, unique=arg)

   152 

   153         # No source reference system defined in the shapefile, should raise an error.

   154         self.assertRaises(LayerMapError, LayerMapping, County, co_shp, co_mapping)

   155 

   156         # Passing in invalid ForeignKey mapping parameters -- must be a dictionary

   157         # mapping for the model the ForeignKey points to.

   158         bad_fk_map1 = copy(co_mapping); bad_fk_map1['state'] = 'name'

   159         bad_fk_map2 = copy(co_mapping); bad_fk_map2['state'] = {'nombre' : 'State'}

   160         self.assertRaises(TypeError, LayerMapping, County, co_shp, bad_fk_map1, transform=False)

   161         self.assertRaises(LayerMapError, LayerMapping, County, co_shp, bad_fk_map2, transform=False)

   162 

   163         # There exist no State models for the ForeignKey mapping to work -- should raise

   164         # a MissingForeignKey exception (this error would be ignored if the `strict`

   165         # keyword is not set).

   166         lm = LayerMapping(County, co_shp, co_mapping, transform=False, unique='name')

   167         self.assertRaises(MissingForeignKey, lm.save, silent=True, strict=True)

   168 

   169         # Now creating the state models so the ForeignKey mapping may work.

   170         co, hi, tx = State(name='Colorado'), State(name='Hawaii'), State(name='Texas')

   171         co.save(), hi.save(), tx.save()

   172 

   173         # If a mapping is specified as a collection, all OGR fields that

   174         # are not collections will be converted into them.  For example,

   175         # a Point column would be converted to MultiPoint. Other things being done

   176         # w/the keyword args:

   177         #  `transform=False`: Specifies that no transform is to be done; this 

   178         #    has the effect of ignoring the spatial reference check (because the

   179         #    county shapefile does not have implicit spatial reference info).

   180         # 

   181         #  `unique='name'`: Creates models on the condition that they have 

   182         #    unique county names; geometries from each feature however will be

   183         #    appended to the geometry collection of the unique model.  Thus,

   184         #    all of the various islands in Honolulu county will be in in one

   185         #    database record with a MULTIPOLYGON type.

   186         lm = LayerMapping(County, co_shp, co_mapping, transform=False, unique='name')

   187         lm.save(silent=True, strict=True)

   188 

   189         # A reference that doesn't use the unique keyword; a new database record will

   190         # created for each polygon.

   191         lm = LayerMapping(CountyFeat, co_shp, cofeat_mapping, transform=False)

   192         lm.save(silent=True, strict=True)

   193 

   194         # The county helper is called to ensure integrity of County models.

   195         self.county_helper()

   196 

   197     def test05_test_fid_range_step(self):

   198         "Tests the `fid_range` keyword and the `step` keyword of .save()."

   199         

   200         # Function for clearing out all the counties before testing.

   201         def clear_counties(): County.objects.all().delete()

   202         

   203         # Initializing the LayerMapping object to use in these tests.

   204         lm = LayerMapping(County, co_shp, co_mapping, transform=False, unique='name')

   205 

   206         # Bad feature id ranges should raise a type error.

   207         clear_counties()

   208         bad_ranges = (5.0, 'foo', co_shp)

   209         for bad in bad_ranges:

   210             self.assertRaises(TypeError, lm.save, fid_range=bad)

   211 

   212         # Step keyword should not be allowed w/`fid_range`.

   213         fr = (3, 5) # layer[3:5]

   214         self.assertRaises(LayerMapError, lm.save, fid_range=fr, step=10) 

   215         lm.save(fid_range=fr)

   216         

   217         # Features IDs 3 & 4 are for Galveston County, Texas -- only

   218         # one model is returned because the `unique` keyword was set.

   219         qs = County.objects.all()

   220         self.assertEqual(1, qs.count())

   221         self.assertEqual('Galveston', qs[0].name)

   222 

   223         # Features IDs 5 and beyond for Honolulu County, Hawaii, and

   224         # FID 0 is for Pueblo County, Colorado.

   225         clear_counties()

   226         lm.save(fid_range=slice(5, None), silent=True, strict=True) # layer[5:]

   227         lm.save(fid_range=slice(None, 1), silent=True, strict=True) # layer[:1]

   228 

   229         # Only Pueblo & Honolulu counties should be present because of

   230         # the `unique` keyword.

   231         qs = County.objects.all()

   232         self.assertEqual(2, qs.count())

   233         hi, co = tuple(qs)

   234         hi_idx, co_idx = tuple(map(NAMES.index, ('Honolulu', 'Pueblo')))

   235         self.assertEqual('Pueblo', co.name); self.assertEqual(NUMS[co_idx], len(co.mpoly))

   236         self.assertEqual('Honolulu', hi.name); self.assertEqual(NUMS[hi_idx], len(hi.mpoly))

   237 

   238         # Testing the `step` keyword -- should get the same counties

   239         # regardless of we use a step that divides equally, that is odd,

   240         # or that is larger than the dataset.

   241         for st in (4,7,1000):

   242             clear_counties()

   243             lm.save(step=st, strict=True)

   244             self.county_helper(county_feat=False)

   245 

   246 def suite():

   247     s = unittest.TestSuite()

   248     s.addTest(unittest.makeSuite(LayerMapTest))

   249     return s