Package osgeo :: Package utils :: Module ogrmerge |
[hide private]
[frames] | no frames]

# Source Code for Module osgeo.utils.ogrmerge

      1  #!/usr/bin/env python3
      2  # -*- coding: utf-8 -*-
      3  ###############################################################################
      4  # $Id$
      5  #
      6  # Project:  GDAL/OGR samples
      7  # Purpose:  Merge the content of several vector datasets into a single one.
      8  # Author:   Even Rouault <even dot rouault at spatialys dot com>
      9  #
     10  ###############################################################################
     11  # Copyright (c) 2017, Even Rouault <even dot rouault at spatialys dot com>
     12  #
     13  # Permission is hereby granted, free of charge, to any person obtaining a
     14  # copy of this software and associated documentation files (the "Software"),
     15  # to deal in the Software without restriction, including without limitation
     16  # the rights to use, copy, modify, merge, publish, distribute, sublicense,
     17  # and/or sell copies of the Software, and to permit persons to whom the
     18  # Software is furnished to do so, subject to the following conditions:
     19  #
     20  # The above copyright notice and this permission notice shall be included
     21  # in all copies or substantial portions of the Software.
     22  #
     23  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
     24  # OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
     25  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
     26  # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
     27  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
     28  # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
     29  # DEALINGS IN THE SOFTWARE.
     30  ###############################################################################
     31
     32  import glob
     33  import os
     34  import os.path
     35  import sys
     36
     37  from osgeo import gdal
     38  from osgeo import ogr
     39
     40  ###############################################################
     41  # Usage()
     42
     43
     44 -def Usage():
     45      print('ogrmerge.py -o out_dsname src_dsname [src_dsname]*')
     46      print('            [-f format] [-single] [-nln layer_name_template]')
     47      print('            [-update | -overwrite_ds] [-append | -overwrite_layer]')
     48      print('            [-src_geom_type geom_type_name[,geom_type_name]*]')
     49      print('            [-dsco NAME=VALUE]* [-lco NAME=VALUE]*')
     50      print('            [-s_srs srs_def] [-t_srs srs_def | -a_srs srs_def]')
     51      print('            [-progress] [-skipfailures] [--help-general]')
     52      print('')
     53      print('Options specific to -single:')
     54      print('            [-field_strategy FirstLayer|Union|Intersection]')
     55      print('            [-src_layer_field_name name]')
     56      print('            [-src_layer_field_content layer_name_template]')
     57      print('')
     58      print('* layer_name_template can contain the following substituable '
     59            'variables:')
     60      print('     {AUTO_NAME}  : {DS_BASENAME}_{LAYER_NAME} if they are '
     61            'different')
     62      print('                    or {LAYER_NAME} if they are identical')
     63      print('     {DS_NAME}    : name of the source dataset')
     64      print('     {DS_BASENAME}: base name of the source dataset')
     65      print('     {DS_INDEX}   : index of the source dataset')
     66      print('     {LAYER_NAME} : name of the source layer')
     67      print('     {LAYER_INDEX}: index of the source layer')
     68
     69      return 1
     70
     71
     72 -def DoesDriverHandleExtension(drv, ext):
     73      exts = drv.GetMetadataItem(gdal.DMD_EXTENSIONS)
     74      return exts is not None and exts.lower().find(ext.lower()) >= 0
     75
     76
     77 -def GetExtension(filename):
     78      if filename.lower().endswith('.shp.zip'):
     79          return 'shp.zip'
     80      ext = os.path.splitext(filename)[1]
     81      if ext.startswith('.'):
     82          ext = ext[1:]
     83      return ext
     84
     85
     86 -def GetOutputDriversFor(filename):
     87      drv_list = []
     88      ext = GetExtension(filename)
     89      if ext.lower() == 'vrt':
     90          return ['VRT']
     91      for i in range(gdal.GetDriverCount()):
     92          drv = gdal.GetDriver(i)
     93          if (drv.GetMetadataItem(gdal.DCAP_CREATE) is not None or
     94              drv.GetMetadataItem(gdal.DCAP_CREATECOPY) is not None) and \
     95             drv.GetMetadataItem(gdal.DCAP_VECTOR) is not None:
     96              if ext and DoesDriverHandleExtension(drv, ext):
     97                  drv_list.append(drv.ShortName)
     98              else:
     99                  prefix = drv.GetMetadataItem(gdal.DMD_CONNECTION_PREFIX)
    100                  if prefix is not None and filename.lower().startswith(prefix.lower()):
    101                      drv_list.append(drv.ShortName)
    102
    103      return drv_list
    104
    105
    106 -def GetOutputDriverFor(filename):
    107      drv_list = GetOutputDriversFor(filename)
    108      ext = GetExtension(filename)
    109      if not drv_list:
    110          if not ext:
    111              return 'ESRI Shapefile'
    112          else:
    113              raise Exception("Cannot guess driver for %s" % filename)
    114      elif len(drv_list) > 1:
    115          print("Several drivers matching %s extension. Using %s" % (ext if ext else '', drv_list[0]))
    116      return drv_list[0]
    117
    118  #############################################################################
    119
    120
    121 -def _VSIFPrintfL(f, s):
    122      gdal.VSIFWriteL(s, 1, len(s), f)
    123
    124  #############################################################################
    125
    126
    127 -def EQUAL(x, y):
    128      return x.lower() == y.lower()
    129
    130  #############################################################################
    131
    132
    133 -def _GetGeomType(src_geom_type_name):
    134      if EQUAL(src_geom_type_name, "GEOMETRY"):
    135          return ogr.wkbUnknown
    136      try:
    137          max_geom_type = ogr.wkbTriangle
    138      except:
    139          # GDAL 2.1 compat
    140          max_geom_type = ogr.wkbSurface
    141      for i in range(max_geom_type + 1):
    142          if EQUAL(src_geom_type_name,
    143                   ogr.GeometryTypeToName(i).replace(' ', '')):
    144              return i
    145      return None
    146
    147  #############################################################################
    148
    149
    150 -def _Esc(x):
    151      return gdal.EscapeString(x, gdal.CPLES_XML)
    152
    153
    154 -class XMLWriter(object):
    155
    156 -    def __init__(self, f):
    157          self.f = f
    158          self.inc = 0
    159          self.elements = []
    160
    161 -    def _indent(self):
    162          return '  ' * self.inc
    163
    164 -    def open_element(self, name, attrs=None):
    165          xml_attrs = ''
    166          if attrs is not None:
    167              for key in attrs:
    168                  xml_attrs = xml_attrs + ' %s=\"%s\"' % (key, _Esc(attrs[key].encode('utf-8')))
    169          x = '%s<%s%s>\n' % (self._indent(), name, xml_attrs)
    170          x = x.encode('utf-8')
    171          _VSIFPrintfL(self.f, x)
    172          self.inc = self.inc + 1
    173          self.elements.append(name)
    174
    175 -    def write_element_value(self, name, value, attrs=None):
    176          xml_attrs = ''
    177          if attrs is not None:
    178              for key in attrs:
    179                  xml_attrs = xml_attrs + ' %s=\"%s\"' % (key, _Esc(attrs[key].encode('utf-8')))
    180          x = '%s<%s%s>%s</%s>\n' % (self._indent(), name, xml_attrs,
    181                        _Esc(value.encode('utf-8')), name)
    182          x = x.encode('utf-8')
    183          _VSIFPrintfL(self.f, x)
    184
    185 -    def close_element(self, closing_name=None):
    186          self.inc = self.inc - 1
    187          name = self.elements[-1]
    188          if closing_name is not None:
    189              assert name == closing_name
    190          self.elements = self.elements[0:-1]
    191          _VSIFPrintfL(self.f, '%s</%s>\n' % (self._indent(), name))
    192
    193
    194  ###############################################################
    195  # process()
    196
    197
    198 -def process(argv, progress=None, progress_arg=None):
    199
    200      if not argv:
    201          return Usage()
    202
    203      dst_filename = None
    204      output_format = None
    205      src_datasets = []
    206      overwrite_ds = False
    207      overwrite_layer = False
    208      update = False
    209      append = False
    210      single_layer = False
    211      layer_name_template = None
    212      skip_failures = False
    213      src_geom_types = []
    214      field_strategy = None
    215      src_layer_field_name = None
    216      src_layer_field_content = None
    217      a_srs = None
    218      s_srs = None
    219      t_srs = None
    220      dsco = []
    221      lco = []
    222
    223      i = 0
    224      while i < len(argv):
    225          arg = argv[i]
    226          if (arg == '-f' or arg == '-of') and i + 1 < len(argv):
    227              i = i + 1
    228              output_format = argv[i]
    229          elif arg == '-o' and i + 1 < len(argv):
    230              i = i + 1
    231              dst_filename = argv[i]
    232          elif arg == '-progress':
    233              progress = ogr.TermProgress_nocb
    234              progress_arg = None
    235          elif arg == '-q' or arg == '-quiet':
    236              pass
    237          elif arg[0:5] == '-skip':
    238              skip_failures = True
    239          elif arg == '-update':
    240              update = True
    241          elif arg == '-overwrite_ds':
    242              overwrite_ds = True
    243          elif arg == '-overwrite_layer':
    244              overwrite_layer = True
    245              update = True
    246          elif arg == '-append':
    247              append = True
    248              update = True
    249          elif arg == '-single':
    250              single_layer = True
    251          elif arg == '-a_srs' and i + 1 < len(argv):
    252              i = i + 1
    253              a_srs = argv[i]
    254          elif arg == '-s_srs' and i + 1 < len(argv):
    255              i = i + 1
    256              s_srs = argv[i]
    257          elif arg == '-t_srs' and i + 1 < len(argv):
    258              i = i + 1
    259              t_srs = argv[i]
    260          elif arg == '-nln' and i + 1 < len(argv):
    261              i = i + 1
    262              layer_name_template = argv[i]
    263          elif arg == '-field_strategy' and i + 1 < len(argv):
    264              i = i + 1
    265              field_strategy = argv[i]
    266          elif arg == '-src_layer_field_name' and i + 1 < len(argv):
    267              i = i + 1
    268              src_layer_field_name = argv[i]
    269          elif arg == '-src_layer_field_content' and i + 1 < len(argv):
    270              i = i + 1
    271              src_layer_field_content = argv[i]
    272          elif arg == '-dsco' and i + 1 < len(argv):
    273              i = i + 1
    274              dsco.append(argv[i])
    275          elif arg == '-lco' and i + 1 < len(argv):
    276              i = i + 1
    277              lco.append(argv[i])
    278          elif arg == '-src_geom_type' and i + 1 < len(argv):
    279              i = i + 1
    280              src_geom_type_names = argv[i].split(',')
    281              for src_geom_type_name in src_geom_type_names:
    282                  src_geom_type = _GetGeomType(src_geom_type_name)
    283                  if src_geom_type is None:
    284                      print('ERROR: Unrecognized geometry type: %s' %
    285                            src_geom_type_name)
    286                      return 1
    287                  src_geom_types.append(src_geom_type)
    288          elif arg[0] == '-':
    289              print('ERROR: Unrecognized argument : %s' % arg)
    290              return Usage()
    291          else:
    292              if '*' in arg:
    293                  if sys.version_info < (3,0,0):
    294                      src_datasets += [fn.decode(sys.getfilesystemencoding()) for fn in glob.glob(arg)]
    295                  else:
    296                      src_datasets += glob.glob(arg)
    297              else:
    298                  src_datasets.append(arg)
    299          i = i + 1
    300
    301      if dst_filename is None:
    302          print('Missing -o')
    303          return 1
    304
    305      if update:
    306          if output_format is not None:
    307              print('ERROR: -f incompatible with -update')
    308              return 1
    309          if dsco:
    310              print('ERROR: -dsco incompatible with -update')
    311              return 1
    312          output_format = ''
    313      else:
    314          if output_format is None:
    315              output_format = GetOutputDriverFor(dst_filename)
    316
    317      if src_layer_field_content is None:
    318          src_layer_field_content = '{AUTO_NAME}'
    319      elif src_layer_field_name is None:
    320          src_layer_field_name = 'source_ds_lyr'
    321
    322      if not single_layer and output_format == 'ESRI Shapefile' and \
    323         dst_filename.lower().endswith('.shp'):
    324          print('ERROR: Non-single layer mode incompatible with non-directory '
    325                'shapefile output')
    326          return 1
    327
    328      if not src_datasets:
    329          print('ERROR: No source datasets')
    330          return 1
    331
    332      if layer_name_template is None:
    333          if single_layer:
    334              layer_name_template = 'merged'
    335          else:
    336              layer_name_template = '{AUTO_NAME}'
    337
    338      vrt_filename = None
    339      if not EQUAL(output_format, 'VRT'):
    340          dst_ds = gdal.OpenEx(dst_filename, gdal.OF_VECTOR | gdal.OF_UPDATE)
    341          if dst_ds is not None:
    342              if not update and not overwrite_ds:
    343                  print('ERROR: Destination dataset already exists, ' +
    344                        'but -update nor -overwrite_ds are specified')
    345                  return 1
    346              if overwrite_ds:
    347                  drv = dst_ds.GetDriver()
    348                  dst_ds = None
    349                  if drv.GetDescription() == 'OGR_VRT':
    350                      # We don't want to destroy the sources of the VRT
    351                      gdal.Unlink(dst_filename)
    352                  else:
    353                      drv.Delete(dst_filename)
    354          elif update:
    355              print('ERROR: Destination dataset does not exist')
    356              return 1
    357          if dst_ds is None:
    358              drv = gdal.GetDriverByName(output_format)
    359              if drv is None:
    360                  print('ERROR: Invalid driver: %s' % output_format)
    361                  return 1
    362              dst_ds = drv.Create(
    363                  dst_filename, 0, 0, 0, gdal.GDT_Unknown, dsco)
    364              if dst_ds is None:
    365                  return 1
    366
    367          vrt_filename = '/vsimem/_ogrmerge_.vrt'
    368      else:
    369          if gdal.VSIStatL(dst_filename) and not overwrite_ds:
    370              print('ERROR: Destination dataset already exists, ' +
    371                    'but -overwrite_ds are specified')
    372              return 1
    373          vrt_filename = dst_filename
    374
    375      f = gdal.VSIFOpenL(vrt_filename, 'wb')
    376      if f is None:
    377          print('ERROR: Cannot create %s' % vrt_filename)
    378          return 1
    379
    380      writer = XMLWriter(f)
    381      writer.open_element('OGRVRTDataSource')
    382
    383      if single_layer:
    384
    385          ogr_vrt_union_layer_written = False
    386
    387          for src_ds_idx, src_dsname in enumerate(src_datasets):
    388              src_ds = ogr.Open(src_dsname)
    389              if src_ds is None:
    390                  print('ERROR: Cannot open %s' % src_dsname)
    391                  if skip_failures:
    392                      continue
    393                  gdal.VSIFCloseL(f)
    394                  gdal.Unlink(vrt_filename)
    395                  return 1
    396              for src_lyr_idx, src_lyr in enumerate(src_ds):
    397                  if src_geom_types:
    398                      gt = ogr.GT_Flatten(src_lyr.GetGeomType())
    399                      if gt not in src_geom_types:
    400                          continue
    401
    402                  if not ogr_vrt_union_layer_written:
    403                      ogr_vrt_union_layer_written = True
    404                      writer.open_element('OGRVRTUnionLayer',
    405                                          attrs={'name': layer_name_template})
    406
    407                      if src_layer_field_name is not None:
    408                          writer.write_element_value('SourceLayerFieldName',
    409                                                     src_layer_field_name)
    410
    411                      if field_strategy is not None:
    412                          writer.write_element_value('FieldStrategy',
    413                                                     field_strategy)
    414
    415                  layer_name = src_layer_field_content
    416
    417                  src_lyr_name = src_lyr.GetName()
    418                  try:
    419                      src_lyr_name = src_lyr_name.decode('utf-8')
    420                  except AttributeError:
    421                      pass
    422
    423                  basename = None
    424                  if os.path.exists(src_dsname):
    425                      basename = os.path.basename(src_dsname)
    426                      if '.' in basename:
    427                          basename = '.'.join(basename.split(".")[0:-1])
    428
    429                  if basename == src_lyr_name:
    430                      layer_name = layer_name.replace('{AUTO_NAME}', basename)
    431                  elif basename is None:
    432                      layer_name = layer_name.replace(
    433                          '{AUTO_NAME}',
    434                          'Dataset%d_%s' % (src_ds_idx, src_lyr_name))
    435                  else:
    436                      layer_name = layer_name.replace(
    437                          '{AUTO_NAME}', basename + '_' + src_lyr_name)
    438
    439                  if basename is not None:
    440                      layer_name = layer_name.replace('{DS_BASENAME}', basename)
    441                  else:
    442                      layer_name = layer_name.replace('{DS_BASENAME}',
    443                                                      src_dsname)
    444                  layer_name = layer_name.replace('{DS_NAME}', '%s' %
    445                                                  src_dsname)
    446                  layer_name = layer_name.replace('{DS_INDEX}', '%d' %
    447                                                  src_ds_idx)
    448                  layer_name = layer_name.replace('{LAYER_NAME}',
    449                                                  src_lyr_name)
    450                  layer_name = layer_name.replace('{LAYER_INDEX}', '%d' %
    451                                                  src_lyr_idx)
    452
    453                  if t_srs is not None:
    454                      writer.open_element('OGRVRTWarpedLayer')
    455
    456                  writer.open_element('OGRVRTLayer',
    457                                      attrs={'name': layer_name})
    458                  attrs = {}
    459                  if EQUAL(output_format, 'VRT') and \
    460                     os.path.exists(src_dsname) and \
    461                     not os.path.isabs(src_dsname) and \
    462                     '/' not in vrt_filename and \
    463                     '\\' not in vrt_filename:
    464                      attrs['relativeToVRT'] = '1'
    465                  if single_layer:
    466                      attrs['shared'] = '1'
    467                  writer.write_element_value('SrcDataSource', src_dsname,
    468                                             attrs=attrs)
    469                  writer.write_element_value('SrcLayer', src_lyr.GetName())
    470
    471                  if a_srs is not None:
    472                      writer.write_element_value('LayerSRS', a_srs)
    473
    474                  writer.close_element('OGRVRTLayer')
    475
    476                  if t_srs is not None:
    477                      if s_srs is not None:
    478                          writer.write_element_value('SrcSRS', s_srs)
    479
    480                      writer.write_element_value('TargetSRS', t_srs)
    481
    482                      writer.close_element('OGRVRTWarpedLayer')
    483
    484          if ogr_vrt_union_layer_written:
    485              writer.close_element('OGRVRTUnionLayer')
    486
    487      else:
    488
    489          for src_ds_idx, src_dsname in enumerate(src_datasets):
    490              src_ds = ogr.Open(src_dsname)
    491              if src_ds is None:
    492                  print('ERROR: Cannot open %s' % src_dsname)
    493                  if skip_failures:
    494                      continue
    495                  gdal.VSIFCloseL(f)
    496                  gdal.Unlink(vrt_filename)
    497                  return 1
    498              for src_lyr_idx, src_lyr in enumerate(src_ds):
    499                  if src_geom_types:
    500                      gt = ogr.GT_Flatten(src_lyr.GetGeomType())
    501                      if gt not in src_geom_types:
    502                          continue
    503
    504                  src_lyr_name = src_lyr.GetName()
    505                  try:
    506                      src_lyr_name = src_lyr_name.decode('utf-8')
    507                  except AttributeError:
    508                      pass
    509
    510                  layer_name = layer_name_template
    511                  basename = None
    512                  if os.path.exists(src_dsname):
    513                      basename = os.path.basename(src_dsname)
    514                      if '.' in basename:
    515                          basename = '.'.join(basename.split(".")[0:-1])
    516
    517                  if basename == src_lyr_name:
    518                      layer_name = layer_name.replace('{AUTO_NAME}', basename)
    519                  elif basename is None:
    520                      layer_name = layer_name.replace(
    521                          '{AUTO_NAME}',
    522                          'Dataset%d_%s' % (src_ds_idx, src_lyr_name))
    523                  else:
    524                      layer_name = layer_name.replace(
    525                          '{AUTO_NAME}', basename + '_' + src_lyr_name)
    526
    527                  if basename is not None:
    528                      layer_name = layer_name.replace('{DS_BASENAME}', basename)
    529                  elif '{DS_BASENAME}' in layer_name:
    530                      if skip_failures:
    531                          if '{DS_INDEX}' not in layer_name:
    532                              layer_name = layer_name.replace(
    533                                  '{DS_BASENAME}', 'Dataset%d' % src_ds_idx)
    534                      else:
    535                          print('ERROR: Layer name template %s '
    536                                'includes {DS_BASENAME} '
    537                                'but %s is not a file' %
    538                                (layer_name_template, src_dsname))
    539
    540                          gdal.VSIFCloseL(f)
    541                          gdal.Unlink(vrt_filename)
    542                          return 1
    543                  layer_name = layer_name.replace('{DS_NAME}', '%s' %
    544                                                  src_dsname)
    545                  layer_name = layer_name.replace('{DS_INDEX}', '%d' %
    546                                                  src_ds_idx)
    547                  layer_name = layer_name.replace('{LAYER_NAME}',
    548                                                  src_lyr_name)
    549                  layer_name = layer_name.replace('{LAYER_INDEX}', '%d' %
    550                                                  src_lyr_idx)
    551
    552                  if t_srs is not None:
    553                      writer.open_element('OGRVRTWarpedLayer')
    554
    555                  writer.open_element('OGRVRTLayer',
    556                                      attrs={'name': layer_name})
    557                  attrs = {}
    558                  if EQUAL(output_format, 'VRT') and \
    559                     os.path.exists(src_dsname) and \
    560                     not os.path.isabs(src_dsname) and \
    561                     '/' not in vrt_filename and \
    562                     '\\' not in vrt_filename:
    563                      attrs['relativeToVRT'] = '1'
    564                  if single_layer:
    565                      attrs['shared'] = '1'
    566                  writer.write_element_value('SrcDataSource', src_dsname,
    567                                             attrs=attrs)
    568                  writer.write_element_value('SrcLayer', src_lyr_name)
    569
    570                  if a_srs is not None:
    571                      writer.write_element_value('LayerSRS', a_srs)
    572
    573                  writer.close_element('OGRVRTLayer')
    574
    575                  if t_srs is not None:
    576                      if s_srs is not None:
    577                          writer.write_element_value('SrcSRS', s_srs)
    578
    579                      writer.write_element_value('TargetSRS', t_srs)
    580
    581                      writer.close_element('OGRVRTWarpedLayer')
    582
    583      writer.close_element('OGRVRTDataSource')
    584
    585      gdal.VSIFCloseL(f)
    586
    587      ret = 0
    588      if not EQUAL(output_format, 'VRT'):
    589          accessMode = None
    590          if append:
    591              accessMode = 'append'
    592          elif overwrite_layer:
    593              accessMode = 'overwrite'
    594          ret = gdal.VectorTranslate(dst_ds, vrt_filename,
    595                                     accessMode=accessMode,
    596                                     layerCreationOptions=lco,
    597                                     skipFailures=skip_failures,
    598                                     callback=progress,
    599                                     callback_data=progress_arg)
    600          if ret == 1:
    601              ret = 0
    602          else:
    603              ret = 1
    604          gdal.Unlink(vrt_filename)
    605
    606      return ret
    607
    608  ###############################################################
    609  # Entry point
    610
    611
    612 -def main(argv):
    613      if sys.version_info < (3,0,0):
    614          argv = [fn.decode(sys.getfilesystemencoding()) for fn in argv]
    615      argv = ogr.GeneralCmdLineProcessor(argv)
    616      if argv is None:
    617          return 1
    618      return process(argv[1:])
    619
    620
    621  if __name__ == '__main__':
    622      sys.exit(main(sys.argv))
    623