arc2mb.py

#!/usr/bin/env python

# Standard library modules
import base64
import glob
import json
import os.path
import re
import shutil
import subprocess
import sys
import zipfile

# Required third-party modules
from boto3.session import Session as boto3_session
from cachecontrol import CacheControl
import ogr
import osr
import requests
from uritemplate import URITemplate

# ------------------------------------------------------------------
"""Converts shapefile or zipped shapefile to Mapbox-hosted vector tile layer
(named after the input file and overwriting/updating if existing), building
vector pyramids through tippecanoe and then following the upload procedure
described at https://www.mapbox.com/api-documentation/#uploads. The max zoom
level parameter sets the final precision of the output product; the geometry
will be lossless up to the specified zoom level, after which some
simplification may occur. A zoom level of at least 10 is reccomended for
polygons and lines; higher zoom levels will increase processing time
significantly. Leave this parameter blank to skip the vector tile generation
step (reccomended for all points; not recomended for complex lines and
polygons).

Call from the command line using:
python arc2mb.py {input file path} {output layer name} {Mapbox token} {max zoom level}
if polygons or lines, for which vector tiling is necessary, or
python arc2mb.py {input file path} {output layer name} {Mapbox token}
if points, for which vector tiling should be skipped.

Call from other scripts stored in the same directory using:
import arc2mb
arc2mb({input file path}, {output layer name} {Mapbox token}, {max zoom level})

The Mapbox token must be a private key with write permissions, which can be
generated by choosing to add an access token in the online Mapbox account
manager and enabling the optional secret token scopes.
"""
# ------------------------------------------------------------------
# Roughly, the procedure for each unique shapefile file is to:
# - Unzip shapefile archive, if necessary.
# - Convert shapefile to WGS-84 projected GeoJSON using OGR and OSR.
# - Generate vector tile layers (.mbtile format) using tippecanoe.
# - Poll Mapbox's API for the credentials of an Amazon S3 bucket in
#   which to stage the upload.
# - Upload to the staging bucket.
# - Send staged data to Mapbox's Uploader API.
#
# If no max zoom level is specified, the script will skip the GeoJSON
# and tiling step, and send the zipped shapefile package directly to
# the Mapbox Uploader API.
# ------------------------------------------------------------------

"""Error classes"""


class ValidationError(ValueError):
    pass


class HTTPError(ValidationError):
    pass


class InvalidFileError(ValidationError):
    pass


class TokenError(ValidationError):
    pass


"""Base Service class"""


class Service(object):
    """Service base class."""

    def __init__(self, access_token=None, cache=None):
        """Constructs a Service object.
        """
        self.session = requests.Session()
        self.session.params.update(access_token=access_token)
        self.session.headers.update({
            "User-Agent": "mapbox-sdk-py/{0} {1}".format(
                "0.8.0", requests.utils.default_user_agent())})
        if cache:
            self.session = CacheControl(self.session, cache=cache)

    @property
    def username(self):
        """Get username from access token. Token contains base64 encoded
        JSON object with username.
        """
        token = self.session.params.get("access_token")
        if not token:
            raise TokenError(
                "session does not have a valid access_token param")
        data = token.split('.')[1]
        # replace url chars and add padding
        data = data.replace("-", "+").replace("_", "/") + "==="
        try:
            return json.loads(base64.b64decode(data).decode("utf-8"))["u"]
        except (ValueError, KeyError):
            raise TokenError(
                "access_token does not contain username")

    def handle_http_error(self, response, custom_messages=None,
                          raise_for_status=False):
        if not custom_messages:
            custom_messages = {}
        if response.status_code in custom_messages.keys():
            raise HTTPError(custom_messages[response.status_code])
        if raise_for_status:
            response.raise_for_status()


"""Uploader class"""


class Uploader(Service):
    """Access to the Mapbox Upload API.
    """

    baseuri = "https://api.mapbox.com/uploads/v1"

    def _get_credentials(self):
        """Gets temporary S3 credentials to stage user-uploaded files.
        """
        print " -- Getting credentials for Amazon S3 staging bucket"

        uri = URITemplate(self.baseuri + "/{username}/credentials").expand(
            username=self.username)
        resp = self.session.get(uri)
        self.handle_http_error(
            resp,
            custom_messages={
                401: "Token is not authorized",
                404: "Token does not have upload scope"})
        return resp

    def stage(self, fileobj, creds=None):
        """Stages the user's file on S3. Returns the URL to the staged
        resource.
        """
        if not hasattr(fileobj, "read"):
            raise InvalidFileError(
                "Object `{0}` has no .read method, "
                "a file-like object is required".format(fileobj))

        if not creds:
            res = self._get_credentials()
            creds = res.json()

        print " -- Uploading package to Amazon S3 staging bucket"

        session = boto3_session(
            aws_access_key_id=creds["accessKeyId"],
            aws_secret_access_key=creds["secretAccessKey"],
            aws_session_token=creds["sessionToken"],
            region_name="us-east-1")

        s3 = session.resource("s3")
        if hasattr(fileobj, "name") and os.path.exists(fileobj.name):
            s3.Bucket(creds["bucket"]).upload_file(fileobj.name, creds["key"])
        else:
            res = s3.Object(creds["bucket"], creds["key"]).put(Body=fileobj)

        return creds["url"]

    def create(self, stage_url, tileset_name):
        """Initiates the creation process from the staging S3 bucket into
        the user's tileset.
        """
        if not tileset_name.startswith(self.username + "."):
            tileset_name = "{0}.{1}".format(
                self.username, tileset_name.replace(" ", "_")[0:31])

        msg = {"tileset": tileset_name,
               "url": stage_url}

        if tileset_name is not None:
            msg["name"] = tileset_name.replace(" ", "_")[0:31]

        uri = URITemplate(self.baseuri + "/{username}").expand(
            username=self.username)

        resp = self.session.post(uri, json=msg)
        self.handle_http_error(resp)

        print " -- Uploaded and sent to Mapbox Upload API"

        return resp

    def upload(self, fileobj, tileset_name):
        """High level function to upload a file object to Mapbox tileset
        Returns a response object where the json() is a dict with upload
        metadata.
        """
        url = self.stage(fileobj)
        return self.create(url, tileset_name)


"""Shapefile packager helper function"""


def zip_shapefile(input_shp):
    """ Packaging helper function which given a path to a shapefile,
    creates an archive of all of its support files (overwriting existing)
    and returns its path.
    """
    print " -- Skipping tile generation and packaging shapefile"

    base_path = os.path.splitext(input_shp)[0]
    package = base_path + ".zip"

    if os.path.exists(package):
        os.remove(package)

    f_list = glob.glob(base_path + "*")
    for f in f_list:
        zf = zipfile.ZipFile(package, "a", zipfile.ZIP_DEFLATED)
        zf.write(f, os.path.basename(f))
        zf.close()

    return package


"""Shapefile unzipper helper function"""


def unzip_shapefile(input_zip):
    """ Unzipping helper function which unzips an archive of shapefile
    support files and returns the path to the primary shapefile.
    """
    print " -- Unzipping shapefile"

    dest = os.path.splitext(input_zip)[0]
    if not os.path.exists(dest):
        os.mkdir(dest)

    with zipfile.ZipFile(input_zip, "r") as z:
        z.extractall(dest)

    output_shp = glob.glob(os.path.join(dest, "*.shp"))[0]

    return output_shp


"""Vector tile tippecanoe dispatcher"""


def generate_mbtiles(base_path, json_path, name, max_zoom):
    """Uses tippecanoe to pre-generate vector tile layers, because
    Mapbox's online Uploader API does not offer control over generated
    zoom detail levels.
    """
    print " -- Launching tippecanoe to generate vector tiles. " + \
          "Tippecanoe output:\n"

    tile_path = os.path.join(base_path, name) + '.mbtiles'

    args = ["/usr/local/bin/tippecanoe", "-f", "-o",
            os.path.abspath(tile_path),
            "-z", str(int(max_zoom)),
            os.path.abspath(json_path)]
    subprocess.call(args)

    print "\n"

    return tile_path


"""Shapefile to JSON converter"""


def shp_to_json(base_path, shp_path, name):
    """ Conversion helper function uses ogr to project shapefiles
    and convert them to GeoJSON, and returns the filesystem path
    to the result.
    """
    print " -- Projecting shapefile to WGS-84 and converting to JSON"

    # define ogr drivers
    shp_driver = ogr.GetDriverByName('ESRI Shapefile')
    json_driver = ogr.GetDriverByName('GeoJSON')

    # define the input layer
    shp = shp_driver.Open(shp_path)
    shp_lyr = shp.GetLayer()

    # create the output layer
    json_path = os.path.join(base_path, name + ".geojson")
    if os.path.exists(json_path):
        json_driver.DeleteDataSource(json_path)
    json = json_driver.CreateDataSource(json_path)
    json_lyr = json.CreateLayer(json_path, geom_type=ogr.wkbMultiPolygon)
    json_lyr_defn = json_lyr.GetLayerDefn()

    # create the CoordinateTransformation
    json_ref = osr.SpatialReference()
    json_ref.ImportFromEPSG(4326)
    coord_trans = osr.CoordinateTransformation(
        shp_lyr.GetSpatialRef(), json_ref)

    # add fields to output layer
    shp_lyr_defn = shp_lyr.GetLayerDefn()
    for i in range(0, shp_lyr_defn.GetFieldCount()):
        field_defn = shp_lyr_defn.GetFieldDefn(i)
        json_lyr.CreateField(field_defn)

    # loop through the input features
    shp_feat = shp_lyr.GetNextFeature()
    while shp_feat:
        # reproject the input geometry
        geom = shp_feat.GetGeometryRef()
        geom.Transform(coord_trans)
        # create a new feature
        json_feat = ogr.Feature(json_lyr_defn)
        # set the feature's geometry and attributes
        json_feat.SetGeometry(geom)
        for i in range(0, json_lyr_defn.GetFieldCount()):
            json_feat.SetField(
                json_lyr_defn.GetFieldDefn(i).GetNameRef(),
                shp_feat.GetField(i))
        # add new feature to output Layer
        json_lyr.CreateFeature(json_feat)
        # destroy the features and get the next input feature
        json_feat.Destroy()
        shp_feat.Destroy()
        shp_feat = shp_lyr.GetNextFeature()

    # close the datasets
    shp.Destroy()
    json.Destroy()

    return json_path


"""Main driver"""


def main(shp_path, tileset_name, mb_token, max_zoom):
    """Main loop which creates Uploader class instance, converts shapefiles
    or zipped shapefiles to WGS-84 projected GeoJSON, generates vector tile
    layers, uploads tile layers to Amazon S3 staging bucket, and dispatches
    the packages to the Mapbox Uploader API. If no max zoom is specified,
    the function simply zips the shapefile and submits the package to the
    Uploader API.
    """
    base_path, name = os.path.split(shp_path)
    name, ext = os.path.splitext(name)

    print "\nProcessing " + name

    # unzip shapefile if necessary
    if ext == ".zip":
        shp_path = unzip_shapefile(shp_path)
        base_path, name = os.path.split(shp_path)
        name = os.path.splitext(name)[0]

    # instantiate the uploader service
    service = Uploader()
    service.session.params["access_token"] = mb_token

    if max_zoom:
        # convert shapefiles to intermediate GeoJSON data
        json_path = shp_to_json(base_path, shp_path, name)
        # generate vector tiles from GeoJSON using tippecanoe
        package = generate_mbtiles(base_path, json_path, name, max_zoom)
        # cleanup the intermediate GeoJSON data
        os.remove(json_path)
    else:
        # if no max zoom specified, skip GeoJSON/ tippecanoe step, zip
        # shapefile into package, and continue uploading
        package = zip_shapefile(os.path.join(base_path, name))
    # dispatch mbtiles to the uploader service
    with open(package, "r") as src:
        upload_resp = service.upload(src, tileset_name)
        try:
            print " -- Hosted tileset name is {0}\n".format(
                upload_resp.json()["tileset"])
        except:
            print " !! Mapbox server returned error:"
            print upload_resp.json()

    # cleanup final product and processing files after upload
    os.remove(package)
    if ext == ".zip":
        shutil.rmtree(base_path)


if __name__ == "__main__":
    args = sys.argv[1:]
    if len(args) == 4:
        shp_path, tileset_name, mb_token, max_zoom = args
        # sanitize non-numeric characters from max_zoom
        max_zoom = int(re.sub("[^0-9]", "", max_zoom))
    elif len(args) == 3:
        shp_path, tileset_name, mb_token = args
        max_zoom = False
    else:
        print 'Please enter arguments in the form "python arc2mb.py ' + \
              '{input file path} {output layer name} {Mapbox token} ' + \
              '{max zoom level}", or "python arc2mb.py {input file path} ' + \
              '{output layer name} {Mapbox token}" if skipping tile ' + \
              'generation (only recommended for points).'
    
    if os.path.exists(shp_path):
        main(shp_path, tileset_name, mb_token, max_zoom)
    else:
        print "Input shapefile not found. Please enter a valid path."