QGIS » QGIS Application

#!/usr/bin/python

# shapemerger.py

# v.1.2 by toni @ http://gistncase.blogspot.com

#

# shapefile merger utility 

import os, sys, glob

try:

    from osgeo import ogr, gdal

except:

    try:

        import ogr, gdal

    except:

        print 'Error: shapemerge.py cannot find python OGR and GDAL modules'

        print '       for handling geometries. OsGeo4W python is recommended.'

        sys.exit(1)

def Usage():

    print

    print 'Shapefile merger'

    print 'Usage: shapemerger.py\t[-o <output.shp>] [-n] [-i] [-s <fieldname>] [-f <fieldname>]'

    print '\t\t\t[-e] [-2] [-r] [-t polygon|linestring|point]'

    print '\t\t\t<shapefiles> [-x <exclude shapefile list>]'

    print 

    print '  <shapefiles>: input shapefile list (accepts wildcards)'

    print

    print '  -o:\toutput shapefile name (default: shapemerge.shp)'

    print '  -n:\tdo not add fields for the source filenames and fids'

    print '  -i:\tdbf schemes intersection (default: union)'

    print '  -s:\tcustom name for the source filenames field (default: SOURCESHP)'

    print '  -f:\tcustom name for the source fids field (default: SOURCEFID)'

    print '  -e:\tuse existing source filenames and fids fields, if any'

    print '  -2:\tforce output geometry type to 2D'

    print '  -x:\tlist of shapefiles to exclude (accepts wildcards)'

    print '  -r:\trecursive (grabs every path of input files)'

    print '  -t:\tfilter input shapefiles by geometry type'

    sys.exit(1)

strCurrentPath = os.path.abspath('')

os.chdir(strCurrentPath)

strOutDirname = ''

strMergeName = 'shapemerge'

strOutBasename = 'shapemerge.shp'

pOgrShpDriver = ogr.GetDriverByName('Esri Shapefile')

lsFilenames=[]

lsExcludeFilenames=[]

lsVerifiedFilenames=[]

lsMergeFieldsToRetain=[]

lsMergeFieldNames=[]

lsMergeFieldOgrDefns=[]

dGeomTypeFilters={'POLYGON': ogr.wkbPolygon, 'AREAL': ogr.wkbPolygon, 'LINESTRING': ogr.wkbLineString, 'LINE': ogr.wkbLineString, 'POINT': ogr.wkbPoint, 'POINTS': ogr.wkbPoint}

bFlagExclude = 0

bFlagAddSourcefields = 1

bFlagUseExistingSfields = 0

bFlagIntersection = 0

bFlagRecursive = 0

bFlagGeomTFilter = 0

bFlag2D = 0

bGeomTypeRefIsSet = 0

strFilenameFieldname = 'SOURCESHP'

strFidFieldname = 'SOURCEFID'

i = 1

# parse commandline (first check for recursive flag)

while i < len(sys.argv): 

    arg = sys.argv[i]

    if arg.upper() == '-R':

        bFlagRecursive = 1

        del sys.argv[i]

        break

    i = i + 1

i = 1

while i < len(sys.argv): 

    arg = sys.argv[i]

    if arg.upper() == '-N':

        bFlagAddSourcefields = 0

        bFlagExclude = 0

    elif arg.upper() == '-I':

        try:

            strGdalVersion = int(gdal.VersionInfo()[:2])

        except:

            strGdalVersion = 0

        if (strGdalVersion < 19):

            print 'Warning: dbf schemes intersection supported only with gdal >= 1.9.0, ignoring.'

            bFlagIntersection = 0

        else:

            bFlagIntersection = 1

        bFlagExclude = 0

    elif arg.upper() == '-S':

        i = i + 1

        strFilenameFieldname = ''.join(sys.argv[i].split())[:10].upper() #upper, tronca al decimo, toglie spazi

        bFlagExclude = 0

    elif arg.upper() == '-F':

        i = i + 1

        strFidFieldname = ''.join(sys.argv[i].split())[:10].upper() #upper, tronca al decimo, toglie spazi

        bFlagExclude = 0

    elif arg.upper() == '-E':

        bFlagUseExistingSfields = 1

        bFlagExclude = 0

    elif arg.upper() == '-2':

        bFlag2D = 1

        bFlagExclude = 0

    elif arg.upper() == '-T':

        i = i + 1

        strGeomTFilter = (sys.argv[i]).upper()

        if (not (strGeomTFilter in dGeomTypeFilters)):

            print "Warning: invalid geometry type declared, ignoring." 

        else:

            intGeometryTypeRef = dGeomTypeFilters[strGeomTFilter]

            bFlagGeomTFilter = 1

        bFlagExclude = 0

    elif arg.upper() == '-X':

        bFlagExclude = 1

    elif arg.upper() == '-O':

        i = i + 1

        bFlagExclude = 0

        strOutDirname, strOutBasename = os.path.split(sys.argv[i])

        strMergeName, strExtension = os.path.splitext(strOutBasename)

    else:

        # expands wildcards: exclude, input list/recursive input list

        if (bFlagRecursive):

            strInputShapePath, strInputShapeBasename = os.path.split(arg)

            strInputShapePath = strInputShapePath or '.'

            lsInputShapeRoots=[]

            lsInputShapeRoots.extend(glob.glob(strInputShapePath)) #expands paths with wildcards

        if (bFlagExclude and bFlagRecursive):

            for strRootPath in lsInputShapeRoots:

                for strRoot,strDirs,strFiles in os.walk(strRootPath):

                    lsExcludeFilenames.extend ( glob.glob(strRoot + os.sep + strInputShapeBasename) )

        elif (bFlagExclude and (not bFlagRecursive)):

            lsExcludeFilenames.extend ( glob.glob(arg) )

        elif (bFlagRecursive):

            for strRootPath in lsInputShapeRoots:

                for strRoot,strDirs,strFiles in os.walk(strRootPath):

                    lsFilenames.extend ( glob.glob(strRoot + os.sep + strInputShapeBasename) )

        else:

            lsFilenames.extend ( glob.glob(arg) )

    i = i + 1

for strItem in lsExcludeFilenames:

    try:

        lsFilenames.remove(strItem)

    except:

        pass

if len(lsFilenames) == 0:

    Usage()

    sys.exit(1)

if not os.path.isdir(strOutDirname or '.'):

    print strOutDirname + ': invalid directory, outputting here'

    strOutDirname = ''

if os.path.isfile( os.path.join(strOutDirname, strOutBasename) ):

    strSuffix = ''

    j = 0

    while os.path.isfile( os.path.join(strOutDirname,strMergeName + strSuffix + '.shp') ):

        j = j + 1

        strSuffix = '_' + str(j)

    if (strMergeName != 'shapemerge'):

        print strOutBasename + ': file exists, outputting in ' + strMergeName + strSuffix + '.shp'

    if (j > 0):

        strMergeName = strMergeName + strSuffix

        strOutBasename = strMergeName + '.shp'  

# geometry type check / "source attributes" renaming

strSuffix = ''

j = 0 

for i,strFile in enumerate(lsFilenames):

    try:

        pOgrDatasource = pOgrShpDriver.Open(strFile)

        pOgrLayer = pOgrDatasource.GetLayer()

        pOgrLayerSchema = pOgrLayer.GetLayerDefn()

        if ( (not bGeomTypeRefIsSet) and (not bFlagGeomTFilter) ):

            # picks up the first geometry type as reference if geometry type filter is not set

            intGeometryTypeRef = pOgrLayer.GetGeomType()

            if (bFlag2D):

                # fetch first feature to flatten geometry and grab TypeRef

                pOgrFeature = pOgrLayer.GetNextFeature()

                pOgrFeatureGeometry = pOgrFeature.GetGeometryRef()

                pOgrClonedGeometry = pOgrFeatureGeometry.Clone()

                pOgrClonedGeometry.FlattenTo2D()

                intGeometryTypeRef = pOgrClonedGeometry.GetGeometryType() 

                #pOgrFeature.Destroy()

            bGeomTypeRefIsSet = i # register index of the reference shapefile

        #check if layers the picked geometry type (ignore 3D flag)

        if pOgrLayer.GetGeomType() != intGeometryTypeRef:

            if ((pOgrLayer.GetGeomType() == ogr.wkbPolygon) or (pOgrLayer.GetGeomType() == ogr.wkbPolygon25D)) and ((intGeometryTypeRef == ogr.wkbPolygon) or (intGeometryTypeRef == ogr.wkbPolygon25D)) \

                or ((pOgrLayer.GetGeomType() == ogr.wkbLineString) or (pOgrLayer.GetGeomType() == ogr.wkbLineString25D)) and ((intGeometryTypeRef == ogr.wkbLineString) or (intGeometryTypeRef == ogr.wkbLineString25D)) \

                or ((pOgrLayer.GetGeomType() == ogr.wkbPoint) or (pOgrLayer.GetGeomType() == ogr.wkbPoint25D)) and ((intGeometryTypeRef == ogr.wkbPoint) or (intGeometryTypeRef == ogr.wkbPoint25D)) :

                    pass

            else:

                if (not bFlagGeomTFilter):

                    #print 'Error: shapefiles must have the same geometry type.'

                    #print '\t' + os.path.basename(lsFilenames[0]) + ': ' + ogr.GeometryTypeToName(intGeometryTypeRef).upper()

                    #print '\t' + os.path.basename(lsFilenames[i]) + ': ' + ogr.GeometryTypeToName(pOgrLayer.GetGeomType()).upper() 

                    #sys.exit(1)

                    print 'Warning: ' + os.path.basename( strFile ) + ' type ' + ogr.GeometryTypeToName(pOgrLayer.GetGeomType()).upper() + ' unfits first file type ' + ogr.GeometryTypeToName(intGeometryTypeRef).upper() + ' (' + os.path.basename(lsFilenames[bGeomTypeRefIsSet]) +'), skipping.'

                continue

        if ( (not bFlagUseExistingSfields) and bFlagAddSourcefields ):

            # check for "source fields" in current shapefile to eventually rename them

            while ( (pOgrLayerSchema.GetFieldIndex(strFilenameFieldname) != -1) or (pOgrLayerSchema.GetFieldIndex(strFidFieldname) != -1) ):

                j = j + 1

                strSuffix = '_' + str(j)

                strFilenameFieldname = strFilenameFieldname[:(10 - len(strSuffix))] + strSuffix

                strFidFieldname      = strFidFieldname[:(10 - len(strSuffix))] + strSuffix

        lsVerifiedFilenames.append( strFile )

        pOgrDatasource.Destroy()

    except:

        print 'Warning: ' + strFile + ' invalid shapefile, skipping.'

        #continue

if (len(lsVerifiedFilenames) == 0):

    print 'Error: all the input shapefiles are invalid.'

    sys.exit(1)

else:

    lsFilenames = lsVerifiedFilenames

if (j >0):

    print "Warning: \"source fieldnames\" renamed to: " +  strFilenameFieldname + ", " + strFidFieldname

# building attributes

if (bFlagAddSourcefields):

    # "source fields" attributes

    lsMergeFieldOgrDefns.append( ogr.FieldDefn(strFilenameFieldname, ogr.OFTString) )

    lsMergeFieldOgrDefns[0].SetWidth(255)

    lsMergeFieldOgrDefns.append( ogr.FieldDefn(strFidFieldname, ogr.OFTInteger) )

    lsMergeFieldOgrDefns[1].SetWidth(9) #max one billion features! to be fixed!

    lsMergeFieldNames.append(strFilenameFieldname)

    lsMergeFieldNames.append(strFidFieldname)

for i,strFile in enumerate(lsFilenames):

    # shapefiles attributes

    if ( i>0 and bFlagIntersection and ( (bFlagAddSourcefields and (len(lsMergeFieldNames)==2)) or (len(lsMergeFieldNames)==0) ) ):

        # no more scheme intersection to do, exit loop 

        break

    strFullname = strFile

    strDirname, strBasename = os.path.split(strFullname)

    strShapename, strExtension = os.path.splitext(strBasename)

    pOgrShpDatasource = pOgrShpDriver.Open (strFullname)

    pOgrShpLayer = pOgrShpDatasource.GetLayer()

    pOgrShpLayerSchema = pOgrShpLayer.GetLayerDefn()

    for j in range( pOgrShpLayerSchema.GetFieldCount() ):

        # parse shapefile attributes

        pOgrShpFielddefn = pOgrShpLayerSchema.GetFieldDefn(j)

        strShpFieldname = pOgrShpFielddefn.GetNameRef().upper()

        try:

            iMergeListsFieldIdx = lsMergeFieldNames.index(strShpFieldname)

        except:

            iMergeListsFieldIdx = -1

        if (iMergeListsFieldIdx == -1):

            # field doesn't exist: create it. If flag intersection is on, parse only the first file

            if ( (not bFlagIntersection) or (bFlagIntersection and (i == 0)) ):

                pOgrFielddefn = ogr.FieldDefn( strShpFieldname, pOgrShpFielddefn.GetType() )

                pOgrFielddefn.SetWidth( pOgrShpFielddefn.GetWidth() )

                pOgrFielddefn.SetPrecision( pOgrShpFielddefn.GetPrecision() )

                #ONLY TEXT:

                #pOgrFielddefn = ogr.FieldDefn( strShpFieldname, ogr.OFTString )

                #pOgrFielddefn.SetWidth( pOgrShpFielddefn.GetWidth() + pOgrShpFielddefn.GetPrecision() )

                lsMergeFieldOgrDefns.append( pOgrFielddefn )

                lsMergeFieldNames.append( strShpFieldname )

                # add field to the retain list for intersection

                lsMergeFieldsToRetain.append( len(lsMergeFieldNames) - 1 )

        else:

            # field exists, check properties

            pOgrFielddefn = lsMergeFieldOgrDefns[iMergeListsFieldIdx]

            #same data type

            if ( pOgrFielddefn.GetType() == pOgrShpFielddefn.GetType() ):

                if pOgrFielddefn.GetWidth() < pOgrShpFielddefn.GetWidth():

                    pOgrFielddefn.SetWidth( pOgrShpFielddefn.GetWidth() )

                if pOgrFielddefn.GetPrecision() < pOgrShpFielddefn.GetPrecision():

                    pOgrFielddefn.SetPrecision( pOgrShpFielddefn.GetPrecision() )

            else:

                # different data type: "translate" it into string

                pOgrFielddefn.SetType( ogr.OFTString )

                pOgrFielddefn.SetWidth( pOgrFielddefn.GetWidth() + pOgrFielddefn.GetPrecision() )

                if ( pOgrFielddefn.GetWidth() < pOgrShpFielddefn.GetWidth() + pOgrShpFielddefn.GetPrecision() ): 

                    pOgrFielddefn.SetWidth( pOgrShpFielddefn.GetWidth() + pOgrShpFielddefn.GetPrecision() )

            if (bFlagIntersection):

                # add field to the retain list for intersection

                lsMergeFieldsToRetain.append(iMergeListsFieldIdx)

    if ( bFlagIntersection ):

        # performing dbf intersection

        if (bFlagAddSourcefields):

            lsMergeFieldsToRetain.append( lsMergeFieldNames.index(strFilenameFieldname) )

            lsMergeFieldsToRetain.append( lsMergeFieldNames.index(strFidFieldname) )

        lsMergeFieldsToDelete = range( len(lsMergeFieldNames)-1, -1, -1) # reverse: from last index to 0

        for iIdx in lsMergeFieldsToRetain:

            try:

                lsMergeFieldsToDelete.remove(iIdx)

            except:

                pass

        for iIdx in lsMergeFieldsToDelete:

            # add or remove fields items

            del lsMergeFieldOgrDefns[iIdx]

            del lsMergeFieldNames[iIdx]            

    pOgrShpDatasource.Destroy()

    lsMergeFieldsToRetain = []

# generate the merge layer

strMergeShpFullname = os.path.join(strOutDirname, strOutBasename)

pOgrMergeShpDatasource = pOgrShpDriver.CreateDataSource( strMergeShpFullname )

pOgrMergeShpLayer = pOgrMergeShpDatasource.CreateLayer( strMergeName, geom_type = intGeometryTypeRef )

for i in lsMergeFieldOgrDefns:

    # create the fields

    pOgrMergeShpLayer.CreateField(i,1)

pOgrMergeShpLayerSchema = pOgrMergeShpLayer.GetLayerDefn()

# Loop for populating the merge shapefile

for i,strFile in enumerate(lsFilenames):

    strFullname = strFile

    strDirname, strBasename = os.path.split(strFullname)

    strShapename, strExtension = os.path.splitext(strBasename)

    pOgrShpDatasource = pOgrShpDriver.Open (strFullname)

    pOgrShpLayer = pOgrShpDatasource.GetLayer()

    pOgrShpLayerSchema = pOgrShpLayer.GetLayerDefn()

    pOgrShpLayer.ResetReading()

    pOgrShpFeature = pOgrShpLayer.GetNextFeature()    

    bFlagExistsFilenameFld = pOgrShpLayerSchema.GetFieldIndex(strFilenameFieldname)

    bFlagExistsFidFld = pOgrShpLayerSchema.GetFieldIndex(strFidFieldname)

    while pOgrShpFeature:

        pOgrMergeShpfeature = ogr.Feature(pOgrMergeShpLayerSchema)

        pOgrMergeShpfeature.SetFrom(pOgrShpFeature,1)

        if (bFlag2D):

            # performing 2d conversion

            pOgrMergeGeometry2d = pOgrMergeShpfeature.GetGeometryRef().Clone()

            pOgrMergeGeometry2d.FlattenTo2D()

            pOgrMergeShpfeature.SetGeometryDirectly(pOgrMergeGeometry2d)

        if (bFlagAddSourcefields):

            # "source fields" values, if any and if they should be populated

            if ( (bFlagUseExistingSfields and bFlagExistsFilenameFld == -1) or (not bFlagUseExistingSfields) ):

                #pOgrMergeShpfeature.SetField( strFilenameFieldname, strBasename )

                pOgrMergeShpfeature.SetField( strFilenameFieldname, strShapename)

            if ( (bFlagUseExistingSfields and bFlagExistsFidFld == -1) or (not bFlagUseExistingSfields) ):

                pOgrMergeShpfeature.SetField( strFidFieldname, pOgrShpFeature.GetFID() )

        pOgrMergeShpLayer.CreateFeature(pOgrMergeShpfeature)

        pOgrMergeShpfeature.Destroy()

        pOgrShpFeature.Destroy()

        pOgrShpFeature = pOgrShpLayer.GetNextFeature()

    pOgrShpDatasource.Destroy()

# flush all to disk before quitting

pOgrMergeShpDatasource.SyncToDisk()

pOgrMergeShpDatasource.Destroy()

print

print strMergeShpFullname

sys.exit(0)