<commit_before># a class to manage SRTM surfaces

import json
import numpy as np
import os
from pylab import *
import random
import scipy.interpolate
import struct
import urllib
import zipfile

import navpy

# return the lower left corner of the 1x1 degree tile containing
# the specified lla coordinate
def lla_ll_corner(lat_deg, lon_deg):
    return int(floor(lat_deg)), int(floor(lon_deg))

# return the tile base name for the specified coordinate
def make_tile_name(lat, lon):
    ll_lat, ll_lon = lla_ll_corner(lat, lon)
    if ll_lat < 0:
        slat = "S%2d" % ll_lat
    else:
        slat = "N%2d" % ll_lat
    if ll_lon < 0:
        slon = "W%03d" % -ll_lon
    else:
        slon = "W%03d" % ll_lon
    return slat + slon

class SRTM():
    def __init__(self, lat, lon, dict_path):
        self.lat, self.lon = lla_ll_corner(lat, lon)
        self.srtm_dict = {}
        self.srtm_cache_dir = '/var/tmp' # unless set otherwise
        self.srtm_z = None
        self.i = None
        self.load_srtm_dict(dict_path)

    # load the directory download dictionary (mapping a desired file
    # to a download path.)
    def load_srtm_dict(self, dict_path):
        dict_file = dict_path + '/srtm.json'
        try:
            f = open(dict_file, 'r')
            self.srtm_dict = json.load(f)
            f.close()
        except:
            print "Notice: unable to read =", dict_file

    # if we'd like a persistant place to cache srtm files so they don't
    # need to be re-download every time we run anything
    def set_srtm_cache_dir(self, srtm_cache_dir):
        if not os.path.exists(srtm_cache_dir):
            print "Notice: srtm cache path doesn't exist =", srtm_cache_dir
        else:
            self.srtm_cache_dir = srtm_cache_dir

    # download and extract srtm file into cache directory
    def download_srtm(self, fileroot):
        if fileroot in self.srtm_dict:
            url = self.srtm_dict[fileroot]
            download_file = self.srtm_cache_dir + '/' + fileroot + '.hgt.zip'
            print "Notice: downloading:", url
            file = urllib.URLopener()
            file.retrieve(url, download_file)
            #print "Extracting:", download_file
            #zip = zipfile.ZipFile(download_file)
            #zip.extractall(self.srtm_cache_dir)
            #os.remove(download_file)
            return True
        else:
            print "Notice: requested srtm that is outside catalog"
            return False

    def parse(self):
        tilename = make_tile_name(self.lat, self.lon)
        cache_file = self.srtm_cache_dir + '/' + tilename + '.hgt.zip'
        if not os.path.exists(cache_file):
            if not self.download_srtm(tilename):
                return False
        print "Notice: parsing SRTM file:", cache_file
        #f = open(cache_file, "rb")
        zip = zipfile.ZipFile(cache_file)
        f = zip.open(tilename + '.hgt', 'r')
        contents = f.read()
        f.close()
        # read 1,442,401 (1201x1201) high-endian
        # signed 16-bit words into self.z
        self.srtm_z = struct.unpack(">1442401H", contents)
        return True

    def make_lla_interpolator(self):
        print "Notice: constructing lla interpolator"

        srtm_pts = np.zeros((1201, 1201))
        for r in range(0,1201):
            for c in range(0,1201):
                idx = (1201*r)+c
                va = self.srtm_z[idx]
                if va == 65535 or va < 0 or va > 10000:
                    va = 0.0
                z = va
                srtm_pts[c,1200-r] = z
        x = np.linspace(self.lon, self.lon+1, 1201)
        y = np.linspace(self.lat, self.lat+1, 1201)
        #print x
        #print y
        self.lla_interp = scipy.interpolate.RegularGridInterpolator((x, y), srtm_pts, bounds_error=False, fill_value=-32768)
        #print self.lla_interp([-93.14530573529404, 45.220697421008396])
        #for i in range(20):
        #    x = -94.0 + random.random()
        #    y = 45 + random.random()
        #    z = self.lla_interp([x,y])
        #    print [x, y, z[0]]

    # notice: this is *really* slow.  Adding the lla2ned conversion
    # makes even more slow.  I think what I want to do is build as
    # many SRTM instances as we need to cover an area, then resample
    # just the needed area, at some reasonable grid size, convert just
    # the resampled points to ned, and build the interpolator out of
    # the subset of points.  A DEM file can cover 60nm in latitude and
    # 60nm*cos(lat) in the longitude direction.

    def make_ned_interpolator(self, ref):
        print "Notice: constructing ned interpolator"

        # The LinearNDInterpolator works well, but is slow to
        # construct, ned data is not 'regularly' gridded so we need to
        # use this fancier interpolator
        ned_pts = np.zeros((1201*1201, 2))
        ned_vals = np.zeros((1201*1201))
        for r in range(0,1201):
            print "%.1f%%" % (float(r)/float(12.01))
            for c in range(0,1201):
                idx = (1201*r)+c
                va = self.srtm_z[idx]
                if va == 65535 or va < 0 or va > 10000:
                    va = 0.0
                lon = -94.0 + (float(c) / 1200.0)
                lat = 45 + (float(1200-r) / 1200.0)
                alt = va
                ned = navpy.lla2ned( lat, lon, alt, ref[0], ref[1], ref[2] )
                ned_pts[idx] = ned[0:2]
                ned_vals[idx] = ned[2]
        self.ned_interp = scipy.interpolate.LinearNDInterpolator(ned_pts, ned_vals)

    def lla_interpolate(self, point_list):
        return self.lla_interp(point_list)

    def ned_interpolate(self, point_list):
        return self.ned_interp(point_list)

    def plot_raw(self):
        zzz = np.zeros((1201,1201))
        for r in range(0,1201):
            for c in range(0,1201):
                va=self.srtm_z[(1201*r)+c]
                if (va==65535 or va<0 or va>10000):
                    va=0.0
                zzz[r][c]=float(va)

        #zz=np.log1p(zzz)
        imshow(zzz, interpolation='bilinear',cmap=cm.gray,alpha=1.0)
        grid(False)
        show()

# Build a gridded elevation interpolation table centered at lla_ref
# with width and height.  This is a little bit of quick feet dancing,
# but allows areas to span corners or edges of srtm tiles and attempts
# to stay on a fast path of regular grids, even though a regularly lla
# grid != a regular ned grid.
class NEDGround():
    def __init__(self, lla_ref, width_m, height_m, step_m):
        self.tile_dict = {}
        self.load_tiles(lla_ref, width_m, height_m)
        self.make_interpolator(lla_ref, width_m, height_m, step_m)

    def load_tiles(self, lla_ref, width_m, height_m):
        print "Notice: loading tiles"
        ll_ned = np.array([[-height_m*0.5, -width_m*0.5, 0.0]])
        ur_ned = np.array([[height_m*0.5, width_m*0.5, 0.0]])
        ll_lla = navpy.ned2lla(ll_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        ur_lla = navpy.ned2lla(ur_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        print ll_lla
        print ur_lla
        lat1, lon1 = lla_ll_corner( ll_lla[0], ll_lla[1] )
        lat2, lon2 = lla_ll_corner( ur_lla[0], ur_lla[1] )
        for lat in range(lat1, lat2+1):
            for lon in range(lon1, lon2+1):
                srtm = SRTM(lat, lon, '../srtm')
                if srtm.parse():
                    srtm.make_lla_interpolator()
                    #srtm.plot_raw()
                    tile_name = make_tile_name(lat, lon)
                    self.tile_dict[tile_name] = srtm

    def make_interpolator(self, lla_ref, width_m, height_m, step_m):
        print "Notice: constructing interpolator"
        rows = (height_m / step_m) + 1
        cols = (width_m / step_m) + 1
        #ned_pts = np.zeros((cols, rows))
        #for r in range(0,rows):
        #    for c in range(0,cols):
        #        idx = (cols*r)+c
        #        #va = self.srtm_z[idx]
        #        #if va == 65535 or va < 0 or va > 10000:
        #        #    va = 0.0
        #        #z = va
        #        ned_pts[c,r] = 0.0

        # build regularly gridded x,y coordinate list and ned_pts array
        n_list = np.linspace(-height_m*0.5, height_m*0.5, rows)
        e_list = np.linspace(-width_m*0.5, width_m*0.5, cols)
        #print "e's:", e_list
        #print "n's:", n_list
        ned_pts = []
        for e in e_list:
            for n in n_list:
                ned_pts.append( [n, e, 0] )

        # convert ned_pts list to lla coordinates (so it's not
        # necessarily an exact grid anymore, but we can now
        # interpolate elevations out of the lla interpolators for each
        # tile.
        navpy_pts = navpy.ned2lla(ned_pts, lla_ref[0], lla_ref[1], lla_ref[2])
        #print navpy_pts

        # build list of (lat, lon) points for doing actual lla
        # elevation lookup
        ll_pts = []
        for i in range( len(navpy_pts[0]) ):
            lat = navpy_pts[0][i]
            lon = navpy_pts[1][i]
            ll_pts.append( [ lon, lat ] )
        #print "ll_pts:", ll_pts

        # set all the elevations in the ned_ds list to the extreme
        # minimum value. (rows,cols) might seem funny, but (ne)d is
        # reversed from (xy)z ... don't think about it too much or
        # you'll get a headache. :-)
        ned_ds = np.zeros((rows,cols))
        ned_ds[:][:] = -32768
        #print "ned_ds:", ned_ds

        # for each tile loaded, interpolate as many elevation values
        # as we can, then copy the good values into ned_ds.  When we
        # finish all the loaded tiles, we should have elevations for
        # the entire range of points.
        for tile in self.tile_dict:
            zs = self.tile_dict[tile].lla_interpolate(np.array(ll_pts))
            #print zs
            # copy the good altitudes back to the corresponding ned points
            for r in range(0,rows):
                for c in range(0,cols):
                    idx = (rows*c)+r
                    if zs[idx] > -10000:
                        ned_ds[r,c] = zs[idx]

        # quick sanity check
        for r in range(0,rows):
            for c in range(0,cols):
                idx = (rows*c)+r
                if ned_ds[r,c] < -10000:
                    print "Problem interpolating elevation for:", ll_pts[idx]
                    ned_ds[r,c] = 0.0
        #print "ned_ds:", ned_ds

        # now finally build the actual grid interpolator with evenly
        # spaced ned n, e values and elevations interpolated out of
        # the srtm lla interpolator.
        self.interp = scipy.interpolate.RegularGridInterpolator((n_list, e_list), ned_ds, bounds_error=False, fill_value=-32768)

        do_plot = False
        if do_plot:
            imshow(ned_ds, interpolation='bilinear', origin='lower', cmap=cm.gray, alpha=1.0)
            grid(False)
            show()

        do_test = False
        if do_test:
            for i in range(40):
                ned = [(random.random()-0.5)*height_m,
                       (random.random()-0.5)*width_m,
                       0.0]
                lla = navpy.ned2lla(ned, lla_ref[0], lla_ref[1], lla_ref[2])
                #print "ned=%s, lla=%s" % (ned, lla)
                nedz = self.interp([ned[0], ned[1]])
                tile = make_tile_name(lla[0], lla[1])
                llaz = self.tile_dict[tile].lla_interpolate(np.array([lla[1], lla[0]]))
                print "nedz=%.2f llaz=%.2f" % (nedz, llaz)

    # while error > eps: find altitude at current point, new pt = proj
    # vector to current alt.
    def interpolate_vector(self, pose_dict, v):
        pose = pose_dict['ned']
        p = pose[:] # copy hopefully

        # sanity check (always assume camera pose is above ground!)
        if v[2] <= 0.0:
            return p

        eps = 0.01

        #print "start:", p
        #print "vec:", v
        #print "pose:", pose
        ground = self.interp([p[0], p[1]])
        error = abs(p[2] + ground[0])
        #print "  p=%s ground=%s error=%s" % (p, ground, error)
        while error > eps:
            d_proj = -(pose[2] + ground[0])
            factor = d_proj / v[2]
            n_proj = v[0] * factor
            e_proj = v[1] * factor
            #print "proj = %s %s" % (n_proj, e_proj)
            p = [ pose[0] + n_proj, pose[1] + e_proj, pose[2] + d_proj ]
            #print "new p:", p
            ground = self.interp([p[0], p[1]])
            error = abs(p[2] + ground[0])
            #print "  p=%s ground=%.2f error = %.3f" % (p, ground, error)
        return p

    def interpolate_vectors(self, pose, v_list):
        pose_ned = pose['ned']
        pt_list = []
        for v in v_list:
            p = self.interpolate_vector(pose, v)
            pt_list.append(p)
        return pt_list

<commit_msg>comment++<commit_after># a class to manage SRTM surfaces

import json
import numpy as np
import os
from pylab import *
import random
import scipy.interpolate
import struct
import urllib
import zipfile

import navpy

# return the lower left corner of the 1x1 degree tile containing
# the specified lla coordinate
def lla_ll_corner(lat_deg, lon_deg):
    return int(floor(lat_deg)), int(floor(lon_deg))

# return the tile base name for the specified coordinate
def make_tile_name(lat, lon):
    ll_lat, ll_lon = lla_ll_corner(lat, lon)
    if ll_lat < 0:
        slat = "S%2d" % ll_lat
    else:
        slat = "N%2d" % ll_lat
    if ll_lon < 0:
        slon = "W%03d" % -ll_lon
    else:
        slon = "W%03d" % ll_lon
    return slat + slon

class SRTM():
    def __init__(self, lat, lon, dict_path):
        self.lat, self.lon = lla_ll_corner(lat, lon)
        self.srtm_dict = {}
        self.srtm_cache_dir = '/var/tmp' # unless set otherwise
        self.srtm_z = None
        self.i = None
        self.load_srtm_dict(dict_path)

    # load the directory download dictionary (mapping a desired file
    # to a download path.)
    def load_srtm_dict(self, dict_path):
        dict_file = dict_path + '/srtm.json'
        try:
            f = open(dict_file, 'r')
            self.srtm_dict = json.load(f)
            f.close()
        except:
            print "Notice: unable to read =", dict_file

    # if we'd like a persistant place to cache srtm files so they don't
    # need to be re-download every time we run anything
    def set_srtm_cache_dir(self, srtm_cache_dir):
        if not os.path.exists(srtm_cache_dir):
            print "Notice: srtm cache path doesn't exist =", srtm_cache_dir
        else:
            self.srtm_cache_dir = srtm_cache_dir

    # download and extract srtm file into cache directory
    def download_srtm(self, fileroot):
        if fileroot in self.srtm_dict:
            url = self.srtm_dict[fileroot]
            download_file = self.srtm_cache_dir + '/' + fileroot + '.hgt.zip'
            print "Notice: downloading:", url
            file = urllib.URLopener()
            file.retrieve(url, download_file)
            #print "Extracting:", download_file
            #zip = zipfile.ZipFile(download_file)
            #zip.extractall(self.srtm_cache_dir)
            #os.remove(download_file)
            return True
        else:
            print "Notice: requested srtm that is outside catalog"
            return False

    def parse(self):
        tilename = make_tile_name(self.lat, self.lon)
        cache_file = self.srtm_cache_dir + '/' + tilename + '.hgt.zip'
        if not os.path.exists(cache_file):
            if not self.download_srtm(tilename):
                return False
        print "Notice: parsing SRTM file:", cache_file
        #f = open(cache_file, "rb")
        zip = zipfile.ZipFile(cache_file)
        f = zip.open(tilename + '.hgt', 'r')
        contents = f.read()
        f.close()
        # read 1,442,401 (1201x1201) high-endian
        # signed 16-bit words into self.z
        self.srtm_z = struct.unpack(">1442401H", contents)
        return True

    def make_lla_interpolator(self):
        print "Notice: constructing lla interpolator"

        srtm_pts = np.zeros((1201, 1201))
        for r in range(0,1201):
            for c in range(0,1201):
                idx = (1201*r)+c
                va = self.srtm_z[idx]
                if va == 65535 or va < 0 or va > 10000:
                    va = 0.0
                z = va
                srtm_pts[c,1200-r] = z
        x = np.linspace(self.lon, self.lon+1, 1201)
        y = np.linspace(self.lat, self.lat+1, 1201)
        #print x
        #print y
        self.lla_interp = scipy.interpolate.RegularGridInterpolator((x, y), srtm_pts, bounds_error=False, fill_value=-32768)
        #print self.lla_interp([-93.14530573529404, 45.220697421008396])
        #for i in range(20):
        #    x = -94.0 + random.random()
        #    y = 45 + random.random()
        #    z = self.lla_interp([x,y])
        #    print [x, y, z[0]]

    # notice: this is *really* slow.  Adding the lla2ned conversion
    # makes even more slow.  I think what I want to do is build as
    # many SRTM instances as we need to cover an area, then resample
    # just the needed area, at some reasonable grid size, convert just
    # the resampled points to ned, and build the interpolator out of
    # the subset of points.  A DEM file can cover 60nm in latitude and
    # 60nm*cos(lat) in the longitude direction.

    def make_ned_interpolator(self, ref):
        print "Notice: constructing ned interpolator"

        # The LinearNDInterpolator works well, but is slow to
        # construct, ned data is not 'regularly' gridded so we need to
        # use this fancier interpolator
        ned_pts = np.zeros((1201*1201, 2))
        ned_vals = np.zeros((1201*1201))
        for r in range(0,1201):
            print "%.1f%%" % (float(r)/float(12.01))
            for c in range(0,1201):
                idx = (1201*r)+c
                va = self.srtm_z[idx]
                if va == 65535 or va < 0 or va > 10000:
                    va = 0.0
                lon = -94.0 + (float(c) / 1200.0)
                lat = 45 + (float(1200-r) / 1200.0)
                alt = va
                ned = navpy.lla2ned( lat, lon, alt, ref[0], ref[1], ref[2] )
                ned_pts[idx] = ned[0:2]
                ned_vals[idx] = ned[2]
        self.ned_interp = scipy.interpolate.LinearNDInterpolator(ned_pts, ned_vals)

    def lla_interpolate(self, point_list):
        return self.lla_interp(point_list)

    def ned_interpolate(self, point_list):
        return self.ned_interp(point_list)

    def plot_raw(self):
        zzz = np.zeros((1201,1201))
        for r in range(0,1201):
            for c in range(0,1201):
                va=self.srtm_z[(1201*r)+c]
                if (va==65535 or va<0 or va>10000):
                    va=0.0
                zzz[r][c]=float(va)

        #zz=np.log1p(zzz)
        imshow(zzz, interpolation='bilinear',cmap=cm.gray,alpha=1.0)
        grid(False)
        show()

# Build a gridded elevation interpolation table centered at lla_ref
# with width and height.  This is a little bit of quick feet dancing,
# but allows areas to span corners or edges of srtm tiles and attempts
# to stay on a fast path of regular grids, even though a regularly lla
# grid != a regular ned grid.
class NEDGround():
    def __init__(self, lla_ref, width_m, height_m, step_m):
        self.tile_dict = {}
        self.load_tiles(lla_ref, width_m, height_m)
        self.make_interpolator(lla_ref, width_m, height_m, step_m)

    def load_tiles(self, lla_ref, width_m, height_m):
        print "Notice: loading tiles"
        ll_ned = np.array([[-height_m*0.5, -width_m*0.5, 0.0]])
        ur_ned = np.array([[height_m*0.5, width_m*0.5, 0.0]])
        ll_lla = navpy.ned2lla(ll_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        ur_lla = navpy.ned2lla(ur_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        print ll_lla
        print ur_lla
        lat1, lon1 = lla_ll_corner( ll_lla[0], ll_lla[1] )
        lat2, lon2 = lla_ll_corner( ur_lla[0], ur_lla[1] )
        for lat in range(lat1, lat2+1):
            for lon in range(lon1, lon2+1):
                srtm = SRTM(lat, lon, '../srtm')
                if srtm.parse():
                    srtm.make_lla_interpolator()
                    #srtm.plot_raw()
                    tile_name = make_tile_name(lat, lon)
                    self.tile_dict[tile_name] = srtm

    def make_interpolator(self, lla_ref, width_m, height_m, step_m):
        print "Notice: constructing interpolator"
        rows = (height_m / step_m) + 1
        cols = (width_m / step_m) + 1
        #ned_pts = np.zeros((cols, rows))
        #for r in range(0,rows):
        #    for c in range(0,cols):
        #        idx = (cols*r)+c
        #        #va = self.srtm_z[idx]
        #        #if va == 65535 or va < 0 or va > 10000:
        #        #    va = 0.0
        #        #z = va
        #        ned_pts[c,r] = 0.0

        # build regularly gridded x,y coordinate list and ned_pts array
        n_list = np.linspace(-height_m*0.5, height_m*0.5, rows)
        e_list = np.linspace(-width_m*0.5, width_m*0.5, cols)
        #print "e's:", e_list
        #print "n's:", n_list
        ned_pts = []
        for e in e_list:
            for n in n_list:
                ned_pts.append( [n, e, 0] )

        # convert ned_pts list to lla coordinates (so it's not
        # necessarily an exact grid anymore, but we can now
        # interpolate elevations out of the lla interpolators for each
        # tile.
        navpy_pts = navpy.ned2lla(ned_pts, lla_ref[0], lla_ref[1], lla_ref[2])
        #print navpy_pts

        # build list of (lat, lon) points for doing actual lla
        # elevation lookup
        ll_pts = []
        for i in range( len(navpy_pts[0]) ):
            lat = navpy_pts[0][i]
            lon = navpy_pts[1][i]
            ll_pts.append( [ lon, lat ] )
        #print "ll_pts:", ll_pts

        # set all the elevations in the ned_ds list to the extreme
        # minimum value. (rows,cols) might seem funny, but (ne)d is
        # reversed from (xy)z ... don't think about it too much or
        # you'll get a headache. :-)
        ned_ds = np.zeros((rows,cols))
        ned_ds[:][:] = -32768
        #print "ned_ds:", ned_ds

        # for each tile loaded, interpolate as many elevation values
        # as we can, then copy the good values into ned_ds.  When we
        # finish all the loaded tiles, we should have elevations for
        # the entire range of points.
        for tile in self.tile_dict:
            zs = self.tile_dict[tile].lla_interpolate(np.array(ll_pts))
            #print zs
            # copy the good altitudes back to the corresponding ned points
            for r in range(0,rows):
                for c in range(0,cols):
                    idx = (rows*c)+r
                    if zs[idx] > -10000:
                        ned_ds[r,c] = zs[idx]

        # quick sanity check
        for r in range(0,rows):
            for c in range(0,cols):
                idx = (rows*c)+r
                if ned_ds[r,c] < -10000:
                    print "Problem interpolating elevation for:", ll_pts[idx]
                    ned_ds[r,c] = 0.0
        #print "ned_ds:", ned_ds

        # now finally build the actual grid interpolator with evenly
        # spaced ned n, e values and elevations interpolated out of
        # the srtm lla interpolator.
        self.interp = scipy.interpolate.RegularGridInterpolator((n_list, e_list), ned_ds, bounds_error=False, fill_value=-32768)

        do_plot = False
        if do_plot:
            imshow(ned_ds, interpolation='bilinear', origin='lower', cmap=cm.gray, alpha=1.0)
            grid(False)
            show()

        do_test = False
        if do_test:
            for i in range(40):
                ned = [(random.random()-0.5)*height_m,
                       (random.random()-0.5)*width_m,
                       0.0]
                lla = navpy.ned2lla(ned, lla_ref[0], lla_ref[1], lla_ref[2])
                #print "ned=%s, lla=%s" % (ned, lla)
                nedz = self.interp([ned[0], ned[1]])
                tile = make_tile_name(lla[0], lla[1])
                llaz = self.tile_dict[tile].lla_interpolate(np.array([lla[1], lla[0]]))
                print "nedz=%.2f llaz=%.2f" % (nedz, llaz)

    # while error > eps: find altitude at current point, new pt = proj
    # vector to current alt.
    def interpolate_vector(self, pose_dict, v):
        pose = pose_dict['ned']
        p = pose[:] # copy hopefully

        # sanity check (always assume camera pose is above ground!)
        if v[2] <= 0.0:
            return p

        eps = 0.01

        #print "start:", p
        #print "vec:", v
        #print "pose:", pose
        ground = self.interp([p[0], p[1]])
        error = abs(p[2] + ground[0])
        #print "  p=%s ground=%s error=%s" % (p, ground, error)
        while error > eps:
            d_proj = -(pose[2] + ground[0])
            factor = d_proj / v[2]
            n_proj = v[0] * factor
            e_proj = v[1] * factor
            #print "proj = %s %s" % (n_proj, e_proj)
            p = [ pose[0] + n_proj, pose[1] + e_proj, pose[2] + d_proj ]
            #print "new p:", p
            ground = self.interp([p[0], p[1]])
            error = abs(p[2] + ground[0])
            #print "  p=%s ground=%.2f error = %.3f" % (p, ground, error)
        return p

    # return a list of (3d) ground intersection points for the give
    # vector list and camera pose.
    def interpolate_vectors(self, pose, v_list):
        pose_ned = pose['ned']
        pt_list = []
        for v in v_list:
            p = self.interpolate_vector(pose, v)
            pt_list.append(p)
        return pt_list
<|endoftext|>


<commit_before>import asyncio
import logging
import json
from math import ceil

import aiohttp
import requests

from user_sync.error import AssertionException

class SignClient:
    version = 'v5'
    _endpoint_template = 'api/rest/{}/'

    def __init__(self, connection, host, integration_key, admin_email, logger=None):
        self.host = host
        self.integration_key = integration_key
        self.admin_email = admin_email
        self.api_url = None
        self.groups = None
        self.max_sign_retries = connection.get('retry_count') or 5
        self.concurrency_limit = connection.get('request_concurrency') or 1
        timeout = connection.get('timeout') or 120
        self.batch_size = connection.get('batch_size') or 10000
        self.logger = logger or logging.getLogger("sign_client_{}".format(self.integration_key[0:4]))
        logging.getLogger("urllib3").setLevel(logging.WARNING)
        self.timeout = aiohttp.ClientTimeout(total=None, sock_connect=timeout, sock_read=timeout)
        self.loop = asyncio.get_event_loop()
        self.users = {}

    def _init(self):
        self.api_url = self.base_uri()
        self.groups = self.get_groups()
        self.reverse_groups = {v: k for k, v in self.groups.items()}

    def sign_groups(self):
        if self.api_url is None or self.groups is None:
            self._init()
        return self.groups

    def header(self):
        """
        Return Sign API auth header
        :return: dict()
        """
        if self.version == 'v6':
            return {
                "Authorization": "Bearer {}".format(self.integration_key)
            }
        return {
            "Access-Token": self.integration_key
        }

    def header_json(self):
        """
        Get auth headers with options to PUT/POST JSON
        :return: dict()
        """

        json_headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
        }
        json_headers.update(self.header())
        return json_headers

    def base_uri(self):
        """
        This function validates that the SIGN integration key is valid.
        :return: dict()
        """

        endpoint = self._endpoint_template.format(self.version)
        url = 'https://' + self.host + '/' + endpoint

        if self.version == "v6":
            url_path = 'baseUris'
            access_point_key = 'apiAccessPoint'
        else:
            url_path = 'base_uris'
            access_point_key = 'api_access_point'

        result = requests.get(url + url_path, headers=self.header())
        if result.status_code != 200:
            raise AssertionException(
                "Error getting base URI from Sign API, is API key valid? (error: {}, reason: {}, {})".format
                (result.status_code, result.reason, result.content))

        if access_point_key not in result.json():
            raise AssertionException("Error getting base URI for Sign API, result invalid")

        return result.json()[access_point_key] + endpoint

    def get_users(self):
        """
        Get list of all users from Sign (indexed by email address)
        :return: dict()
        """
        if self.api_url is None or self.groups is None:
            self._init()

        users = {}
        self.logger.debug('getting list of all Sign users')
        user_list, users_res = self.call_with_retry_sync('GET', self.api_url + 'users', self.header())
        user_ids = [u['userId'] for u in user_list['userInfoList']]
        self._handle_calls(self._get_user, self.header(), user_ids)
        if users_res != 200:
            raise AssertionException(
                "Error retrieving Sign user list, (error: {}, reason: {}, {})".format(
                    users_res.status_code, users_res.reason, users_res.content))
        # fix line above
        for user_id in map(lambda u: u['userId'], user_list['userInfoList']):
            user_res = requests.get(self.api_url + 'users/' + user_id, headers=self.header())
            if users_res != 200:
                raise AssertionException("Error retrieving details for Sign user '{}'(error: {}, reason: {}, {})".format
                                         (user_id, users_res.status_code, users_res.reason, users_res.content))
            user = user_res.json()
            if user['userStatus'] != 'ACTIVE':
                continue
            if user['email'] == self.admin_email:
                continue
            user['userId'] = user_id
            user['roles'] = self.user_roles(user)
            users[user['email']] = user
            self.logger.debug('retrieved user details for Sign user {}'.format(user['email']))

        # return users
        return self.users

    def get_groups(self):
        """
        API request to get group information
        :return: dict()
        """
        if self.api_url is None:
            self.api_url = self.base_uri()

        res = requests.get(self.api_url + 'groups', headers=self.header())
        if res.status_code != 200:
            raise AssertionException("Error retrieving Sign group list(error: {}, reason: {}, {})".format(
                res.status_code, res.reason, res.content))
        groups = {}
        sign_groups = res.json()
        for group in sign_groups['groupInfoList']:
            groups[group['groupName'].lower()] = group['groupId']
        return groups

    def create_group(self, group):
        """
        Create a new group in Sign
        :param group: str
        :return:
        """
        if self.api_url is None or self.groups is None:
            self._init()

        res, code = self.call_with_retry_sync('POST', self.api_url + 'groups', self.header_json(), data=json.dumps({'groupName': group}))
        if code != 201:
            raise AssertionException("Failed to create Sign group '{}' (reason: {})".format(group, res.reason))
        self.groups[group] = res['groupId']

    def update_user(self, user_id, data):
        """
        Update Sign user
        :param user_id: str
        :param data: dict()
        :return: dict()
        """
        if self.api_url is None or self.groups is None:
            self._init()

        res = requests.put(self.api_url + 'users/' + user_id, headers=self.header_json(), data=json.dumps(data))
        if res.status_code != 200:
            raise AssertionException("Failed to update user '{}' (reason: {})".format(user_id, res.reason))

    @staticmethod
    def user_roles(user):
        """
        Resolve user roles
        :return: list[]
        """
        return ['NORMAL_USER'] if 'roles' not in user else user['roles']

    def insert_user(self, data):
        """
        Insert Sign user
        :param data: dict()
        """
        if self.api_url is None or self.groups is None:
            self._init()

        res = requests.post(self.api_url + 'users', headers=self.header_json(), data=json.dumps(data))
        # Response status code 201 is successful insertion
        if res.status_code != 201:
                exp_obj = json.loads(res.text)
                exp_obj['reason'] = res.reason
                exp_obj['status_code'] = res.status_code
                raise AssertionException("Failed to insert user '{}' (error response: {})".format(data['email'], exp_obj))

    def deactivate_user(self, user_id):
        """
        Deactivate Sign user
        :param data: dict()
        """
        if self.api_url is None or self.groups is None:
            self._init()
        data = {
            "userStatus": 'INACTIVE'
        }
        res = requests.put(self.api_url + 'users/' + user_id + '/status', headers=self.header_json(), data=json.dumps(data))
        # Response status code 200 is successful update
        if res.status_code != 200:
                exp_obj = json.loads(res.text)
                exp_obj['reason'] = res.reason
                exp_obj['status_code'] = res.status_code
                raise AssertionException(exp_obj)

    def _handle_calls(self, handle, headers, objects):
        """
        Batches and executes handle for each of o in objects
        handle: reference to function which will be called
        headers: api headers (common to all requests)
        objects: list of objects, which will be iterated through - and handle called on each
        """

        if self.api_url is None or self.groups is None:
            self._init()

        # Execute calls by batches.  This reduces the memory stack, since we do not need to create all
        # coroutines before starting execution.  We call run_until_complete for each set until all sets have run
        set_number = 1
        batch_count = ceil(len(objects) / self.batch_size)
        for i in range(0, len(objects), self.batch_size):
            self.logger.info("{}s - batch {}/{}".format(handle.__name__, set_number, batch_count))
            self.loop.run_until_complete(self._await_calls(handle, headers, objects[i:i + self.batch_size]))
            set_number += 1

    async def _await_calls(self, handle, headers, objects):
        """
        Where we actually await the coroutines. Must be own method, in order to be handled by loop
        """

        if not objects:
            return

        # Semaphore specifies number of allowed calls at one time
        sem = asyncio.Semaphore(value=self.concurrency_limit)

        # We must use only 1 session, else will hang
        async with aiohttp.ClientSession(trust_env=True, timeout=self.timeout) as session:
            # prepare a list of calls to make * Note: calls are prepared by using call
            # syntax (eg, func() and not func), but they will not be run until executed by the wait
            # split into batches of self.bach_size to avoid taking too much memory
            calls = [handle(sem, o, headers, session) for o in objects]
            await asyncio.wait(calls)

    async def _get_user(self, semaphore, user_id, header, session):

        # This will block the method from executing until a position opens
        async with semaphore:
            user_url = self.api_url + 'users/' + user_id
            user, code = await self.call_with_retry_async('GET', user_url, header, session=session)
            if code != 200:
                self.logger.error("Error fetching user '{}' with response: {}".format(user_id, user))
                return
            if user['userStatus'] != 'ACTIVE':
                return
            if user['email'] == self.admin_email:
                return
            user['userId'] = user_id
            user['roles'] = self.user_roles(user)
            self.users[user['email']] = user
            self.logger.debug('retrieved user details for Sign user {}'.format(user['email']))

    async def _update_user(self, semaphore, user, headers, session):
        """
        Update Sign user
        """
        # This will block the method from executing until a position opens
        async with semaphore:
            url = self.api_url + 'users/' + user['userId']
            group = self.reverse_groups[user['groupId']]
            body, code = await self.call_with_retry_async('PUT', url, headers, data=json.dumps(user), session=session)
            self.logger.info(
                "Updated Sign user '{}', Group: '{}', Roles: {}".format(user['email'], group, user['roles']))
            if code != 200:
                self.logger.error("Error updating user '{}' with response: {}".format(user['email'], body))

    def call_with_retry_sync(self, method, url, header, data=None):
        """
        Need to define this method, so that it can be called outside async context
        loop will execute a single synchronous call, but sharing code with the async retry method
        """
        return self.loop.run_until_complete(self.call_with_retry_async(method, url, header, data=data or {}))

    async def call_with_retry_async(self, method, url, header, data=None, session=None):
        """
        Call manager with exponential retry
        :return: Response <Response> object
        """
        retry_nb = 0
        waiting_time = 10
        close = session is None
        session = session or aiohttp.ClientSession(trust_env=True, timeout=self.timeout)
        session.headers.update(header)
        while True:
            try:
                waiting_time *= 3
                self.logger.debug('Attempt {} to call: {}'.format(retry_nb, url))
                async with session.request(method=method, url=url, data=data or {}) as r:
                    if r.status >= 500:
                        raise Exception('{}, Headers: {}'.format(r.status, r.headers))
                    elif r.status == 429:
                        raise Exception('{} - too many calls. Headers: {}'.format(r.status, r.headers))
                    # elif r.status > 400 and r.status < 500:
                    #     self.logger.critical(' {} - {}. Headers: {}'.format(r.status, r.reason, r.headers))
                    #     raise AssertionException('')
                    body = await r.json()
                    return body, r.status
            except Exception as exp:
                retry_nb += 1
                self.logger.warning('Failed: {} - {}'.format(type(exp), exp.args))
                if retry_nb == (self.max_sign_retries + 1):
                    raise AssertionException('Quitting after {} retries'.format(self.max_sign_retries))
                self.logger.warning('Waiting for {} seconds'.format(waiting_time))
                await asyncio.sleep(waiting_time)
            finally:
                if close:
                    await session.close()<commit_msg><commit_after>import asyncio
import logging
import json
from math import ceil

import aiohttp
import requests

from user_sync.error import AssertionException

class SignClient:
    version = 'v5'
    _endpoint_template = 'api/rest/{}/'

    def __init__(self, connection, host, integration_key, admin_email, logger=None):
        self.host = host
        self.integration_key = integration_key
        self.admin_email = admin_email
        self.api_url = None
        self.groups = None
        self.max_sign_retries = connection.get('retry_count') or 5
        self.concurrency_limit = connection.get('request_concurrency') or 1
        timeout = connection.get('timeout') or 120
        self.batch_size = connection.get('batch_size') or 10000
        self.logger = logger or logging.getLogger("sign_client_{}".format(self.integration_key[0:4]))
        logging.getLogger("urllib3").setLevel(logging.WARNING)
        self.timeout = aiohttp.ClientTimeout(total=None, sock_connect=timeout, sock_read=timeout)
        self.loop = asyncio.get_event_loop()
        self.users = {}

    def _init(self):
        self.api_url = self.base_uri()
        self.groups = self.get_groups()
        self.reverse_groups = {v: k for k, v in self.groups.items()}

    def sign_groups(self):
        if self.api_url is None or self.groups is None:
            self._init()
        return self.groups

    def header(self):
        """
        Return Sign API auth header
        :return: dict()
        """
        if self.version == 'v6':
            return {
                "Authorization": "Bearer {}".format(self.integration_key)
            }
        return {
            "Access-Token": self.integration_key
        }

    def header_json(self):
        """
        Get auth headers with options to PUT/POST JSON
        :return: dict()
        """

        json_headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
        }
        json_headers.update(self.header())
        return json_headers

    def base_uri(self):
        """
        This function validates that the SIGN integration key is valid.
        :return: dict()
        """

        endpoint = self._endpoint_template.format(self.version)
        url = 'https://' + self.host + '/' + endpoint

        if self.version == "v6":
            url_path = 'baseUris'
            access_point_key = 'apiAccessPoint'
        else:
            url_path = 'base_uris'
            access_point_key = 'api_access_point'

        result = requests.get(url + url_path, headers=self.header())
        if result.status_code != 200:
            raise AssertionException(
                "Error getting base URI from Sign API, is API key valid? (error: {}, reason: {}, {})".format
                (result.status_code, result.reason, result.content))

        if access_point_key not in result.json():
            raise AssertionException("Error getting base URI for Sign API, result invalid")

        return result.json()[access_point_key] + endpoint

    def get_users(self):
        """
        Gets the full user list, and then extracts the user ID's for making calls
        We return self.users because it will be filled by the _get_user method.  This is
        necessary to avoid returning futures of calls which we cannot predict.
        """
        if self.api_url is None or self.groups is None:
            self._init()

        self.logger.info('Getting list of all Sign users')
        user_list, _ = self.call_with_retry_sync('GET', self.api_url + 'users', self.header())

        user_ids = [u['userId'] for u in user_list['userInfoList']]
        self._handle_calls(self._get_user, self.header(), user_ids)
        return self.users

    def update_users(self, users):
        """
        Passthrough for call handling
        """
        self._handle_calls(self._update_user, self.header_json(), users)

    def get_groups(self):
        """
        API request to get group information
        :return: dict()
        """
        if self.api_url is None:
            self.api_url = self.base_uri()

        url = self.api_url + 'groups'
        header = self.header()
        sign_groups, code = self.call_with_retry_sync('GET', url, header)
        self.logger.info('getting Sign user groups')
        groups = {}
        for group in sign_groups['groupInfoList']:
            groups[group['groupName'].lower()] = group['groupId']
        return groups

    def create_group(self, group):
        """
        Create a new group in Sign
        :param group: str
        :return:
        """
        if self.api_url is None or self.groups is None:
            self._init()
        url = self.api_url + 'groups'
        header = self.header_json()
        data = json.dumps({'groupName': group})
        self.logger.info('Creating Sign group {} '.format(group))
        res, code = self.call_with_retry_sync('POST', url, header, data)
        self.groups[group] = res['groupId']

    @staticmethod
    def user_roles(user):
        """
        Resolve user roles
        :return: list[]
        """
        return ['NORMAL_USER'] if 'roles' not in user else user['roles']

    def insert_user(self, data):
        """
        Insert Sign user
        :param data: dict()
        """
        if self.api_url is None or self.groups is None:
            self._init()

        res = requests.post(self.api_url + 'users', headers=self.header_json(), data=json.dumps(data))
        # Response status code 201 is successful insertion
        if res.status_code != 201:
                exp_obj = json.loads(res.text)
                exp_obj['reason'] = res.reason
                exp_obj['status_code'] = res.status_code
                raise AssertionException("Failed to insert user '{}' (error response: {})".format(data['email'], exp_obj))

    def deactivate_user(self, user_id):
        """
        Deactivate Sign user
        :param data: dict()
        """
        if self.api_url is None or self.groups is None:
            self._init()
        data = {
            "userStatus": 'INACTIVE'
        }
        res = requests.put(self.api_url + 'users/' + user_id + '/status', headers=self.header_json(), data=json.dumps(data))
        # Response status code 200 is successful update
        if res.status_code != 200:
                exp_obj = json.loads(res.text)
                exp_obj['reason'] = res.reason
                exp_obj['status_code'] = res.status_code
                raise AssertionException(exp_obj)

    def _handle_calls(self, handle, headers, objects):
        """
        Batches and executes handle for each of o in objects
        handle: reference to function which will be called
        headers: api headers (common to all requests)
        objects: list of objects, which will be iterated through - and handle called on each
        """

        if self.api_url is None or self.groups is None:
            self._init()

        # Execute calls by batches.  This reduces the memory stack, since we do not need to create all
        # coroutines before starting execution.  We call run_until_complete for each set until all sets have run
        set_number = 1
        batch_count = ceil(len(objects) / self.batch_size)
        for i in range(0, len(objects), self.batch_size):
            self.logger.info("{}s - batch {}/{}".format(handle.__name__, set_number, batch_count))
            self.loop.run_until_complete(self._await_calls(handle, headers, objects[i:i + self.batch_size]))
            set_number += 1

    async def _await_calls(self, handle, headers, objects):
        """
        Where we actually await the coroutines. Must be own method, in order to be handled by loop
        """

        if not objects:
            return

        # Semaphore specifies number of allowed calls at one time
        sem = asyncio.Semaphore(value=self.concurrency_limit)

        # We must use only 1 session, else will hang
        async with aiohttp.ClientSession(trust_env=True, timeout=self.timeout) as session:
            # prepare a list of calls to make * Note: calls are prepared by using call
            # syntax (eg, func() and not func), but they will not be run until executed by the wait
            # split into batches of self.bach_size to avoid taking too much memory
            calls = [handle(sem, o, headers, session) for o in objects]
            await asyncio.wait(calls)

    async def _get_user(self, semaphore, user_id, header, session):

        # This will block the method from executing until a position opens
        async with semaphore:
            user_url = self.api_url + 'users/' + user_id
            user, code = await self.call_with_retry_async('GET', user_url, header, session=session)
            if code != 200:
                self.logger.error("Error fetching user '{}' with response: {}".format(user_id, user))
                return
            if user['userStatus'] != 'ACTIVE':
                return
            if user['email'] == self.admin_email:
                return
            user['userId'] = user_id
            user['roles'] = self.user_roles(user)
            self.users[user['email']] = user
            self.logger.debug('retrieved user details for Sign user {}'.format(user['email']))

    async def _update_user(self, semaphore, user, headers, session):
        """
        Update Sign user
        """
        # This will block the method from executing until a position opens
        async with semaphore:
            url = self.api_url + 'users/' + user['userId']
            group = self.reverse_groups[user['groupId']]
            body, code = await self.call_with_retry_async('PUT', url, headers, data=json.dumps(user), session=session)
            self.logger.info(
                "Updated Sign user '{}', Group: '{}', Roles: {}".format(user['email'], group, user['roles']))
            if code != 200:
                self.logger.error("Error updating user '{}' with response: {}".format(user['email'], body))

    def call_with_retry_sync(self, method, url, header, data=None):
        """
        Need to define this method, so that it can be called outside async context
        loop will execute a single synchronous call, but sharing code with the async retry method
        """
        return self.loop.run_until_complete(self.call_with_retry_async(method, url, header, data=data or {}))

    async def call_with_retry_async(self, method, url, header, data=None, session=None):
        """
        Call manager with exponential retry
        :return: Response <Response> object
        """
        retry_nb = 0
        waiting_time = 10
        close = session is None
        session = session or aiohttp.ClientSession(trust_env=True, timeout=self.timeout)
        session.headers.update(header)
        while True:
            try:
                waiting_time *= 3
                self.logger.debug('Attempt {} to call: {}'.format(retry_nb, url))
                async with session.request(method=method, url=url, data=data or {}) as r:
                    if r.status >= 500:
                        raise Exception('{}, Headers: {}'.format(r.status, r.headers))
                    elif r.status == 429:
                        raise Exception('{} - too many calls. Headers: {}'.format(r.status, r.headers))
                    # elif r.status > 400 and r.status < 500:
                    #     self.logger.critical(' {} - {}. Headers: {}'.format(r.status, r.reason, r.headers))
                    #     raise AssertionException('')
                    body = await r.json()
                    return body, r.status
            except Exception as exp:
                retry_nb += 1
                self.logger.warning('Failed: {} - {}'.format(type(exp), exp.args))
                if retry_nb == (self.max_sign_retries + 1):
                    raise AssertionException('Quitting after {} retries'.format(self.max_sign_retries))
                self.logger.warning('Waiting for {} seconds'.format(waiting_time))
                await asyncio.sleep(waiting_time)
            finally:
                if close:
                    await session.close()<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
Commands that update or process the application data.
"""
from datetime import datetime, timedelta
from itertools import count
import json
import os
import random
import shutil

import boto.dynamodb
from boto.dynamodb.condition import GE
import copytext
from dateutil.parser import parse
from fabric.api import env, local, settings, task
from facebook import GraphAPI
from twitter import Twitter, OAuth

import app_config
import admin_app
import servers

def server_postgres_command(cmd ):
    """
    Run a postgres command on the server.
    """
    local('export PGPASSWORD=$elections14_POSTGRES_PASSWORD && %s --username=$elections14_POSTGRES_USER --host=$elections14_POSTGRES_HOST --port=$elections14_POSTGRES_PORT' % cmd)

@task
def bootstrap():
    """
    Resets the local environment to a fresh copy of the db and data.
    """
    secrets = app_config.get_secrets()

    if env.settings:
        with settings(warn_only=True):
            service_name = servers._get_installed_service_name('uwsgi')
            local('sudo service %s stop' % service_name)

            server_postgres_command('dropdb %s' % app_config.PROJECT_SLUG)
            server_postgres_command('createdb %s' % app_config.PROJECT_SLUG)

            local('sudo service %s start' % service_name)
    else:
        with settings(warn_only=True):
            local('dropdb %s' % app_config.PROJECT_SLUG)
            local('echo "CREATE USER %s WITH PASSWORD \'%s\';" | psql' % (app_config.PROJECT_SLUG, secrets['POSTGRES_PASSWORD']))

        local('createdb -O %s %s' % (app_config.PROJECT_SLUG, app_config.PROJECT_SLUG))

    import models

    models.Race.create_table()
    models.Candidate.create_table()
    models.Slide.create_table()
    models.SlideSequence.create_table()

    admin_app.auth.User.create_table()
    admin_user = admin_app.auth.User(username='admin', email='', admin=True, active=True)
    admin_user.set_password(secrets.get('ADMIN_PASSWORD'))
    admin_user.save()

    mock_slides()

    with open('data/races.json') as f:
        races = json.load(f)

        for race in races:
            models.Race.create(
                state_postal = race['state_postal'],
                office_id = race['office_id'],
                office_name = race['office_name'],
                seat_name = race['seat_name'],
                seat_number = race['seat_number'],
                race_id = race['race_id'],
                race_type = race['race_type'],
                last_updated = race['last_updated'],
            )

    with open('data/candidates.json') as f:
        candidates = json.load(f)

        for candidate in candidates:
            models.Candidate.create(
                first_name = candidate['first_name'],
                last_name = candidate['last_name'],
                party = candidate['party'],
                race = models.Race.get(models.Race.race_id == candidate['race_id']),
                candidate_id = candidate['candidate_id'],
            )

@task(default=True)
def update(test=False):
    """
    Update models with elections data from interediary files.
    """
    import models

    #update_featured_social()

    if test:
        shutil.copyfile('data/fake_update.json', 'data/update.json')

    if not os.path.exists('data/update.json'):
        return

    update_flat = []

    with open('data/update.json') as f:
        races = json.load(f)

        for race in races:
            race_model = models.Race.get(models.Race.race_id == race['race_id'])

            # If race has not been updated, skip
            last_updated = parse(race['last_updated']).replace(tzinfo=None)

            if race_model.last_updated == last_updated:
                continue

            race_model.is_test = race['is_test']
            race_model.precincts_reporting = race['precincts_reporting']
            race_model.precincts_total = race['precincts_total']
            race_model.last_updated = last_updated

            race_model.save()

            for candidate in race['candidates']:
                # Select candidate by candidate_id AND race_id, since they can appear in multiple races
                candidate_model = models.Candidate.get(models.Candidate.candidate_id == candidate['candidate_id'], models.Candidate.race == race_model)

                candidate_model.vote_count = candidate['vote_count']
                candidate_model.ap_winner = candidate.get('ap_winner', False)

                candidate_model.save()

            update_flat.append(race_model.flatten(update_only=True))

    with open('www/live-data/update.json', 'w') as f:
        json.dump(update_flat, f, cls=models.ModelEncoder)

    print 'Updated %i races' % len(update_flat)
    print 'Updated %i candidates' % sum([len(race['candidates']) for race in update_flat])

@task
def get_quiz_answers():
    """
    Read all quiz answers from Dynamo.
    """
    conn = boto.dynamodb.connect_to_region('us-west-2')

    table = conn.get_table('elections14-game')

    recent_answers = table.scan({ 'timestamp': GE('1410800598040') })

    print recent_answers.count

    for answer in recent_answers:
        print answer

@task
def update_featured_social():
    """
    Update featured tweets
    """
    COPY = copytext.Copy(app_config.COPY_PATH)
    secrets = app_config.get_secrets()

    # Twitter
    print 'Fetching tweets...'

    twitter_api = Twitter(
        auth=OAuth(
            secrets['TWITTER_API_OAUTH_TOKEN'],
            secrets['TWITTER_API_OAUTH_SECRET'],
            secrets['TWITTER_API_CONSUMER_KEY'],
            secrets['TWITTER_API_CONSUMER_SECRET']
        )
    )

    tweets = []

    for i in range(1, 4):
        tweet_url = COPY['share']['featured_tweet%i' % i]

        if isinstance(tweet_url, copytext.Error) or unicode(tweet_url).strip() == '':
            continue

        tweet_id = unicode(tweet_url).split('/')[-1]

        tweet = twitter_api.statuses.show(id=tweet_id)

        creation_date = datetime.strptime(tweet['created_at'],'%a %b %d %H:%M:%S +0000 %Y')
        creation_date = '%s %i' % (creation_date.strftime('%b'), creation_date.day)

        tweet_url = 'http://twitter.com/%s/status/%s' % (tweet['user']['screen_name'], tweet['id'])

        photo = None
        html = tweet['text']
        subs = {}

        for media in tweet['entities'].get('media', []):
            original = tweet['text'][media['indices'][0]:media['indices'][1]]
            replacement = '%s' % (media['url'], app_config.PROJECT_SLUG, tweet_url, media['display_url'])

            subs[original] = replacement

            if media['type'] == 'photo' and not photo:
                photo = {
                    'url': media['media_url']
                }

        for url in tweet['entities'].get('urls', []):
            original = tweet['text'][url['indices'][0]:url['indices'][1]]
            replacement = '%s' % (url['url'], app_config.PROJECT_SLUG, tweet_url, url['display_url'])

            subs[original] = replacement

        for hashtag in tweet['entities'].get('hashtags', []):
            original = tweet['text'][hashtag['indices'][0]:hashtag['indices'][1]]
            replacement = '%s' % (hashtag['text'], app_config.PROJECT_SLUG, tweet_url, '#%s' % hashtag['text'])

            subs[original] = replacement

        for original, replacement in subs.items():
            html =  html.replace(original, replacement)

        # https://dev.twitter.com/docs/api/1.1/get/statuses/show/%3Aid
        tweets.append({
            'id': tweet['id'],
            'url': tweet_url,
            'html': html,
            'favorite_count': tweet['favorite_count'],
            'retweet_count': tweet['retweet_count'],
            'user': {
                'id': tweet['user']['id'],
                'name': tweet['user']['name'],
                'screen_name': tweet['user']['screen_name'],
                'profile_image_url': tweet['user']['profile_image_url'],
                'url': tweet['user']['url'],
            },
            'creation_date': creation_date,
            'photo': photo
        })

    # Facebook
    print 'Fetching Facebook posts...'

    fb_api = GraphAPI(secrets['FACEBOOK_API_APP_TOKEN'])

    facebook_posts = []

    for i in range(1, 4):
        fb_url = COPY['share']['featured_facebook%i' % i]

        if isinstance(fb_url, copytext.Error) or unicode(fb_url).strip() == '':
            continue

        fb_id = unicode(fb_url).split('/')[-1]

        post = fb_api.get_object(fb_id)
        user  = fb_api.get_object(post['from']['id'])
        user_picture = fb_api.get_object('%s/picture' % post['from']['id'])
        likes = fb_api.get_object('%s/likes' % fb_id, summary='true')
        comments = fb_api.get_object('%s/comments' % fb_id, summary='true')
        #shares = fb_api.get_object('%s/sharedposts' % fb_id)

        creation_date = datetime.strptime(post['created_time'],'%Y-%m-%dT%H:%M:%S+0000')
        creation_date = '%s %i' % (creation_date.strftime('%b'), creation_date.day)

        # https://developers.facebook.com/docs/graph-api/reference/v2.0/post
        facebook_posts.append({
            'id': post['id'],
            'message': post['message'],
            'link': {
                'url': post['link'],
                'name': post['name'],
                'caption': (post['caption'] if 'caption' in post else None),
                'description': post['description'],
                'picture': post['picture']
            },
            'from': {
                'name': user['name'],
                'link': user['link'],
                'picture': user_picture['url']
            },
            'likes': likes['summary']['total_count'],
            'comments': comments['summary']['total_count'],
            #'shares': shares['summary']['total_count'],
            'creation_date': creation_date
        })

    # Render to JSON
    output = {
        'tweets': tweets,
        'facebook_posts': facebook_posts
    }

    with open('data/featured.json', 'w') as f:
        json.dump(output, f)

def _mock_slide_from_image(filename, i):
    import models

    body = '' % (app_config.S3_BASE_URL, filename)
    name = 'Test slide: %s' % filename[0:-4]
    slide = models.Slide.create(body=body, name=name)
    models.SlideSequence.create(sequence=i, slide=slide)

@task
def mock_slides():
    """
    Load mockup slides from assets directory.
    """
    import models

    models.SlideSequence.delete().execute()
    models.Slide.delete().execute()

    it = count()
    _mock_slide_from_image('welcome.png', it.next())
    _mock_slide_from_image('senate.png', it.next())
    _mock_slide_from_image('gif1.gif', it.next())
    _mock_slide_from_image('party_pix.png', it.next())

@task
def mock_election_results():
    """
    Fake out some election results
    """
    import models

    for race in models.Race.select():
        _fake_incumbent(race)
        _fake_poll_closing_time(race)
        _fake_precincts_reporting(race)
        _fake_called_status(race)
        _fake_results(race)
        race.save()

def _fake_incumbent(race):
    """
    Fake one incumbent for a race
    """
    import models
    candidates = race.candidates.where((models.Candidate.party == "GOP") | (models.Candidate.party == "Dem"))
    try:
        candidate = candidates[random.randint(0,1)]
        candidate.incumbent = True
        candidate.save()
    except IndexError:
        pass

def _fake_poll_closing_time(race):
    """
    Fake poll closing time
    """
    first_close = datetime(2014, 11, 4, 7)
    closing_times = [first_close + timedelta(hours=delta) for delta in range(6)]
    race.poll_closing_time = random.choice(closing_times)

def _fake_precincts_reporting(race):
    """
    Fake precincts reporting
    """
    race.precincts_total = random.randint(2000, 4000)
    if random.choice([True, False, False]):
        race.precincts_reporting = random.randint(200, race.precincts_total)
    else:
        race.precincts_reporting = 0

def _fake_called_status(race):
    """
    Fake AP called status, requires race to have closing time
    """
    if race.precincts_reporting > 0:
        race.ap_called = random.choice([True, False, False, False])
        if race.ap_called:
            race.accept_ap_call = True
            race.ap_called_time = race.poll_closing_time + timedelta(hours=random.randint(1,3), minutes=random.randint(0,59))

def _fake_results(race):
    max_votes = 0
    for candidate in race.candidates:
        if candidate.party in ["GOP", "Dem"]:
            votes = random.randint(400000, 500000)
            candidate.vote_count = votes
            if votes > max_votes:
                max_candidate = candidate
                max_votes = votes
        else:
            candidate.vote_count = 0

        candidate.save()

    if race.ap_called:
        max_candidate.ap_winner = True
        max_candidate.save()
<commit_msg>refactor some data mocking functions<commit_after>#!/usr/bin/env python

"""
Commands that update or process the application data.
"""
from datetime import datetime, timedelta
from itertools import count
import json
import os
import random
import shutil

import boto.dynamodb
from boto.dynamodb.condition import GE
import copytext
from dateutil.parser import parse
from fabric.api import env, local, settings, task
from facebook import GraphAPI
from twitter import Twitter, OAuth

import app_config
import admin_app
import servers

def server_postgres_command(cmd ):
    """
    Run a postgres command on the server.
    """
    local('export PGPASSWORD=$elections14_POSTGRES_PASSWORD && %s --username=$elections14_POSTGRES_USER --host=$elections14_POSTGRES_HOST --port=$elections14_POSTGRES_PORT' % cmd)

@task
def bootstrap():
    """
    Resets the local environment to a fresh copy of the db and data.
    """
    secrets = app_config.get_secrets()

    if env.settings:
        with settings(warn_only=True):
            service_name = servers._get_installed_service_name('uwsgi')
            local('sudo service %s stop' % service_name)

            server_postgres_command('dropdb %s' % app_config.PROJECT_SLUG)
            server_postgres_command('createdb %s' % app_config.PROJECT_SLUG)

            local('sudo service %s start' % service_name)
    else:
        with settings(warn_only=True):
            local('dropdb %s' % app_config.PROJECT_SLUG)
            local('echo "CREATE USER %s WITH PASSWORD \'%s\';" | psql' % (app_config.PROJECT_SLUG, secrets['POSTGRES_PASSWORD']))

        local('createdb -O %s %s' % (app_config.PROJECT_SLUG, app_config.PROJECT_SLUG))

    import models

    models.Race.create_table()
    models.Candidate.create_table()
    models.Slide.create_table()
    models.SlideSequence.create_table()

    admin_app.auth.User.create_table()
    admin_user = admin_app.auth.User(username='admin', email='', admin=True, active=True)
    admin_user.set_password(secrets.get('ADMIN_PASSWORD'))
    admin_user.save()

    mock_slides()

    with open('data/races.json') as f:
        races = json.load(f)

        for race in races:
            models.Race.create(
                state_postal = race['state_postal'],
                office_id = race['office_id'],
                office_name = race['office_name'],
                seat_name = race['seat_name'],
                seat_number = race['seat_number'],
                race_id = race['race_id'],
                race_type = race['race_type'],
                last_updated = race['last_updated'],
            )

    with open('data/candidates.json') as f:
        candidates = json.load(f)

        for candidate in candidates:
            models.Candidate.create(
                first_name = candidate['first_name'],
                last_name = candidate['last_name'],
                party = candidate['party'],
                race = models.Race.get(models.Race.race_id == candidate['race_id']),
                candidate_id = candidate['candidate_id'],
            )

@task(default=True)
def update(test=False):
    """
    Update models with elections data from interediary files.
    """
    import models

    #update_featured_social()

    if test:
        shutil.copyfile('data/fake_update.json', 'data/update.json')

    if not os.path.exists('data/update.json'):
        return

    update_flat = []

    with open('data/update.json') as f:
        races = json.load(f)

        for race in races:
            race_model = models.Race.get(models.Race.race_id == race['race_id'])

            # If race has not been updated, skip
            last_updated = parse(race['last_updated']).replace(tzinfo=None)

            if race_model.last_updated == last_updated:
                continue

            race_model.is_test = race['is_test']
            race_model.precincts_reporting = race['precincts_reporting']
            race_model.precincts_total = race['precincts_total']
            race_model.last_updated = last_updated

            race_model.save()

            for candidate in race['candidates']:
                # Select candidate by candidate_id AND race_id, since they can appear in multiple races
                candidate_model = models.Candidate.get(models.Candidate.candidate_id == candidate['candidate_id'], models.Candidate.race == race_model)

                candidate_model.vote_count = candidate['vote_count']
                candidate_model.ap_winner = candidate.get('ap_winner', False)

                candidate_model.save()

            update_flat.append(race_model.flatten(update_only=True))

    with open('www/live-data/update.json', 'w') as f:
        json.dump(update_flat, f, cls=models.ModelEncoder)

    print 'Updated %i races' % len(update_flat)
    print 'Updated %i candidates' % sum([len(race['candidates']) for race in update_flat])

@task
def get_quiz_answers():
    """
    Read all quiz answers from Dynamo.
    """
    conn = boto.dynamodb.connect_to_region('us-west-2')

    table = conn.get_table('elections14-game')

    recent_answers = table.scan({ 'timestamp': GE('1410800598040') })

    print recent_answers.count

    for answer in recent_answers:
        print answer

@task
def update_featured_social():
    """
    Update featured tweets
    """
    COPY = copytext.Copy(app_config.COPY_PATH)
    secrets = app_config.get_secrets()

    # Twitter
    print 'Fetching tweets...'

    twitter_api = Twitter(
        auth=OAuth(
            secrets['TWITTER_API_OAUTH_TOKEN'],
            secrets['TWITTER_API_OAUTH_SECRET'],
            secrets['TWITTER_API_CONSUMER_KEY'],
            secrets['TWITTER_API_CONSUMER_SECRET']
        )
    )

    tweets = []

    for i in range(1, 4):
        tweet_url = COPY['share']['featured_tweet%i' % i]

        if isinstance(tweet_url, copytext.Error) or unicode(tweet_url).strip() == '':
            continue

        tweet_id = unicode(tweet_url).split('/')[-1]

        tweet = twitter_api.statuses.show(id=tweet_id)

        creation_date = datetime.strptime(tweet['created_at'],'%a %b %d %H:%M:%S +0000 %Y')
        creation_date = '%s %i' % (creation_date.strftime('%b'), creation_date.day)

        tweet_url = 'http://twitter.com/%s/status/%s' % (tweet['user']['screen_name'], tweet['id'])

        photo = None
        html = tweet['text']
        subs = {}

        for media in tweet['entities'].get('media', []):
            original = tweet['text'][media['indices'][0]:media['indices'][1]]
            replacement = '%s' % (media['url'], app_config.PROJECT_SLUG, tweet_url, media['display_url'])

            subs[original] = replacement

            if media['type'] == 'photo' and not photo:
                photo = {
                    'url': media['media_url']
                }

        for url in tweet['entities'].get('urls', []):
            original = tweet['text'][url['indices'][0]:url['indices'][1]]
            replacement = '%s' % (url['url'], app_config.PROJECT_SLUG, tweet_url, url['display_url'])

            subs[original] = replacement

        for hashtag in tweet['entities'].get('hashtags', []):
            original = tweet['text'][hashtag['indices'][0]:hashtag['indices'][1]]
            replacement = '%s' % (hashtag['text'], app_config.PROJECT_SLUG, tweet_url, '#%s' % hashtag['text'])

            subs[original] = replacement

        for original, replacement in subs.items():
            html =  html.replace(original, replacement)

        # https://dev.twitter.com/docs/api/1.1/get/statuses/show/%3Aid
        tweets.append({
            'id': tweet['id'],
            'url': tweet_url,
            'html': html,
            'favorite_count': tweet['favorite_count'],
            'retweet_count': tweet['retweet_count'],
            'user': {
                'id': tweet['user']['id'],
                'name': tweet['user']['name'],
                'screen_name': tweet['user']['screen_name'],
                'profile_image_url': tweet['user']['profile_image_url'],
                'url': tweet['user']['url'],
            },
            'creation_date': creation_date,
            'photo': photo
        })

    # Facebook
    print 'Fetching Facebook posts...'

    fb_api = GraphAPI(secrets['FACEBOOK_API_APP_TOKEN'])

    facebook_posts = []

    for i in range(1, 4):
        fb_url = COPY['share']['featured_facebook%i' % i]

        if isinstance(fb_url, copytext.Error) or unicode(fb_url).strip() == '':
            continue

        fb_id = unicode(fb_url).split('/')[-1]

        post = fb_api.get_object(fb_id)
        user  = fb_api.get_object(post['from']['id'])
        user_picture = fb_api.get_object('%s/picture' % post['from']['id'])
        likes = fb_api.get_object('%s/likes' % fb_id, summary='true')
        comments = fb_api.get_object('%s/comments' % fb_id, summary='true')
        #shares = fb_api.get_object('%s/sharedposts' % fb_id)

        creation_date = datetime.strptime(post['created_time'],'%Y-%m-%dT%H:%M:%S+0000')
        creation_date = '%s %i' % (creation_date.strftime('%b'), creation_date.day)

        # https://developers.facebook.com/docs/graph-api/reference/v2.0/post
        facebook_posts.append({
            'id': post['id'],
            'message': post['message'],
            'link': {
                'url': post['link'],
                'name': post['name'],
                'caption': (post['caption'] if 'caption' in post else None),
                'description': post['description'],
                'picture': post['picture']
            },
            'from': {
                'name': user['name'],
                'link': user['link'],
                'picture': user_picture['url']
            },
            'likes': likes['summary']['total_count'],
            'comments': comments['summary']['total_count'],
            #'shares': shares['summary']['total_count'],
            'creation_date': creation_date
        })

    # Render to JSON
    output = {
        'tweets': tweets,
        'facebook_posts': facebook_posts
    }

    with open('data/featured.json', 'w') as f:
        json.dump(output, f)

def _mock_slide_from_image(filename, i):
    import models

    body = '' % (app_config.S3_BASE_URL, filename)
    name = 'Test slide: %s' % filename[0:-4]
    slide = models.Slide.create(body=body, name=name)
    models.SlideSequence.create(sequence=i, slide=slide)

@task
def mock_slides():
    """
    Load mockup slides from assets directory.
    """
    import models

    models.SlideSequence.delete().execute()
    models.Slide.delete().execute()

    it = count()
    _mock_slide_from_image('welcome.png', it.next())
    _mock_slide_from_image('senate.png', it.next())
    _mock_slide_from_image('gif1.gif', it.next())
    _mock_slide_from_image('party_pix.png', it.next())

@task
def mock_election_results():
    """
    Fake out some election results
    """
    import models

    for race in models.Race.select():
        _fake_incumbent(race)
        _fake_poll_closing_time(race)
        _fake_precincts_reporting(race)
        _fake_called_status(race)
        _fake_results(race)
        race.save()

def _fake_incumbent(race):
    """
    Fake one incumbent for a race
    """
    import models
    candidates = race.candidates.where((models.Candidate.party == "GOP") | (models.Candidate.party == "Dem"))
    try:
        candidate = candidates[random.randint(0,1)]
        candidate.incumbent = True
        candidate.save()
    except IndexError:
        pass

def _fake_poll_closing_time(race):
    """
    Fake poll closing time
    """
    first_close = datetime(2014, 11, 4, 7)
    closing_times = [first_close + timedelta(hours=delta) for delta in range(6)]
    race.poll_closing_time = random.choice(closing_times)

def _fake_precincts_reporting(race):
    """
    Fake precincts reporting
    """
    race.precincts_total = random.randint(2000, 4000)
    if random.choice([True, False, False]):
        race.precincts_reporting = random.randint(200, race.precincts_total)
    else:
        race.precincts_reporting = 0

def _fake_called_status(race):
    """
    Fake AP called status, requires race to have closing time
    """
    if race.precincts_reporting > 0:
        race.ap_called = random.choice([True, False, False, False])
        if race.ap_called:
            race.accept_ap_call = True
            race.ap_called_time = race.poll_closing_time + timedelta(hours=random.randint(1,3), minutes=random.randint(0,59))

def _fake_results(race):
    max_votes = 0
    for candidate in race.candidates:
        if candidate.party in ["GOP", "Dem"] and race.precincts_reporting > 0:
            votes = random.randint(400000, 500000)
            candidate.vote_count = votes
            if votes > max_votes:
                max_candidate = candidate
                max_votes = votes
        else:
            candidate.vote_count = 0

        candidate.save()

    if race.precincts_reporting > 0 and race.ap_called:
        max_candidate.ap_winner = True
        max_candidate.save()
<|endoftext|>


<commit_before># Copyright (c) 2016 The Regents of the University of Michigan
# All rights reserved.
# This software is licensed under the BSD 3-Clause License.
import os
import errno
import logging
import json
import shutil
import copy

from ..common import six
from ..core.jsondict import JSonDict
from .hashing import calc_id
from .utility import _mkdir_p

logger = logging.getLogger(__name__)

class Job(object):
    """The job instance is a handle to the data of a unique statepoint.

    Application developers should usually not need to directly
    instantiate this class, but use :meth:`~.project.Project.open_job`
    instead."""
    FN_MANIFEST = 'signac_statepoint.json'
    """The job's manifest filename.

    The job manifest, this means a human-readable dump of the job's\
    statepoint is stored in each workspace directory.
    """
    FN_DOCUMENT = 'signac_job_document.json'
    "The job's document filename."

    def __init__(self, project, statepoint):
        self._project = project
        self._statepoint = json.loads(json.dumps(statepoint))
        self._id = calc_id(self._statepoint)
        self._document = None
        self._wd = os.path.join(project.workspace(), str(self))
        self._cwd = list()

    def get_id(self):
        """The unique identifier for the job's statepoint.

        :return: The job id.
        :rtype: str"""
        return self._id

    def __str__(self):
        "Returns the job's id."
        return str(self.get_id())

    def __repr__(self):
        return "{}(project={}, id='{}')".format(
            self.__class__.__module__ + '.' + self.__class__.__name__,
            repr(self._project), self._id)

    def __eq__(self, other):
        return repr(self) == repr(other)

    def workspace(self):
        """Each job is associated with a unique workspace directory.

        :return: The path to the job's workspace directory.
        :rtype: str"""
        return self._wd

    def statepoint(self):
        """The statepoint associated with this job.

        :return: The statepoint mapping.
        :rtype: dict"""
        return copy.deepcopy(self._statepoint)

    @property
    def document(self):
        """The document associated with this job.

        :return: The job document handle.
        :rtype: :class:`~.JSonDict`"""
        if self._document is None:
            self._create_directory()
            fn = os.path.join(self.workspace(), self.FN_DOCUMENT)
            self._document = JSonDict(
                fn, synchronized=True, write_concern=True)
        return self._document

    def _create_directory(self, overwrite=False):
        "Create the workspace directory and write the manifest file."
        fn_manifest = os.path.join(self.workspace(), self.FN_MANIFEST)

        # Create the workspace directory if it did not exist yet.
        _mkdir_p(self.workspace())

        try:
            # Ensure to create the binary to write before file creation
            blob = json.dumps(self.statepoint(), indent=2)

            try:
                # Open the file for writing only if it does not exist yet.
                mode = 'w' if overwrite else 'wx' if six.PY2 else 'x'
                with open(fn_manifest, mode) as file:
                    file.write(blob)
            except IOError as error:
                if not error.errno == errno.EEXIST:
                    raise
        except Exception as error:
            # Attempt to delete the file on error, to prevent corruption.
            try:
                os.remove(fn_manifest)
            except Exception:  # ignore all errors here
                pass
            raise error
        else:
            self._check_manifest()

    def _check_manifest(self):
        "Check whether the manifest file, if it exists, is correct."
        fn_manifest = os.path.join(self.workspace(), self.FN_MANIFEST)
        try:
            try:
                with open(fn_manifest) as file:
                    assert calc_id(json.load(file)) == self._id
            except IOError as error:
                if not error.errno == errno.ENOENT:
                    raise error
        except Exception as error:
            msg = "Manifest file of job '{}' is corrupted: {}."
            raise RuntimeError(msg.format(self, error))

    def init(self):
        """Initialize the job's workspace directory.

        This function will do nothing if the directory and
        the job manifest already exist."""
        self._create_directory()

    def remove(self):
        """Remove the job's workspace including the job document.

        This function will do nothing if the workspace directory
        does not exist."""
        if self._document is not None:
            self._document.clear()
        try:
            shutil.rmtree(self.workspace())
        except OSError as error:
            if error.errno != errno.ENOENT:
                raise

    def fn(self, filename):
        """Prepend a filename with the job's workspace directory path.

        :param filename: The filename of the file.
        :type filename: str
        :return: The full workspace path of the file."""
        return os.path.join(self.workspace(), filename)

    def isfile(self, filename):
        """Return True if file exists in the job's workspace.

        :param filename: The filename of the file.
        :type filename: str
        :return: True if file with filename exists in workspace.
        :rtype: bool"""
        return os.path.isfile(self.fn(filename))

    def open(self):
        """Enter the job's workspace directory.

        You can use the :class:`~.Job` class as context manager:

        .. code-block:: python

            with project.open_job(my_statepoint) as job:
                # manipulate your job data

        Opening the context will switch into the job's workspace,
        leaving it will switch back to the previous working directory.
        """
        self._cwd.append(os.getcwd())
        self._create_directory()
        logger.info("Enter workspace '{}'.".format(self.workspace()))
        os.chdir(self.workspace())

    def close(self):
        "Close the job and switch to the previous working directory."
        try:
            os.chdir(self._cwd.pop())
            logger.info("Leave workspace.")
        except IndexError:
            pass

    def __enter__(self):
        self.open()
        return self

    def __exit__(self, err_type, err_value, tb):
        self.close()
        return False
<commit_msg>Fix error when using a job instance after removal.<commit_after># Copyright (c) 2016 The Regents of the University of Michigan
# All rights reserved.
# This software is licensed under the BSD 3-Clause License.
import os
import errno
import logging
import json
import shutil
import copy

from ..common import six
from ..core.jsondict import JSonDict
from .hashing import calc_id
from .utility import _mkdir_p

logger = logging.getLogger(__name__)

class Job(object):
    """The job instance is a handle to the data of a unique statepoint.

    Application developers should usually not need to directly
    instantiate this class, but use :meth:`~.project.Project.open_job`
    instead."""
    FN_MANIFEST = 'signac_statepoint.json'
    """The job's manifest filename.

    The job manifest, this means a human-readable dump of the job's\
    statepoint is stored in each workspace directory.
    """
    FN_DOCUMENT = 'signac_job_document.json'
    "The job's document filename."

    def __init__(self, project, statepoint):
        self._project = project
        self._statepoint = json.loads(json.dumps(statepoint))
        self._id = calc_id(self._statepoint)
        self._document = None
        self._wd = os.path.join(project.workspace(), str(self))
        self._cwd = list()

    def get_id(self):
        """The unique identifier for the job's statepoint.

        :return: The job id.
        :rtype: str"""
        return self._id

    def __str__(self):
        "Returns the job's id."
        return str(self.get_id())

    def __repr__(self):
        return "{}(project={}, id='{}')".format(
            self.__class__.__module__ + '.' + self.__class__.__name__,
            repr(self._project), self._id)

    def __eq__(self, other):
        return repr(self) == repr(other)

    def workspace(self):
        """Each job is associated with a unique workspace directory.

        :return: The path to the job's workspace directory.
        :rtype: str"""
        return self._wd

    def statepoint(self):
        """The statepoint associated with this job.

        :return: The statepoint mapping.
        :rtype: dict"""
        return copy.deepcopy(self._statepoint)

    @property
    def document(self):
        """The document associated with this job.

        :return: The job document handle.
        :rtype: :class:`~.JSonDict`"""
        if self._document is None:
            self._create_directory()
            fn = os.path.join(self.workspace(), self.FN_DOCUMENT)
            self._document = JSonDict(
                fn, synchronized=True, write_concern=True)
        return self._document

    def _create_directory(self, overwrite=False):
        "Create the workspace directory and write the manifest file."
        fn_manifest = os.path.join(self.workspace(), self.FN_MANIFEST)

        # Create the workspace directory if it did not exist yet.
        _mkdir_p(self.workspace())

        try:
            # Ensure to create the binary to write before file creation
            blob = json.dumps(self.statepoint(), indent=2)

            try:
                # Open the file for writing only if it does not exist yet.
                mode = 'w' if overwrite else 'wx' if six.PY2 else 'x'
                with open(fn_manifest, mode) as file:
                    file.write(blob)
            except IOError as error:
                if not error.errno == errno.EEXIST:
                    raise
        except Exception as error:
            # Attempt to delete the file on error, to prevent corruption.
            try:
                os.remove(fn_manifest)
            except Exception:  # ignore all errors here
                pass
            raise error
        else:
            self._check_manifest()

    def _check_manifest(self):
        "Check whether the manifest file, if it exists, is correct."
        fn_manifest = os.path.join(self.workspace(), self.FN_MANIFEST)
        try:
            try:
                with open(fn_manifest) as file:
                    assert calc_id(json.load(file)) == self._id
            except IOError as error:
                if not error.errno == errno.ENOENT:
                    raise error
        except Exception as error:
            msg = "Manifest file of job '{}' is corrupted: {}."
            raise RuntimeError(msg.format(self, error))

    def init(self):
        """Initialize the job's workspace directory.

        This function will do nothing if the directory and
        the job manifest already exist."""
        self._create_directory()

    def remove(self):
        """Remove the job's workspace including the job document.

        This function will do nothing if the workspace directory
        does not exist."""
        try:
            shutil.rmtree(self.workspace())
        except OSError as error:
            if error.errno != errno.ENOENT:
                raise
        else:
            if self._document is not None:
                self._document.clear()
                self._document = None

    def fn(self, filename):
        """Prepend a filename with the job's workspace directory path.

        :param filename: The filename of the file.
        :type filename: str
        :return: The full workspace path of the file."""
        return os.path.join(self.workspace(), filename)

    def isfile(self, filename):
        """Return True if file exists in the job's workspace.

        :param filename: The filename of the file.
        :type filename: str
        :return: True if file with filename exists in workspace.
        :rtype: bool"""
        return os.path.isfile(self.fn(filename))

    def open(self):
        """Enter the job's workspace directory.

        You can use the :class:`~.Job` class as context manager:

        .. code-block:: python

            with project.open_job(my_statepoint) as job:
                # manipulate your job data

        Opening the context will switch into the job's workspace,
        leaving it will switch back to the previous working directory.
        """
        self._cwd.append(os.getcwd())
        self._create_directory()
        logger.info("Enter workspace '{}'.".format(self.workspace()))
        os.chdir(self.workspace())

    def close(self):
        "Close the job and switch to the previous working directory."
        try:
            os.chdir(self._cwd.pop())
            logger.info("Leave workspace.")
        except IndexError:
            pass

    def __enter__(self):
        self.open()
        return self

    def __exit__(self, err_type, err_value, tb):
        self.close()
        return False
<|endoftext|>