csv-wkt_to_shp.py

# -*- coding: utf-8 -*-
"""
@author: mthh
Convert a csv file (containing WKT geometry) into a shapefile
"""

from osr import SpatialReference
import ogr
from csv import DictReader
import sys
from os.path import exists
from os import remove


def usage(name):
    print("{} input_csv wkt_field_name output_shp".format(name))
    sys.exit(0)


def main(input_path, GEOM_FIELD, output_path):
    if exists(output_path):
        remove(output_path)
    # Set the spatial reference :
    spatialref = SpatialReference()
    spatialref.SetWellKnownGeogCS('WGS84')
    # Set the destination driver :
    driver = ogr.GetDriverByName("ESRI Shapefile")
    dstfile = driver.CreateDataSource(output_path)
    dstlayer = dstfile.CreateLayer("layer", spatialref)
    # Read the feature in the csv file:
    with open(input_path) as file_input:
        dreader = DictReader(file_input)
        fieldnames = list(dreader.fieldnames)  # Grab fieldname list
        fieldnames.remove(GEOM_FIELD)
        for field in fieldnames:  # Define all the new field (string format)
            fielddef = ogr.FieldDefn(field, ogr.OFTString)
            fielddef.SetWidth(254)
            dstlayer.CreateField(fielddef)
        for nb, row in enumerate(dreader):
            poly = ogr.CreateGeometryFromWkt(row[GEOM_FIELD])  # Read geometry from the defined field
            feature = ogr.Feature(dstlayer.GetLayerDefn())  # and write it to the shapefile
            feature.SetGeometry(poly)
            for field in fieldnames:  # Write other fields
                feature.SetField(field, row[field])
            dstlayer.CreateFeature(feature)
        feature.Destroy()
        dstfile.Destroy()

if __name__ == '__main__':
    if len(sys.argv) != 4:
        print("Csv (with WKT geom) to shapefile\n\nUsage :")
        usage(sys.argv[0])
    elif exists(sys.argv[1]):
        main(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print('Error: can\'t find input file\nUsage:')
        usage(sys.argv[0])