CopyAsGeoJSON SequelPro bundle

copy_as_geo_json.py

#!/usr/local/bin/python3
import io
import json
import random
import re
import subprocess
import sys

from osgeo import gdal
from ogr import CreateGeometryFromWkt

gdal.UseExceptions()

def split(line):
	# parse TSV according to http://www.sequelpro.com/bundles/reference/
	return [x.replace("\u21e5", "\t").replace("\u21b5", "\n") if x != 'NULL' else None for x in line.split("\t")]

def parseWkt(wkt):
	if wkt is None:
		return None

	try:
		geometry = CreateGeometryFromWkt(wkt)
		if geometry is None:
			return None
	except:
		return None

	geojson = geometry.ExportToJson()
	if geojson is None:
		return None

	return json.loads(geojson)

def handleLine(line, columnNames, wktColumns):
	columns = split(line)

	properties = dict((columnNames[i], x) for i, x in enumerate(columns) if i not in wktColumns)

	# generate random color and mix it in with the properties
	color = "#%06x" % random.randint(0,0xFFFFFF)
	properties['fill'] = color
	properties['stroke'] = color
	properties['marker-color'] = color

	# create features
	for wktColumn in wktColumns:
		geometry = parseWkt(columns[wktColumn])
		if geometry is None:
			continue

		# add column name to feature properties
		properties['@column'] = columnNames[wktColumn]

		yield {
			"type": "Feature",
			"properties": properties.copy(),
			"geometry": geometry
		}

def sendToClipboard(text):
	p = subprocess.Popen(['pbcopy'], stdin=subprocess.PIPE)
	p.stdin.write(text.encode('utf8'))
	p.stdin.close()

input = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8')

firstLine = next(input)
columnNames = split(firstLine)

wktColumns = []
wktRegex = re.compile(r'^(MULTI)?(POINT|POLYGON|LINESTRING)\s*\(')
secondLine = next(input)
for index, column in enumerate(split(secondLine)):
	if column is not None and wktRegex.match(column):
		wktColumns.append(index)

if len(wktColumns) == 0:
	raise RuntimeError("Couldn't find WKT column!")

features = []
features.extend(list(handleLine(secondLine, columnNames, wktColumns)))
for line in input:
	features.extend(list(handleLine(line, columnNames, wktColumns)))

featureCollection = {
	"type": "FeatureCollection",
	"features": features
}
sendToClipboard(json.dumps(featureCollection))