1 import os, os.path, unittest

     2 from django.contrib.gis.gdal import DataSource, Envelope, OGRException, OGRIndexError

     3 from django.contrib.gis.gdal.field import OFTReal, OFTInteger, OFTString

     4 

     5 # Path for SHP files

     6 data_path = os.path.join(os.path.dirname(__file__), 'data')

     7 def get_ds_file(name, ext):

     8     return os.sep.join([data_path, name, name + '.%s' % ext])

     9 

    10 # Test SHP data source object

    11 class TestDS:

    12     def __init__(self, name, **kwargs):

    13         ext = kwargs.pop('ext', 'shp')

    14         self.ds = get_ds_file(name, ext)

    15         for key, value in kwargs.items():

    16             setattr(self, key, value)

    17 

    18 # List of acceptable data sources.

    19 ds_list = (TestDS('test_point', nfeat=5, nfld=3, geom='POINT', gtype=1, driver='ESRI Shapefile',

    20                   fields={'dbl' : OFTReal, 'int' : OFTInteger, 'str' : OFTString,},

    21                   extent=(-1.35011,0.166623,-0.524093,0.824508), # Got extent from QGIS

    22                   srs_wkt='GEOGCS["GCS_WGS_1984",DATUM["WGS_1984",SPHEROID["WGS_1984",6378137,298.257223563]],PRIMEM["Greenwich",0],UNIT["Degree",0.017453292519943295]]',

    23                   field_values={'dbl' : [float(i) for i in range(1, 6)], 'int' : range(1, 6), 'str' : [str(i) for i in range(1, 6)]},

    24                   fids=range(5)),

    25            TestDS('test_vrt', ext='vrt', nfeat=3, nfld=3, geom='POINT', gtype=1, driver='VRT',

    26                   fields={'POINT_X' : OFTString, 'POINT_Y' : OFTString, 'NUM' : OFTString}, # VRT uses CSV, which all types are OFTString.

    27                   extent=(1.0, 2.0, 100.0, 523.5), # Min/Max from CSV

    28                   field_values={'POINT_X' : ['1.0', '5.0', '100.0'], 'POINT_Y' : ['2.0', '23.0', '523.5'], 'NUM' : ['5', '17', '23']},

    29                   fids=range(1,4)),

    30            TestDS('test_poly', nfeat=3, nfld=3, geom='POLYGON', gtype=3, 

    31                   driver='ESRI Shapefile',

    32                   fields={'float' : OFTReal, 'int' : OFTInteger, 'str' : OFTString,},

    33                   extent=(-1.01513,-0.558245,0.161876,0.839637), # Got extent from QGIS

    34                   srs_wkt='GEOGCS["GCS_WGS_1984",DATUM["WGS_1984",SPHEROID["WGS_1984",6378137,298.257223563]],PRIMEM["Greenwich",0],UNIT["Degree",0.017453292519943295]]'),

    35            )

    36 

    37 bad_ds = (TestDS('foo'),

    38           )

    39 

    40 class DataSourceTest(unittest.TestCase):

    41 

    42     def test01_valid_shp(self):

    43         "Testing valid SHP Data Source files."

    44 

    45         for source in ds_list:

    46             # Loading up the data source

    47             ds = DataSource(source.ds)

    48 

    49             # Making sure the layer count is what's expected (only 1 layer in a SHP file)

    50             self.assertEqual(1, len(ds))

    51 

    52             # Making sure GetName works

    53             self.assertEqual(source.ds, ds.name)

    54 

    55             # Making sure the driver name matches up

    56             self.assertEqual(source.driver, str(ds.driver))

    57 

    58             # Making sure indexing works

    59             try:

    60                 ds[len(ds)]

    61             except OGRIndexError:

    62                 pass

    63             else:

    64                 self.fail('Expected an IndexError!')

    65                         

    66     def test02_invalid_shp(self):

    67         "Testing invalid SHP files for the Data Source."

    68         for source in bad_ds:

    69             self.assertRaises(OGRException, DataSource, source.ds)

    70 

    71     def test03a_layers(self):

    72         "Testing Data Source Layers."

    73         print "\nBEGIN - expecting out of range feature id error; safe to ignore.\n"

    74         for source in ds_list:

    75             ds = DataSource(source.ds)

    76 

    77             # Incrementing through each layer, this tests DataSource.__iter__

    78             for layer in ds:                

    79                 # Making sure we get the number of features we expect

    80                 self.assertEqual(len(layer), source.nfeat)

    81 

    82                 # Making sure we get the number of fields we expect

    83                 self.assertEqual(source.nfld, layer.num_fields)

    84                 self.assertEqual(source.nfld, len(layer.fields))

    85 

    86                 # Testing the layer's extent (an Envelope), and it's properties

    87                 self.assertEqual(True, isinstance(layer.extent, Envelope))

    88                 self.assertAlmostEqual(source.extent[0], layer.extent.min_x, 5)

    89                 self.assertAlmostEqual(source.extent[1], layer.extent.min_y, 5)

    90                 self.assertAlmostEqual(source.extent[2], layer.extent.max_x, 5)

    91                 self.assertAlmostEqual(source.extent[3], layer.extent.max_y, 5)

    92 

    93                 # Now checking the field names.

    94                 flds = layer.fields

    95                 for f in flds: self.assertEqual(True, f in source.fields)

    96                 

    97                 # Negative FIDs are not allowed.

    98                 self.assertRaises(OGRIndexError, layer.__getitem__, -1)

    99                 self.assertRaises(OGRIndexError, layer.__getitem__, 50000)

   100 

   101                 if hasattr(source, 'field_values'):

   102                     fld_names = source.field_values.keys()

   103 

   104                     # Testing `Layer.get_fields` (which uses Layer.__iter__)

   105                     for fld_name in fld_names:

   106                         self.assertEqual(source.field_values[fld_name], layer.get_fields(fld_name))

   107 

   108                     # Testing `Layer.__getitem__`.

   109                     for i, fid in enumerate(source.fids):

   110                         feat = layer[fid]

   111                         self.assertEqual(fid, feat.fid)

   112                         # Maybe this should be in the test below, but we might as well test

   113                         # the feature values here while in this loop.

   114                         for fld_name in fld_names:

   115                             self.assertEqual(source.field_values[fld_name][i], feat.get(fld_name))

   116         print "\nEND - expecting out of range feature id error; safe to ignore."

   117                         

   118     def test03b_layer_slice(self):

   119         "Test indexing and slicing on Layers."

   120         # Using the first data-source because the same slice

   121         # can be used for both the layer and the control values.

   122         source = ds_list[0]

   123         ds = DataSource(source.ds)

   124 

   125         sl = slice(1, 3)

   126         feats = ds[0][sl]

   127 

   128         for fld_name in ds[0].fields:

   129             test_vals = [feat.get(fld_name) for feat in feats]

   130             control_vals = source.field_values[fld_name][sl]

   131             self.assertEqual(control_vals, test_vals)

   132 

   133     def test04_features(self):

   134         "Testing Data Source Features."

   135         for source in ds_list:

   136             ds = DataSource(source.ds)

   137 

   138             # Incrementing through each layer

   139             for layer in ds:

   140                 # Incrementing through each feature in the layer

   141                 for feat in layer:

   142                     # Making sure the number of fields, and the geometry type

   143                     # are what's expected.

   144                     self.assertEqual(source.nfld, len(list(feat)))

   145                     self.assertEqual(source.gtype, feat.geom_type)

   146 

   147                     # Making sure the fields match to an appropriate OFT type.

   148                     for k, v in source.fields.items():

   149                         # Making sure we get the proper OGR Field instance, using

   150                         # a string value index for the feature.

   151                         self.assertEqual(True, isinstance(feat[k], v))

   152 

   153                     # Testing Feature.__iter__

   154                     for fld in feat: self.assertEqual(True, fld.name in source.fields.keys())

   155                         

   156     def test05_geometries(self):

   157         "Testing Geometries from Data Source Features."

   158         for source in ds_list:

   159             ds = DataSource(source.ds)

   160 

   161             # Incrementing through each layer and feature.

   162             for layer in ds:

   163                 for feat in layer:

   164                     g = feat.geom

   165 

   166                     # Making sure we get the right Geometry name & type

   167                     self.assertEqual(source.geom, g.geom_name)

   168                     self.assertEqual(source.gtype, g.geom_type)

   169 

   170                     # Making sure the SpatialReference is as expected.

   171                     if hasattr(source, 'srs_wkt'):

   172                         self.assertEqual(source.srs_wkt, g.srs.wkt)

   173 

   174 

   175 def suite():

   176     s = unittest.TestSuite()

   177     s.addTest(unittest.makeSuite(DataSourceTest))

   178     return s

   179 

   180 def run(verbosity=2):

   181     unittest.TextTestRunner(verbosity=verbosity).run(suite())