<commit_before>from __future__ import division
import logging; _L = logging.getLogger('openaddr.preview')

from zipfile import ZipFile
from io import TextIOWrapper
from csv import DictReader
from tempfile import mkstemp
from math import pow, sqrt, pi, log
from argparse import ArgumentParser
from urllib.parse import urlparse
import json, itertools, os

import requests, uritemplate

from osgeo import osr, ogr
from .ci import setup_logger
from .compat import cairo

TILE_URL = 'http://tile.mapzen.com/mapzen/vector/v1/all/{z}/{x}/{y}.json{?api_key}'
EARTH_DIAMETER = 6378137 * 2 * pi

# WGS 84, http://spatialreference.org/ref/epsg/4326/
EPSG4326 = '+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs'

# Web Mercator, https://trac.osgeo.org/openlayers/wiki/SphericalMercator
EPSG900913 = '+proj=merc +a=6378137 +b=6378137 +lat_ts=0.0 +lon_0=0.0 +x_0=0.0 +y_0=0 +k=1.0 +units=m +nadgrids=@null +no_defs'

def render(zip_filename_or_url, png_filename, width, resolution, mapzen_key):
    '''
    '''
    zip_filename = get_local_zip_filename(zip_filename_or_url)
    points = project_points(iterate_zipfile_points(zip_filename))
    xmin, ymin, xmax, ymax = calculate_bounds(points)

    surface, context, scale = make_context(xmin, ymin, xmax, ymax, width, resolution)

    _L.info('Preview width {:.0f}, scale {:.3f}, zoom {:.2f}'.format(width, scale, calculate_zoom(scale, resolution)))

    # Map units per reference pixel (http://www.w3.org/TR/css3-values/#reference-pixel)
    muppx = resolution / scale

    black = 0, 0, 0
    light_green = 0x74/0xff, 0xA5/0xff, 0x78/0xff

    context.set_source_rgb(1, 1, 1)
    context.rectangle(xmin, ymax, xmax - xmin, ymin - ymax)
    context.fill()

    water_geoms, roads_geoms = \
        get_map_features(xmin, ymin, xmax, ymax, resolution, scale, mapzen_key)

    fill_geometries(context, water_geoms, muppx, (0xdd/0xff, 0xea/0xff, 0xf8/0xff))

    context.set_line_width(.25 * muppx)
    context.set_source_rgb(0xe0/0xff, 0xe3/0xff, 0xe5/0xff)
    stroke_geometries(context, roads_geoms)

    context.set_line_width(.25 * muppx)

    for (x, y) in points:
        context.arc(x, y, 15, 0, 2 * pi)
        context.set_source_rgb(*light_green)
        context.fill()
        context.arc(x, y, 15, 0, 2 * pi)
        context.set_source_rgb(*black)
        context.stroke()

    surface.write_to_png(png_filename)

def get_local_zip_filename(zip_filename_or_url):
    '''
    '''
    if urlparse(zip_filename_or_url).scheme in ('', 'file'):
        return zip_filename_or_url

    if urlparse(zip_filename_or_url).scheme not in ('http', 'https'):
        raise ValueError('Unknown URL type: {}'.format(zip_filename_or_url))

    _L.info('Downloading {}...'.format(zip_filename_or_url))

    got = requests.get(zip_filename_or_url)
    _, zip_filename = mkstemp(prefix='Preview-', suffix='.zip')

    with open(zip_filename, 'wb') as file:
        file.write(got.content)
        _L.debug('Saved to {}'.format(zip_filename))

    return zip_filename

def iterate_zipfile_points(filename):
    '''
    '''
    with open(filename, 'rb') as file:
        zip = ZipFile(file)
        csv_names = [name for name in zip.namelist() if name.endswith('.csv')]
        csv_file = TextIOWrapper(zip.open(csv_names[0]))

        for row in DictReader(csv_file):
            try:
                lon, lat = float(row['LON']), float(row['LAT'])
            except:
                continue

            if -180 <= lon <= 180 and -90 <= lat <= 90:
                yield (lon, lat)

def get_map_features(xmin, ymin, xmax, ymax, resolution, scale, mapzen_key):
    '''
    '''
    zoom = round(calculate_zoom(scale, resolution))
    mincol = 2**zoom * (xmin + EARTH_DIAMETER/2) / EARTH_DIAMETER
    minrow = 2**zoom * (EARTH_DIAMETER/2 - ymax) / EARTH_DIAMETER
    maxcol = 2**zoom * (xmax + EARTH_DIAMETER/2) / EARTH_DIAMETER
    maxrow = 2**zoom * (EARTH_DIAMETER/2 - ymin) / EARTH_DIAMETER

    row_cols = itertools.product(range(int(minrow), int(maxrow) + 1),
                                 range(int(mincol), int(maxcol) + 1))

    water_geoms, roads_geoms, project = list(), list(), get_projection()

    def projected_geom(feature):
        geom = ogr.CreateGeometryFromJson(json.dumps(feature['geometry']))
        geom.Transform(project)
        return geom

    for (row, col) in row_cols:
        url = uritemplate.expand(TILE_URL, dict(z=zoom, x=col, y=row, api_key=mapzen_key))
        got = requests.get(url)

        for feature in got.json()['water']['features']:
            if 'Polygon' in feature['geometry']['type']:
                if feature['properties']['kind'] in ('basin', 'lake', 'ocean', 'riverbank', 'water'):
                    water_geoms.append(projected_geom(feature))

        for feature in got.json()['roads']['features']:
            if 'LineString' in feature['geometry']['type']:
                if feature['properties']['kind'] in ('highway', 'major_road', 'minor_road', 'rail', 'path'):
                    roads_geoms.append(projected_geom(feature))

        _L.debug('Getting tile {}'.format(url))

    return water_geoms, roads_geoms

def get_projection():
    '''
    '''
    osr.UseExceptions()
    sref_geo = osr.SpatialReference(); sref_geo.ImportFromProj4(EPSG4326)
    sref_map = osr.SpatialReference(); sref_map.ImportFromProj4(EPSG900913)
    return osr.CoordinateTransformation(sref_geo, sref_map)

def project_points(lonlats):
    '''
    '''
    project = get_projection()
    points = list()

    for (lon, lat) in lonlats:
        geom = ogr.CreateGeometryFromWkt('POINT({:.7f} {:.7f})'.format(lon, lat))
        geom.Transform(project)
        points.append((geom.GetX(), geom.GetY()))

    return points

def stats(values):
    '''
    '''
    mean = sum(values) / len(values)
    deviations = [pow(val - mean, 2) for val in values]
    stddev = sqrt(sum(deviations) / len(values))

    return mean, stddev

def calculate_zoom(scale, resolution):
    ''' Calculate web map zoom based on scale.
    '''
    scale_at_zero = resolution * 256 / EARTH_DIAMETER
    zoom = log(scale / scale_at_zero) / log(2)

    return zoom

def calculate_bounds(points):
    '''
    '''
    xs, ys = zip(*points)

    # use standard deviation to avoid far-flung mistakes
    (xmean, xsdev), (ymean, ysdev) = stats(xs), stats(ys)
    xmin, xmax = xmean - 5 * xsdev, xmean + 5 * xsdev
    ymin, ymax = ymean - 5 * ysdev, ymean + 5 * ysdev

    # look at the actual points
    okay_xs = [x for (x, y) in points if (xmin <= x <= xmax)]
    okay_ys = [y for (x, y) in points if (ymin <= y <= ymax)]
    left, bottom = min(okay_xs), min(okay_ys)
    right, top = max(okay_xs), max(okay_ys)

    # pad by 2% on all sides
    width, height = right - left, top - bottom
    left -= width / 50
    bottom -= height / 50
    right += width / 50
    top += height / 50

    return left, bottom, right, top

def make_context(left, bottom, right, top, width=668, resolution=1):
    ''' Get Cairo surface, context, and drawing scale.

        668px is the width of a comment box in Github, one place where
        these previews are designed to be used.
    '''
    aspect = (right - left) / (top - bottom)

    hsize = int(resolution * width)
    vsize = int(hsize / aspect)

    hscale = hsize / (right - left)
    vscale = (hsize / aspect) / (bottom - top)

    hoffset = -left
    voffset = -top

    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, hsize, vsize)
    context = cairo.Context(surface)
    context.scale(hscale, vscale)
    context.translate(hoffset, voffset)

    return surface, context, hscale

def stroke_geometries(ctx, geometries):
    '''
    '''
    for geometry in geometries:
        if geometry.GetGeometryType() in (ogr.wkbMultiPolygon, ogr.wkbMultiLineString):
            parts = geometry
        elif geometry.GetGeometryType() in (ogr.wkbPolygon, ogr.wkbLineString):
            parts = [geometry]
        else:
            continue

        for part in parts:
            if part.GetGeometryType() is ogr.wkbPolygon:
                rings = part
            else:
                rings = [part]

            for ring in rings:
                points = ring.GetPoints()
                if geometry.GetGeometryType() in (ogr.wkbPolygon, ogr.wkbMultiPolygon):
                    draw_line(ctx, points[-1], points)
                else:
                    draw_line(ctx, points[0], points[1:])
                ctx.stroke()

def fill_geometries(ctx, geometries, muppx, rgb):
    '''
    '''
    ctx.set_source_rgb(*rgb)

    for geometry in geometries:
        if geometry.GetGeometryType() == ogr.wkbMultiPolygon:
            parts = geometry
        elif geometry.GetGeometryType() == ogr.wkbPolygon:
            parts = [geometry]
        elif geometry.GetGeometryType() == ogr.wkbPoint:
            buffer = geometry.Buffer(2 * muppx, 3)
            parts = [buffer]
        else:
            raise NotImplementedError()

        for part in parts:
            for ring in part:
                points = ring.GetPoints()
                draw_line(ctx, points[-1], points)
            ctx.fill()

def draw_line(ctx, start, points):
    '''
    '''
    ctx.move_to(*start)

    for point in points:
        ctx.line_to(*point)

parser = ArgumentParser(description='Draw a map of a single source preview.')

parser.add_argument('zip_filename', help='Input Zip filename or URL.')
parser.add_argument('png_filename', help='Output PNG filename.')

parser.set_defaults(resolution=1, width=668)

parser.add_argument('--2x', dest='resolution', action='store_const', const=2,
                    help='Draw at double resolution.')

parser.add_argument('--1x', dest='resolution', action='store_const', const=1,
                    help='Draw at normal resolution.')

parser.add_argument('--width', dest='width', type=int,
                    help='Width in pixels.')

parser.add_argument('--mapzen-key', dest='mapzen_key',
                    help='Mapze API Key. See: https://mapzen.com/documentation/overview/')

parser.add_argument('-v', '--verbose', help='Turn on verbose logging',
                    action='store_const', dest='loglevel',
                    const=logging.DEBUG, default=logging.INFO)

parser.add_argument('-q', '--quiet', help='Turn off most logging',
                    action='store_const', dest='loglevel',
                    const=logging.WARNING, default=logging.INFO)

def main():
    args = parser.parse_args()
    setup_logger(None, None, None, log_level=args.loglevel)
    render(args.zip_filename, args.png_filename, args.width, args.resolution, args.mapzen_key)

if __name__ == '__main__':
    exit(main())
<commit_msg>Added green landuse polygons to basemap render<commit_after>from __future__ import division
import logging; _L = logging.getLogger('openaddr.preview')

from zipfile import ZipFile
from io import TextIOWrapper
from csv import DictReader
from tempfile import mkstemp
from math import pow, sqrt, pi, log
from argparse import ArgumentParser
from urllib.parse import urlparse
import json, itertools, os

import requests, uritemplate

from osgeo import osr, ogr
from .ci import setup_logger
from .compat import cairo

TILE_URL = 'http://tile.mapzen.com/mapzen/vector/v1/all/{z}/{x}/{y}.json{?api_key}'
EARTH_DIAMETER = 6378137 * 2 * pi

# WGS 84, http://spatialreference.org/ref/epsg/4326/
EPSG4326 = '+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs'

# Web Mercator, https://trac.osgeo.org/openlayers/wiki/SphericalMercator
EPSG900913 = '+proj=merc +a=6378137 +b=6378137 +lat_ts=0.0 +lon_0=0.0 +x_0=0.0 +y_0=0 +k=1.0 +units=m +nadgrids=@null +no_defs'

def render(zip_filename_or_url, png_filename, width, resolution, mapzen_key):
    '''
    '''
    zip_filename = get_local_zip_filename(zip_filename_or_url)
    points = project_points(iterate_zipfile_points(zip_filename))
    xmin, ymin, xmax, ymax = calculate_bounds(points)

    surface, context, scale = make_context(xmin, ymin, xmax, ymax, width, resolution)

    _L.info('Preview width {:.0f}, scale {:.3f}, zoom {:.2f}'.format(width, scale, calculate_zoom(scale, resolution)))

    # Map units per reference pixel (http://www.w3.org/TR/css3-values/#reference-pixel)
    muppx = resolution / scale

    black = 0, 0, 0
    point_fill = 0x74/0xFF, 0xA5/0xFF, 0x78/0xFF
    water_fill = 0xC7/0xFF, 0xDE/0xFF, 0xF5/0xFF # 0xDD/0xFF, 0xEA/0xFF, 0xF8/0xFF
    road_stroke = 0xC0/0xFF, 0xE0/0xFF, 0xE0/0xFF # 0xE0/0xFF, 0xE3/0xFF, 0xE5/0xFF
    park_fill = 0xDD/0xFF, 0xF6/0xFF, 0xDE/0xFF

    context.set_source_rgb(1, 1, 1)
    context.rectangle(xmin, ymax, xmax - xmin, ymin - ymax)
    context.fill()

    landuse_geoms, water_geoms, roads_geoms = \
        get_map_features(xmin, ymin, xmax, ymax, resolution, scale, mapzen_key)

    fill_geometries(context, landuse_geoms, muppx, park_fill)
    fill_geometries(context, water_geoms, muppx, water_fill)

    context.set_line_width(.5 * muppx)
    context.set_source_rgb(*road_stroke)
    stroke_geometries(context, roads_geoms)

    context.set_line_width(.25 * muppx)

    for (x, y) in points:
        context.arc(x, y, 15, 0, 2 * pi)
        context.set_source_rgb(*point_fill)
        context.fill()
        context.arc(x, y, 15, 0, 2 * pi)
        context.set_source_rgb(*black)
        context.stroke()

    surface.write_to_png(png_filename)

def get_local_zip_filename(zip_filename_or_url):
    '''
    '''
    if urlparse(zip_filename_or_url).scheme in ('', 'file'):
        return zip_filename_or_url

    if urlparse(zip_filename_or_url).scheme not in ('http', 'https'):
        raise ValueError('Unknown URL type: {}'.format(zip_filename_or_url))

    _L.info('Downloading {}...'.format(zip_filename_or_url))

    got = requests.get(zip_filename_or_url)
    _, zip_filename = mkstemp(prefix='Preview-', suffix='.zip')

    with open(zip_filename, 'wb') as file:
        file.write(got.content)
        _L.debug('Saved to {}'.format(zip_filename))

    return zip_filename

def iterate_zipfile_points(filename):
    '''
    '''
    with open(filename, 'rb') as file:
        zip = ZipFile(file)
        csv_names = [name for name in zip.namelist() if name.endswith('.csv')]
        csv_file = TextIOWrapper(zip.open(csv_names[0]))

        for row in DictReader(csv_file):
            try:
                lon, lat = float(row['LON']), float(row['LAT'])
            except:
                continue

            if -180 <= lon <= 180 and -90 <= lat <= 90:
                yield (lon, lat)

def get_map_features(xmin, ymin, xmax, ymax, resolution, scale, mapzen_key):
    '''
    '''
    zoom = round(calculate_zoom(scale, resolution))
    mincol = 2**zoom * (xmin + EARTH_DIAMETER/2) / EARTH_DIAMETER
    minrow = 2**zoom * (EARTH_DIAMETER/2 - ymax) / EARTH_DIAMETER
    maxcol = 2**zoom * (xmax + EARTH_DIAMETER/2) / EARTH_DIAMETER
    maxrow = 2**zoom * (EARTH_DIAMETER/2 - ymin) / EARTH_DIAMETER

    row_cols = itertools.product(range(int(minrow), int(maxrow) + 1),
                                 range(int(mincol), int(maxcol) + 1))

    landuse_geoms, water_geoms, roads_geoms = list(), list(), list()
    project = get_projection()

    def projected_geom(feature):
        geom = ogr.CreateGeometryFromJson(json.dumps(feature['geometry']))
        geom.Transform(project)
        return geom

    for (row, col) in row_cols:
        url = uritemplate.expand(TILE_URL, dict(z=zoom, x=col, y=row, api_key=mapzen_key))
        got = requests.get(url)

        for feature in got.json()['landuse']['features']:
            if 'Polygon' in feature['geometry']['type']:
                if feature['properties'].get('kind') in ('cemetery', 'forest', 'golf_course', 'grave_yard', 'meadow', 'park', 'pitch', 'wood'):
                    landuse_geoms.append(projected_geom(feature))

        for feature in got.json()['water']['features']:
            if 'Polygon' in feature['geometry']['type']:
                if feature['properties']['kind'] in ('basin', 'lake', 'ocean', 'riverbank', 'water'):
                    water_geoms.append(projected_geom(feature))

        for feature in got.json()['roads']['features']:
            if 'LineString' in feature['geometry']['type']:
                if feature['properties']['kind'] in ('highway', 'major_road', 'minor_road', 'rail', 'path'):
                    roads_geoms.append(projected_geom(feature))

        _L.debug('Getting tile {}'.format(url))

    return landuse_geoms, water_geoms, roads_geoms

def get_projection():
    '''
    '''
    osr.UseExceptions()
    sref_geo = osr.SpatialReference(); sref_geo.ImportFromProj4(EPSG4326)
    sref_map = osr.SpatialReference(); sref_map.ImportFromProj4(EPSG900913)
    return osr.CoordinateTransformation(sref_geo, sref_map)

def project_points(lonlats):
    '''
    '''
    project = get_projection()
    points = list()

    for (lon, lat) in lonlats:
        geom = ogr.CreateGeometryFromWkt('POINT({:.7f} {:.7f})'.format(lon, lat))
        geom.Transform(project)
        points.append((geom.GetX(), geom.GetY()))

    return points

def stats(values):
    '''
    '''
    mean = sum(values) / len(values)
    deviations = [pow(val - mean, 2) for val in values]
    stddev = sqrt(sum(deviations) / len(values))

    return mean, stddev

def calculate_zoom(scale, resolution):
    ''' Calculate web map zoom based on scale.
    '''
    scale_at_zero = resolution * 256 / EARTH_DIAMETER
    zoom = log(scale / scale_at_zero) / log(2)

    return zoom

def calculate_bounds(points):
    '''
    '''
    xs, ys = zip(*points)

    # use standard deviation to avoid far-flung mistakes
    (xmean, xsdev), (ymean, ysdev) = stats(xs), stats(ys)
    xmin, xmax = xmean - 5 * xsdev, xmean + 5 * xsdev
    ymin, ymax = ymean - 5 * ysdev, ymean + 5 * ysdev

    # look at the actual points
    okay_xs = [x for (x, y) in points if (xmin <= x <= xmax)]
    okay_ys = [y for (x, y) in points if (ymin <= y <= ymax)]
    left, bottom = min(okay_xs), min(okay_ys)
    right, top = max(okay_xs), max(okay_ys)

    # pad by 2% on all sides
    width, height = right - left, top - bottom
    left -= width / 50
    bottom -= height / 50
    right += width / 50
    top += height / 50

    return left, bottom, right, top

def make_context(left, bottom, right, top, width=668, resolution=1):
    ''' Get Cairo surface, context, and drawing scale.

        668px is the width of a comment box in Github, one place where
        these previews are designed to be used.
    '''
    aspect = (right - left) / (top - bottom)

    hsize = int(resolution * width)
    vsize = int(hsize / aspect)

    hscale = hsize / (right - left)
    vscale = (hsize / aspect) / (bottom - top)

    hoffset = -left
    voffset = -top

    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, hsize, vsize)
    context = cairo.Context(surface)
    context.scale(hscale, vscale)
    context.translate(hoffset, voffset)

    return surface, context, hscale

def stroke_geometries(ctx, geometries):
    '''
    '''
    for geometry in geometries:
        if geometry.GetGeometryType() in (ogr.wkbMultiPolygon, ogr.wkbMultiLineString):
            parts = geometry
        elif geometry.GetGeometryType() in (ogr.wkbPolygon, ogr.wkbLineString):
            parts = [geometry]
        else:
            continue

        for part in parts:
            if part.GetGeometryType() is ogr.wkbPolygon:
                rings = part
            else:
                rings = [part]

            for ring in rings:
                points = ring.GetPoints()
                if geometry.GetGeometryType() in (ogr.wkbPolygon, ogr.wkbMultiPolygon):
                    draw_line(ctx, points[-1], points)
                else:
                    draw_line(ctx, points[0], points[1:])
                ctx.stroke()

def fill_geometries(ctx, geometries, muppx, rgb):
    '''
    '''
    ctx.set_source_rgb(*rgb)

    for geometry in geometries:
        if geometry.GetGeometryType() == ogr.wkbMultiPolygon:
            parts = geometry
        elif geometry.GetGeometryType() == ogr.wkbPolygon:
            parts = [geometry]
        elif geometry.GetGeometryType() == ogr.wkbPoint:
            buffer = geometry.Buffer(2 * muppx, 3)
            parts = [buffer]
        else:
            raise NotImplementedError()

        for part in parts:
            for ring in part:
                points = ring.GetPoints()
                draw_line(ctx, points[-1], points)
            ctx.fill()

def draw_line(ctx, start, points):
    '''
    '''
    ctx.move_to(*start)

    for point in points:
        ctx.line_to(*point)

parser = ArgumentParser(description='Draw a map of a single source preview.')

parser.add_argument('zip_filename', help='Input Zip filename or URL.')
parser.add_argument('png_filename', help='Output PNG filename.')

parser.set_defaults(resolution=1, width=668)

parser.add_argument('--2x', dest='resolution', action='store_const', const=2,
                    help='Draw at double resolution.')

parser.add_argument('--1x', dest='resolution', action='store_const', const=1,
                    help='Draw at normal resolution.')

parser.add_argument('--width', dest='width', type=int,
                    help='Width in pixels.')

parser.add_argument('--mapzen-key', dest='mapzen_key',
                    help='Mapze API Key. See: https://mapzen.com/documentation/overview/')

parser.add_argument('-v', '--verbose', help='Turn on verbose logging',
                    action='store_const', dest='loglevel',
                    const=logging.DEBUG, default=logging.INFO)

parser.add_argument('-q', '--quiet', help='Turn off most logging',
                    action='store_const', dest='loglevel',
                    const=logging.WARNING, default=logging.INFO)

def main():
    args = parser.parse_args()
    setup_logger(None, None, None, log_level=args.loglevel)
    render(args.zip_filename, args.png_filename, args.width, args.resolution, args.mapzen_key)

if __name__ == '__main__':
    exit(main())
<|endoftext|>


<commit_before>from pyqtgraph.Qt import QtGui, QtCore
import pyqtgraph.functions as fn

__all__ = ['ArrowItem']
class ArrowItem(QtGui.QGraphicsPolygonItem):
    """
    For displaying scale-invariant arrows.
    For arrows pointing to a location on a curve, see CurveArrow

    """

    def __init__(self, **opts):
        QtGui.QGraphicsPolygonItem.__init__(self, opts.get('parent', None))
        defOpts = {
            'style': 'tri',
            'pxMode': True,
            'size': 20,
            'angle': -150,   ## If the angle is 0, the arrow points left
            'pos': (0,0),
            'width': None,  ## width is automatically size / 2.
            'tipAngle': 25,
            'baseAngle': 90,
            'pen': (200,200,200),
            'brush': (50,50,200),
        }
        defOpts.update(opts)

        self.setStyle(**defOpts)

        self.setPen(fn.mkPen(defOpts['pen']))
        self.setBrush(fn.mkBrush(defOpts['brush']))

        self.rotate(self.opts['angle'])
        self.moveBy(*self.opts['pos'])

    def setStyle(self, **opts):
        self.opts = opts

        if opts['style'] == 'tri':
            if opts['width'] is None:
                width = opts['size'] / 2.
            else:
                width = opts['width']

            points = [
                QtCore.QPointF(0,0),
                QtCore.QPointF(opts['size'],-width/2.),
                QtCore.QPointF(opts['size'],width/2.),
            ]
            poly = QtGui.QPolygonF(points)

        else:
            raise Exception("Unrecognized arrow style '%s'" % opts['style'])

        self.setPolygon(poly)

        if opts['pxMode']:
            self.setFlags(self.flags() | self.ItemIgnoresTransformations)
        else:
            self.setFlags(self.flags() & ~self.ItemIgnoresTransformations)

    def paint(self, p, *args):
        p.setRenderHint(QtGui.QPainter.Antialiasing)
        QtGui.QGraphicsPolygonItem.paint(self, p, *args)
<commit_msg>ArrowItem has more flexibility in the shapes it can draw<commit_after>from pyqtgraph.Qt import QtGui, QtCore
import pyqtgraph.functions as fn
import numpy as np
__all__ = ['ArrowItem']

class ArrowItem(QtGui.QGraphicsPathItem):
    """
    For displaying scale-invariant arrows.
    For arrows pointing to a location on a curve, see CurveArrow

    """

    def __init__(self, **opts):
        QtGui.QGraphicsPathItem.__init__(self, opts.get('parent', None))
        defOpts = {
            'pxMode': True,
            'angle': -150,   ## If the angle is 0, the arrow points left
            'pos': (0,0),
            'headLen': 20,
            'tipAngle': 25,
            'baseAngle': 0,
            'tailLen': None,
            'tailWidth': 3,
            'pen': (200,200,200),
            'brush': (50,50,200),
        }
        defOpts.update(opts)

        self.setStyle(**defOpts)

        self.setPen(fn.mkPen(defOpts['pen']))
        self.setBrush(fn.mkBrush(defOpts['brush']))

        self.rotate(self.opts['angle'])
        self.moveBy(*self.opts['pos'])

    def setStyle(self, **opts):
        self.opts = opts

        opt = {k:self.opts[k] for k in ['headLen', 'tipAngle', 'baseAngle', 'tailLen', 'tailWidth']}
        self.path = fn.makeArrowPath(**opt)
        self.setPath(self.path)

        if opts['pxMode']:
            self.setFlags(self.flags() | self.ItemIgnoresTransformations)
        else:
            self.setFlags(self.flags() & ~self.ItemIgnoresTransformations)

    def paint(self, p, *args):
        p.setRenderHint(QtGui.QPainter.Antialiasing)
        QtGui.QGraphicsPathItem.paint(self, p, *args)
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
# This file is part of PYBOSSA.
#
# Copyright (C) 2017 Scifabric LTD.
#
# PYBOSSA is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# PYBOSSA is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with PYBOSSA.  If not, see <http://www.gnu.org/licenses/>.
import json
from default import db, with_context
from nose.tools import assert_equal
from test_api import TestAPI

from factories import TaskFactory, UserFactory

from pybossa.repositories import TaskRepository

task_repo = TaskRepository(db)

class TestFavoritesAPI(TestAPI):

    url = '/api/favorites'

    @with_context
    def test_query_favorites_anon(self):
        """Test API Favorites works for anon."""
        user = UserFactory.create()
        TaskFactory.create(fav_user_ids=[user.id])
        res = self.app.get(self.url)
        data = json.loads(res.data)
        assert res.status_code == 401
        assert data['status_code'] == 401

    @with_context
    def test_query_get_favorites_auth(self):
        """Test API GET Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        tasks = TaskFactory.create_batch(30, fav_user_ids=[user.id])
        task = tasks[0]
        TaskFactory.create(fav_user_ids=[user2.id])
        res = self.app.get(self.url + '?api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 20, data
        data = data[0]
        assert data['id'] == task.id, (data, task)
        assert data['fav_user_ids'] == [user.id], data
        assert len(data['fav_user_ids']) == 1, data

        # limit
        res = self.app.get(self.url + '?limit=1&api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['id'] == task.id, (data, task)
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data

        # limit & offset
        res = self.app.get(self.url + '?limit=1&offset=1&api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data
        assert data[0]['id'] == tasks[1].id

        # last_id
        res = self.app.get(self.url + '?limit=1&last_id=%s&api_key=%s' %
                           (tasks[2].id, user.api_key))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data
        assert data[0]['id'] == tasks[3].id, data[0]['id']

        # desc == true
        res = self.app.get(self.url + '?limit=1&desc=1&api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['id'] == tasks[-1].id, data[0]['id']
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data

    @with_context
    def test_query_put_favorites_auth(self):
        """Test API PUT Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        TaskFactory.create(fav_user_ids=[user.id])
        TaskFactory.create(fav_user_ids=[user2.id])
        res = self.app.put(self.url + '/1?api_key=%s' % user.id)
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        res = self.app.put(self.url + '?api_key=%s' % user.id)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_query_put_favorites_anon(self):
        """Test API PUT Favorites works for anon."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        TaskFactory.create(fav_user_ids=[user.id])
        TaskFactory.create(fav_user_ids=[user2.id])
        res = self.app.put(self.url + '/1')
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        res = self.app.put(self.url)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_query_post_favorites_auth(self):
        """Test API POST Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        task = TaskFactory.create()
        url = self.url + '?api_key=%s' % user.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert data['fav_user_ids'] == [user.id], data

        url = self.url + '?api_key=%s' % user2.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id in data['fav_user_ids'], data
        assert user2.id in data['fav_user_ids'], data

        url = self.url + '?api_key=%s' % user2.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=4000000000)))
        data = json.loads(res.data)
        assert res.status_code == 404, res.status_code
        assert data['status_code'] == 404, data

        url = self.url + '?api_key=%s' % user2.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id in data['fav_user_ids'], data
        assert user2.id in data['fav_user_ids'], data

    @with_context
    def test_query_post_wrong_datafavorites_auth(self):
        """Test API POST Favorites wrong data for user."""
        user = UserFactory.create()
        task = TaskFactory.create()
        url = self.url + '?api_key=%s' % user.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id, id=3)))
        print res.data
        data = json.loads(res.data)
        assert res.status_code == 415, res.status_code
        assert data['status_code'] == 415, data

    @with_context
    def test_query_post_favorites_anon(self):
        """Test API POST Favorites works for anon."""
        res = self.app.post(self.url, data=json.dumps(dict(task_id=1)))
        data = json.loads(res.data)
        assert res.status_code == 401, res.status_code
        assert data['status_code'] == 401, res.status_code

    @with_context
    def test_query_delete_favorites_anon(self):
        """Test API DEL Favorites works for anon."""
        res = self.app.delete(self.url + '/1')
        data = json.loads(res.data)
        assert res.status_code == 401, res.status_code
        assert data['status_code'] == 401, res.status_code

        res = self.app.delete(self.url)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_query_delete_favorites_auth(self):
        """Test API DEL Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        task = TaskFactory.create(fav_user_ids=[user.id, user2.id])

        url = self.url + '/%s?api_key=%s' % (task.id, user.api_key)
        res = self.app.delete(url)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id not in data['fav_user_ids'], data
        assert user2.id in data['fav_user_ids'], data

        url = self.url + '/%s?api_key=%s' % (task.id, user2.api_key)
        res = self.app.delete(url)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id not in data['fav_user_ids'], data
        assert user2.id not in data['fav_user_ids'], data

        url = self.url + '/%s?api_key=%s' % (task.id, user2.api_key)
        res = self.app.delete(url)
        data = json.loads(res.data)
        assert res.status_code == 404, res.status_code
<commit_msg>Orderby.<commit_after># -*- coding: utf8 -*-
# This file is part of PYBOSSA.
#
# Copyright (C) 2017 Scifabric LTD.
#
# PYBOSSA is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# PYBOSSA is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with PYBOSSA.  If not, see <http://www.gnu.org/licenses/>.
import json
from default import db, with_context
from nose.tools import assert_equal
from test_api import TestAPI

from factories import TaskFactory, UserFactory

from pybossa.repositories import TaskRepository

task_repo = TaskRepository(db)

class TestFavoritesAPI(TestAPI):

    url = '/api/favorites'

    @with_context
    def test_query_favorites_anon(self):
        """Test API Favorites works for anon."""
        user = UserFactory.create()
        TaskFactory.create(fav_user_ids=[user.id])
        res = self.app.get(self.url)
        data = json.loads(res.data)
        assert res.status_code == 401
        assert data['status_code'] == 401

    @with_context
    def test_query_get_favorites_auth(self):
        """Test API GET Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        tasks = TaskFactory.create_batch(30, fav_user_ids=[user.id])
        task = tasks[0]
        TaskFactory.create(fav_user_ids=[user2.id])
        res = self.app.get(self.url + '?api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 20, data
        data = data[0]
        assert data['id'] == task.id, (data, task)
        assert data['fav_user_ids'] == [user.id], data
        assert len(data['fav_user_ids']) == 1, data

        # limit
        res = self.app.get(self.url + '?limit=1&api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['id'] == task.id, (data, task)
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data

        # limit & offset
        res = self.app.get(self.url + '?limit=1&offset=1&api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data
        assert data[0]['id'] == tasks[1].id

        # last_id
        res = self.app.get(self.url + '?limit=1&last_id=%s&api_key=%s' %
                           (tasks[2].id, user.api_key))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data
        assert data[0]['id'] == tasks[3].id, data[0]['id']

        # desc == true
        res = self.app.get(self.url + '?limit=1&desc=1&api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['id'] == tasks[-1].id, data[0]['id']
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data

        # orderby
        res = self.app.get(self.url + '?limit=1&orderby=created&api_key=%s' % user.api_key)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert len(data) == 1, len(data)
        assert data[0]['id'] == tasks[0].id, data[0]['id']
        assert data[0]['fav_user_ids'] == [user.id], data
        assert len(data[0]['fav_user_ids']) == 1, data

    @with_context
    def test_query_put_favorites_auth(self):
        """Test API PUT Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        TaskFactory.create(fav_user_ids=[user.id])
        TaskFactory.create(fav_user_ids=[user2.id])
        res = self.app.put(self.url + '/1?api_key=%s' % user.id)
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        res = self.app.put(self.url + '?api_key=%s' % user.id)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_query_put_favorites_anon(self):
        """Test API PUT Favorites works for anon."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        TaskFactory.create(fav_user_ids=[user.id])
        TaskFactory.create(fav_user_ids=[user2.id])
        res = self.app.put(self.url + '/1')
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        res = self.app.put(self.url)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_query_post_favorites_auth(self):
        """Test API POST Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        task = TaskFactory.create()
        url = self.url + '?api_key=%s' % user.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert data['fav_user_ids'] == [user.id], data

        url = self.url + '?api_key=%s' % user2.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id in data['fav_user_ids'], data
        assert user2.id in data['fav_user_ids'], data

        url = self.url + '?api_key=%s' % user2.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=4000000000)))
        data = json.loads(res.data)
        assert res.status_code == 404, res.status_code
        assert data['status_code'] == 404, data

        url = self.url + '?api_key=%s' % user2.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id in data['fav_user_ids'], data
        assert user2.id in data['fav_user_ids'], data

    @with_context
    def test_query_post_wrong_datafavorites_auth(self):
        """Test API POST Favorites wrong data for user."""
        user = UserFactory.create()
        task = TaskFactory.create()
        url = self.url + '?api_key=%s' % user.api_key
        res = self.app.post(url, data=json.dumps(dict(task_id=task.id, id=3)))
        print res.data
        data = json.loads(res.data)
        assert res.status_code == 415, res.status_code
        assert data['status_code'] == 415, data

    @with_context
    def test_query_post_favorites_anon(self):
        """Test API POST Favorites works for anon."""
        res = self.app.post(self.url, data=json.dumps(dict(task_id=1)))
        data = json.loads(res.data)
        assert res.status_code == 401, res.status_code
        assert data['status_code'] == 401, res.status_code

    @with_context
    def test_query_delete_favorites_anon(self):
        """Test API DEL Favorites works for anon."""
        res = self.app.delete(self.url + '/1')
        data = json.loads(res.data)
        assert res.status_code == 401, res.status_code
        assert data['status_code'] == 401, res.status_code

        res = self.app.delete(self.url)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_query_delete_favorites_auth(self):
        """Test API DEL Favorites works for user."""
        user = UserFactory.create()
        user2 = UserFactory.create()
        task = TaskFactory.create(fav_user_ids=[user.id, user2.id])

        url = self.url + '/%s?api_key=%s' % (task.id, user.api_key)
        res = self.app.delete(url)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id not in data['fav_user_ids'], data
        assert user2.id in data['fav_user_ids'], data

        url = self.url + '/%s?api_key=%s' % (task.id, user2.api_key)
        res = self.app.delete(url)
        data = json.loads(res.data)
        assert res.status_code == 200, res.status_code
        assert user.id not in data['fav_user_ids'], data
        assert user2.id not in data['fav_user_ids'], data

        url = self.url + '/%s?api_key=%s' % (task.id, user2.api_key)
        res = self.app.delete(url)
        data = json.loads(res.data)
        assert res.status_code == 404, res.status_code
<|endoftext|>


<commit_before>__author__ = 'sdc50'

from job import Job

#TODO: add node type (DATA, JOB)
class Node(object):
    """

    """

    def __init__(self, job,
                 parents=None,
                 children=None,
                 variables=None,
                 pre_script=None,
                 post_script=None,
                 retry=0):
        """

        :param job:
        :param parents:
        :param children:
        :param pre_script:
        :param post_script:
        :param retry:
        :return:
        """
        self._job = job
        self._parent_nodes = parents or set()
        self._link_parent_nodes()
        self._child_nodes = children or set()
        self._link_child_nodes()
        self._variables = variables or dict()
        self._pre_script = pre_script
        self._post_script = post_script
        self._retry = retry

    def __str__(self):
        """

        :return:
        """
        return 'JOB %s %s\n' % (self.job.name, self.job.job_file)

    def __repr__(self):
        """

        :return:
        """
        return '<NODE: Job:%s Parents:%s Children:%s>' % (self.job.name, self._get_parent_names(), self._get_child_names())

    @property
    def job(self):
        """

        :return:
        """
        return self._job

    @job.setter
    def job(self, job):
        """

        :param job:
        :return:
        """
        if isinstance(job, Job):
            self._job = job
        else:
            raise TypeError

    @property
    def pre_script(self):
        """

        :return:
        """
        return self._pre_script

    @pre_script.setter
    def pre_script(self, script):
        """

        :param script:
        :return:
        """
        self._pre_script = script

    @property
    def post_script(self):
        """

        :return:
        """
        return self._pre_script

    @post_script.setter
    def post_script(self, script):
        """

        :param script:
        :return:
        """
        self._post_script = script

    @property
    def parent_nodes(self):
        """

        :return:
        """
        return self._parent_nodes

    @property
    def child_nodes(self):
        """

        :return:
        """
        return self._child_nodes

    @property
    def retry(self):
        """

        :return:
        """
        return self._retry

    @retry.setter
    def retry(self, retry):
        """

        :return:
        """
        self._retry = retry

    def add_parent(self, parent):
        """

        :param parent:
        :return:
        """
        assert isinstance(parent, Node)
        self.parent_nodes.add(parent)
        if self not in parent.child_nodes:
            parent.add_child(self)

    def remove_parent(self, parent):
        """

        :param parent:
        :return:
        """
        assert isinstance(parent, Node)
        self.parent_nodes.discard(parent)
        if self in parent.child_nodes:
            parent.remove_child(self)

    def add_child(self, child):
        """

        :param child:
        :return:
        """
        assert isinstance(child, Node)
        self.child_nodes.add(child)
        if self not in child.parent_nodes:
            child.add_parent(self)

    def remove_child(self, child):
        """

        :param child:
        :return:
        """
        assert isinstance(child, Node)
        self.child_nodes.discard(child)
        if self in child.parent_nodes:
            child.remove_parent(self)

    def _add(self):
        """

        :return:
        """
        pass

    def _remove(self):
        """

        :return:
        """
        pass

    def get_all_family_nodes(self):
        """

        :return:
        """
        ancestors = self._get_all_ancestors()
        descendants = self._get_all_descendants()
        family_nodes = ancestors.union(descendants)
        return family_nodes

    def _get_all_ancestors(self):
        """

        :return:
        """
        ancestors = set()
        ancestors = ancestors.union(self.parent_nodes)
        for parent in self.parent_nodes:
            ancestors = ancestors.union(parent._get_all_ancestors())

        if self in ancestors:
            raise
        #TODO: make loop exception
        return ancestors

    def _get_all_descendants(self):
        """
        traverses all descendants nodes
        :raises: LoopException if self is contained in descendants
        :return: a set containing all descendant nodes
        """
        descendants = set()
        descendants = descendants.union(self.child_nodes)
        for child in self.child_nodes:
            descendants = descendants.union(child._get_all_descendants())

        if self in descendants:
            raise
            #TODO: make loop exception
        return descendants

    def _link_parent_nodes(self):
        """

        :return:
        """
        for parent in self.parent_nodes:
            if self not in parent.child_nodes:
                parent.add_child(self)

    def _link_child_nodes(self):
        """

        :return:
        """
        for child in self.child_nodes:
            if self not in child.parent_nodes:
                child.add_parent(self)

    def list_relations(self):
        """

        :return:
        """
        result = ''
        if len(self.child_nodes):
            result += 'PARENT %s CHILD %s\n' % (self.job.name, self._get_child_names())
        if self.retry:
            result += 'RETRY %s %d\n' % (self.job.name, self.retry)
        return result

    def _get_child_names(self):
        """

        :return:
        """
        return self._get_names(self.child_nodes)

    def _get_parent_names(self):
        """

        :return:
        """
        return self._get_names(self.parent_nodes)

    def _get_names(self, nodes):
        """

        :return:
        """
        names = ''
        for node in nodes:
            names += node.job.name + ' '
        return names<commit_msg>added pre and post script support in node.py<commit_after>__author__ = 'sdc50'

from job import Job

#TODO: add node type (DATA, JOB)
class Node(object):
    """

    """

    def __init__(self, job,
                 parents=None,
                 children=None,
                 variables=None,
                 pre_script=None,
                 pre_script_args=None,
                 post_script=None,
                 post_script_args=None,
                 retry=0):
        """

        :param job:
        :param parents:
        :param children:
        :param pre_script:
        :param post_script:
        :param retry:
        :return:
        """
        self._job = job
        self._parent_nodes = parents or set()
        self._link_parent_nodes()
        self._child_nodes = children or set()
        self._link_child_nodes()
        self._variables = variables or dict()
        self._pre_script = pre_script
        self._pre_script_args = pre_script_args
        self._post_script = post_script
        self._post_script_args = post_script_args
        self._retry = retry

    def __str__(self):
        """

        :return:
        """
        return 'JOB %s %s\n' % (self.job.name, self.job.job_file)

    def __repr__(self):
        """

        :return:
        """
        return '<NODE: Job:%s Parents:%s Children:%s>' % (self.job.name, self._get_parent_names(), self._get_child_names())

    @property
    def job(self):
        """

        :return:
        """
        return self._job

    @job.setter
    def job(self, job):
        """

        :param job:
        :return:
        """
        if isinstance(job, Job):
            self._job = job
        else:
            raise TypeError

    @property
    def pre_script(self):
        """

        :return:
        """
        return self._pre_script

    @pre_script.setter
    def pre_script(self, script):
        """

        :param script:
        :return:
        """
        self._pre_script = script

    @property
    def pre_script_args(self):
        return self._pre_script_args

    @pre_script_args.setter
    def pre_script_args(self, args):
        self._pre_script_args = args

    @property
    def post_script(self):
        """

        :return:
        """
        return self._pre_script

    @post_script.setter
    def post_script(self, script):
        """

        :param script:
        :return:
        """
        self._post_script = script

    @property
    def post_script_args(self):
        return self._post_script_args

    @post_script_args.setter
    def post_script_args(self, args):
        self._post_script_args = args

    @property
    def parent_nodes(self):
        """

        :return:
        """
        return self._parent_nodes

    @property
    def child_nodes(self):
        """

        :return:
        """
        return self._child_nodes

    @property
    def retry(self):
        """

        :return:
        """
        return self._retry

    @retry.setter
    def retry(self, retry):
        """

        :return:
        """
        self._retry = retry

    def add_parent(self, parent):
        """

        :param parent:
        :return:
        """
        assert isinstance(parent, Node)
        self.parent_nodes.add(parent)
        if self not in parent.child_nodes:
            parent.add_child(self)

    def remove_parent(self, parent):
        """

        :param parent:
        :return:
        """
        assert isinstance(parent, Node)
        self.parent_nodes.discard(parent)
        if self in parent.child_nodes:
            parent.remove_child(self)

    def add_child(self, child):
        """

        :param child:
        :return:
        """
        assert isinstance(child, Node)
        self.child_nodes.add(child)
        if self not in child.parent_nodes:
            child.add_parent(self)

    def remove_child(self, child):
        """

        :param child:
        :return:
        """
        assert isinstance(child, Node)
        self.child_nodes.discard(child)
        if self in child.parent_nodes:
            child.remove_parent(self)

    def _add(self):
        """

        :return:
        """
        pass

    def _remove(self):
        """

        :return:
        """
        pass

    def get_all_family_nodes(self):
        """

        :return:
        """
        ancestors = self._get_all_ancestors()
        descendants = self._get_all_descendants()
        family_nodes = ancestors.union(descendants)
        return family_nodes

    def _get_all_ancestors(self):
        """

        :return:
        """
        ancestors = set()
        ancestors = ancestors.union(self.parent_nodes)
        for parent in self.parent_nodes:
            ancestors = ancestors.union(parent._get_all_ancestors())

        if self in ancestors:
            raise
        #TODO: make loop exception
        return ancestors

    def _get_all_descendants(self):
        """
        traverses all descendants nodes
        :raises: LoopException if self is contained in descendants
        :return: a set containing all descendant nodes
        """
        descendants = set()
        descendants = descendants.union(self.child_nodes)
        for child in self.child_nodes:
            descendants = descendants.union(child._get_all_descendants())

        if self in descendants:
            raise
            #TODO: make loop exception
        return descendants

    def _link_parent_nodes(self):
        """

        :return:
        """
        for parent in self.parent_nodes:
            if self not in parent.child_nodes:
                parent.add_child(self)

    def _link_child_nodes(self):
        """

        :return:
        """
        for child in self.child_nodes:
            if self not in child.parent_nodes:
                child.add_parent(self)

    def list_relations(self):
        """

        :return:
        """
        result = ''
        if len(self.child_nodes):
            result += 'PARENT %s CHILD %s\n' % (self.job.name, self._get_child_names())
        if self.retry:
            result += 'RETRY %s %d\n' % (self.job.name, self.retry)
        if self.pre_script:
            result += 'SCRIPT PRE %s %s\n' % (self.job.name, self.pre_script, self.pre_script_args)
        if self.post_script:
            result += 'SCRIPT POST %s %s\n' % (self.job.name, self.post_script, self.post_script_args)
        return result

    def _get_child_names(self):
        """

        :return:
        """
        return self._get_names(self.child_nodes)

    def _get_parent_names(self):
        """

        :return:
        """
        return self._get_names(self.parent_nodes)

    def _get_names(self, nodes):
        """

        :return:
        """
        names = ''
        for node in nodes:
            names += node.job.name + ' '
        return names<|endoftext|>


<commit_before>
try:
    import wpilib
except ImportError:
    from pyfrc import wpilib

class MyRobot(wpilib.SimpleRobot):
    state = 1
    def __init__ (self):
        super().__init__()
<<<<<<< HEAD

        print("Tim Winters")

=======
        print("Tim Winters II")
>>>>>>> 932d5c04c69b854047f3ac275a4cc161f31c31ff
        #self.digitalInput=wpilib.DigitalInput(4)
        self.Joystick=wpilib.Joystick(1)
        self.Joystick2=wpilib.Joystick(2)
        self.Timer=wpilib.Timer()
        self.jaguar=wpilib.Jaguar(1)
        self.jaguar2=wpilib.Jaguar(2)
        self.jaguar3=wpilib.Jaguar(3)
        self.jaguar4=wpilib.Jaguar(4)
        self.drive = wpilib.RobotDrive(self.jaguar, self.jaguar2, self.jaguar3, self.jaguar4)#self.jaguar4=wpilib.Jaguar(4)

        self.drive.SetSafetyEnabled(False)

    challenge= 1

    def OperatorControl(self):
        print(self.IsEnabled())
       # dog = wpilib.GetWatchdog()
        #dog.setEnabled(True)
        #dog.SetExpiration(10)

        while self.IsOperatorControl()and self.IsEnabled():
          #  dog.Feed()
            self.drive.MecanumDrive_Cartesian(self.Joystick.GetY(), self.Joystick.GetX(), self.Joystick2.GetX(), 0)
            wpilib.Wait(0.01)

    def move_motor(self):
       # print(self.state, self.Timer.Get(), self.digitalInput.Get())
        '''if self.state is 1:
            self.Timer.Reset()
            self.jaguar.Set(self.Joystick.GetY())
            if self.digitalInput.Get()==True:
                self.Timer.Reset()
                self.state =  2

        elif self.state is 2:
            self.jaguar.Set(self.Joystick.GetY())
            if self.Timer.HasPeriodPassed(1):
                self.state = 3
            if not self.digitalInput.Get():
                self.state=1
        elif self.state is 3:
            self.jaguar.Set(1)
            if self.Timer.HasPeriodPassed(3):
                self.state = 4
        elif self.state is 4:
             self.jaguar.Set(-1)
             if(self.Timer.HasPeriodPassed(2)):
                 self.state=1
        '''

def run():

    robot = MyRobot()
    robot.StartCompetition()

    return robot

if __name__ == '__main__':
    wpilib.run()

<commit_msg>Added 5th and 6th Jags. May need to switch around the CAN jagaur.<commit_after>
try:
    import wpilib
except ImportError:
    from pyfrc import wpilib

class MyRobot(wpilib.SimpleRobot):
    state = 1
    def __init__ (self):
        super().__init__()

        print("Tim Winters II")
        #self.digitalInput=wpilib.DigitalInput(4)
        self.Joystick=wpilib.Joystick(1)
        self.Joystick2=wpilib.Joystick(2)
        self.Timer=wpilib.Timer()
        self.jaguar1=wpilib.Jaguar(1)
        self.jaguar2=wpilib.Jaguar(2)
        self.jaguar3=wpilib.Jaguar(3)
        self.jaguar4=wpilib.Jaguar(4)
        self.jaguar5=wpilib.CANJaguar(5)
        self.jaguar6=wpilib.Jaguar(6)
        self.drive = wpilib.RobotDrive(self.jaguar, self.jaguar2, self.jaguar3, self.jaguar4)#self.jaguar4=wpilib.Jaguar(4)

        self.drive.SetSafetyEnabled(False)

    challenge= 1

    def OperatorControl(self):
        print(self.IsEnabled())
       # dog = wpilib.GetWatchdog()
        #dog.setEnabled(True)
        #dog.SetExpiration(10)

        while self.IsOperatorControl()and self.IsEnabled():
          #  dog.Feed()
            self.drive.MecanumDrive_Cartesian(self.Joystick.GetY(), self.Joystick.GetX(), self.Joystick2.GetX(), 0)
            x = self.Joystick.GetButton()
            if x==3:
                self.jaguar6.Set(-1)
            if x==5:
                self.jaguar6.Set(1)
            wpilib.Wait(0.01)

'''

def run():

    robot = MyRobot()
    robot.StartCompetition()

    return robot

if __name__ == '__main__':
    wpilib.run()

<|endoftext|>


<commit_before>import struct
import threading
import queue
import socket

# server_port=0xC20A

class ClientCommand(object):
    """ A command to the client thread.
    Each command type has its associated data:
    CONNECT: (host, port) tuple
    SEND: Data string
    RECEIVE: None
    CLOSE: None
    """
    CONNECT, SEND, RECEIVE, CLOSE = range(4)
    def __init__(self, type, data = None):
        self.type = type
        self.data = data

class ClientReply(object):
    """ A reply from the client thread.
    Each reply type has its associated data:
    ERROR: The error string
    SUCCESS: Depends on the command - for RECEIVE it's the received
    data string, for others None.
    """
    ERROR, SUCCESS = range(2)
    def __init__(self, type, data=None):
        self.type = type
        self.data = data

# The socket thread acts as a server.

class SocketThread(threading.Thread):
    def __init__(self, command_queue=queue.Queue(), reply_queue=queue.Queue()):
        super(SocketThread, self).__init__()
        self._command_queue = command_queue
        self._reply_queue = reply_queue
        self.alive = threading.Event()
        self.alive.set()
        self._host = '127.0.0.1'
        self._port = None
        self._server_socket = None
        self._data_socket = None
        self.handlers = {
            ClientCommand.CONNECT: self._handle_CONNECT,
            ClientCommand.CLOSE: self._handle_CLOSE,
            ClientCommand.SEND: self._handle_SEND,
            ClientCommand.RECEIVE: self._handle_RECEIVE,
        }

    def getPort(self):
        return self._port

    def connectTo(self, host, port):
        self._command_queue.put(ClientCommand(ClientCommand.CONNECT, (port)))

    def sendCommand(self,command_id, data = None):
        packed_command = struct.pack('@i', int(command_id))
        if data is not None:
            packed_command += struct.pack('@i',data)
        self._command_queue.put(ClientCommand(ClientCommand.SEND, packed_command))

    ##  Return the next command_id & data that was recieved
    def getNextCommand(self):
        command_data = self._command_queue.get()
        return command_data.command_id , command_data.data

    def run(self):
        while self.alive.isSet():
            try:
                # Queue.get with timeout to allow checking self.alive
                command = self._command_queue.get(True, 0.1)
                self.handlers[command.type](command)
            except queue.Empty as e:
                continue

    def join(self, timeout = None):
        self.alive.clear()
        threading.Thread.join(self, timeout)

    def _handle_CONNECT(self, cmd):
        self._host = "127.0.0.1"
        self._port = 49674 #Hardcoded stuff
        self._server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        while True:
            try:
                self._server_socket.bind(('', self._port))
            except:
                self._port += 1
                if self._port > 0xFFFF:
                    break
            else:
                break
        print("Listening for backend on " + str(self._port))
        self._server_socket.listen(1)

        self._data_socket, address = self._server_socket.acscept()
        print("Backend connected on " + str(address))

    def _handle_CLOSE(self, cmd):
        self._data_socket.close()
        reply = ClientReply(ClientReply.SUCCESS)
        self._reply_queue.put(reply)

    def _handle_SEND(self, cmd):
        print(self._data_socket)
        try:
            self._data_socket.sendall(struct.pack('@i', len(cmd.data)))
            self._data_socket.sendall(cmd.data)
            self._reply_queue.put(self._success_reply())
        except IOError as e:
            print(e)
            self._reply_queue.put(self._error_reply(str(e)))

    def _handle_RECEIVE(self, cmd):
        try:
            message_length = self._recieveInt32()
            data = self._recieve_n_bytes(message_length)
            if len(data) == message_length:
                self._reply_queue.put(self._success_reply(data))
                return

            self._reply_queue.put(self._error_reply('Socket closed prematurely'))
        except IOError as e:
            self._reply_queue.put(self._error_reply(str(e)))

    def _recv_n_bytes(self, n):
        data = b''
        while len(data) < size:
            if self._data_socket is None:
                #Raise an IO error to signal the socket is closed.
                raise IOError()
            recieved = self._socket.recv(size - len(data))
            data += recieved
            if len(recieved) <= 0:
                #Raise an IO error to signal the socket is closed.
                raise IOError()
        return data

    def _recieveInt32(self):
        return struct.unpack('@i', self._recv_n_bytes(4))[0]

    def _error_reply(self, errstr):
        return ClientReply(ClientReply.ERROR, errstr)

    def _success_reply(self, data=None):
        return ClientReply(ClientReply.SUCCESS, data)<commit_msg>Fixed stupid typo & removed debug prints<commit_after>import struct
import threading
import queue
import socket

# server_port=0xC20A

class ClientCommand(object):
    """ A command to the client thread.
    Each command type has its associated data:
    CONNECT: (host, port) tuple
    SEND: Data string
    RECEIVE: None
    CLOSE: None
    """
    CONNECT, SEND, RECEIVE, CLOSE = range(4)
    def __init__(self, type, data = None):
        self.type = type
        self.data = data

class ClientReply(object):
    """ A reply from the client thread.
    Each reply type has its associated data:
    ERROR: The error string
    SUCCESS: Depends on the command - for RECEIVE it's the received
    data string, for others None.
    """
    ERROR, SUCCESS = range(2)
    def __init__(self, type, data=None):
        self.type = type
        self.data = data

# The socket thread acts as a server.

class SocketThread(threading.Thread):
    def __init__(self, command_queue=queue.Queue(), reply_queue=queue.Queue()):
        super(SocketThread, self).__init__()
        self._command_queue = command_queue
        self._reply_queue = reply_queue
        self.alive = threading.Event()
        self.alive.set()
        self._host = '127.0.0.1'
        self._port = None
        self._server_socket = None
        self._data_socket = None
        self.handlers = {
            ClientCommand.CONNECT: self._handle_CONNECT,
            ClientCommand.CLOSE: self._handle_CLOSE,
            ClientCommand.SEND: self._handle_SEND,
            ClientCommand.RECEIVE: self._handle_RECEIVE,
        }

    def getPort(self):
        return self._port

    def connectTo(self, host, port):
        self._command_queue.put(ClientCommand(ClientCommand.CONNECT, (port)))

    def sendCommand(self,command_id, data = None):
        packed_command = struct.pack('@i', int(command_id))
        if data is not None:
            packed_command += struct.pack('@i',data)
        self._command_queue.put(ClientCommand(ClientCommand.SEND, packed_command))

    ##  Return the next command_id & data that was recieved
    def getNextCommand(self):
        command_data = self._command_queue.get()
        return command_data.command_id , command_data.data

    def run(self):
        while self.alive.isSet():
            try:
                # Queue.get with timeout to allow checking self.alive
                command = self._command_queue.get(True, 0.1)
                self.handlers[command.type](command)
            except queue.Empty as e:
                continue

    def join(self, timeout = None):
        self.alive.clear()
        threading.Thread.join(self, timeout)

    def _handle_CONNECT(self, cmd):
        self._host = "127.0.0.1"
        self._port = 49674 #Hardcoded stuff
        self._server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        while True:
            try:
                self._server_socket.bind(('', self._port))
            except:
                self._port += 1
                if self._port > 0xFFFF:
                    break
            else:
                break
        print("Listening for backend on " + str(self._port))
        self._server_socket.listen(1)

        self._data_socket, address = self._server_socket.accept()
        print("Backend connected on " + str(address))

    def _handle_CLOSE(self, cmd):
        self._data_socket.close()
        reply = ClientReply(ClientReply.SUCCESS)
        self._reply_queue.put(reply)

    def _handle_SEND(self, cmd):
        try:
            self._data_socket.sendall(struct.pack('@i', len(cmd.data)))
            self._data_socket.sendall(cmd.data)
            self._reply_queue.put(self._success_reply())
        except IOError as e:
            print(e)
            self._reply_queue.put(self._error_reply(str(e)))

    def _handle_RECEIVE(self, cmd):
        try:
            message_length = self._recieveInt32()
            data = self._recieve_n_bytes(message_length)
            if len(data) == message_length:
                self._reply_queue.put(self._success_reply(data))
                return

            self._reply_queue.put(self._error_reply('Socket closed prematurely'))
        except IOError as e:
            self._reply_queue.put(self._error_reply(str(e)))

    def _recv_n_bytes(self, n):
        data = b''
        while len(data) < size:
            if self._data_socket is None:
                #Raise an IO error to signal the socket is closed.
                raise IOError()
            recieved = self._socket.recv(size - len(data))
            data += recieved
            if len(recieved) <= 0:
                #Raise an IO error to signal the socket is closed.
                raise IOError()
        return data

    def _recieveInt32(self):
        return struct.unpack('@i', self._recv_n_bytes(4))[0]

    def _error_reply(self, errstr):
        return ClientReply(ClientReply.ERROR, errstr)

    def _success_reply(self, data=None):
        return ClientReply(ClientReply.SUCCESS, data)<|endoftext|>


<commit_before>"""
Deployment settings file
"""

from settings import *
import json

DEBUG=False

TIME_BETWEEN_INDEX_REBUILDS = 60 * 30 # seconds

#Tastypie throttle settings
THROTTLE_AT = 100 #Throttle requests after this number in below timeframe
THROTTLE_TIMEFRAME= 60 * 60 #Timeframe in which to throttle N requests, seconds
THROTTLE_EXPIRATION= 24 * 60 * 60 # When to remove throttle entries from cache, seconds

with open(os.path.join(ENV_ROOT,"env.json")) as env_file:
    ENV_TOKENS = json.load(env_file)

with open(os.path.join(ENV_ROOT, "auth.json")) as auth_file:
    AUTH_TOKENS = json.load(auth_file)

DATABASES = AUTH_TOKENS.get('DATABASES', DATABASES)
CACHES = AUTH_TOKENS.get('CACHES', CACHES)

AWS_ACCESS_KEY_ID = AUTH_TOKENS.get('AWS_ACCESS_KEY_ID', AWS_ACCESS_KEY_ID)
AWS_SECRET_ACCESS_KEY = AUTH_TOKENS.get('AWS_SECRET_ACCESS_KEY', AWS_SECRET_ACCESS_KEY)

USE_S3_TO_STORE_MODELS = ENV_TOKENS.get('USE_S3_TO_STORE_MODELS', USE_S3_TO_STORE_MODELS)
S3_BUCKETNAME = ENV_TOKENS.get('S3_BUCKETNAME', S3_BUCKETNAME)

BROKER_URL = AUTH_TOKENS.get('BROKER_URL', BROKER_URL)
CELERY_RESULT_BACKEND = AUTH_TOKENS.get('CELERY_RESULT_BACKEND', CELERY_RESULT_BACKEND)

ELB_HOSTNAME = ENV_TOKENS.get('ELB_HOSTNAME', None)

DNS_HOSTNAME = ENV_TOKENS.get('DNS_HOSTNAME', None)

if ELB_HOSTNAME is not None:
    ALLOWED_HOSTS += [ELB_HOSTNAME]

if DNS_HOSTNAME is not None:
    ALLOWED_HOSTS += [DNS_HOSTNAME]

EMAIL_BACKEND = ENV_TOKENS.get('EMAIL_BACKEND', EMAIL_BACKEND)

DEFAULT_FROM_EMAIL = ENV_TOKENS.get('DEFAULT_FROM_EMAIL')

ACCOUNT_EMAIL_VERIFICATION = ENV_TOKENS.get('ACCOUNT_EMAIL_VERIFICATION', ACCOUNT_EMAIL_VERIFICATION)

AWS_SES_REGION_NAME = ENV_TOKENS.get('AWS_SES_REGION_NAME', 'us-east-1')
if AWS_SES_REGION_NAME is not None:
    AWS_SES_REGION_ENDPOINT = 'email.{0}.amazonaws.com'.format(AWS_SES_REGION_NAME)

#Try to set the domain for the current site
#Needed to get the right site name for email activation
try:
    if DNS_HOSTNAME is not None:
        from django.contrib.sites.models import Site
        current_site = Site.objects.get(id=SITE_ID)

        current_site.domain = DNS_HOSTNAME
        current_site.name = DNS_HOSTNAME
        current_site.save()
except:
    log.info("Could not set site name and domain.  Not a problem if this is a dev/sandbox environment.  May cause confusion with email activation in production.")
<commit_msg>Remove sites hack<commit_after>"""
Deployment settings file
"""

from settings import *
import json

DEBUG=False

TIME_BETWEEN_INDEX_REBUILDS = 60 * 30 # seconds

#Tastypie throttle settings
THROTTLE_AT = 100 #Throttle requests after this number in below timeframe
THROTTLE_TIMEFRAME= 60 * 60 #Timeframe in which to throttle N requests, seconds
THROTTLE_EXPIRATION= 24 * 60 * 60 # When to remove throttle entries from cache, seconds

with open(os.path.join(ENV_ROOT,"env.json")) as env_file:
    ENV_TOKENS = json.load(env_file)

with open(os.path.join(ENV_ROOT, "auth.json")) as auth_file:
    AUTH_TOKENS = json.load(auth_file)

DATABASES = AUTH_TOKENS.get('DATABASES', DATABASES)
CACHES = AUTH_TOKENS.get('CACHES', CACHES)

AWS_ACCESS_KEY_ID = AUTH_TOKENS.get('AWS_ACCESS_KEY_ID', AWS_ACCESS_KEY_ID)
AWS_SECRET_ACCESS_KEY = AUTH_TOKENS.get('AWS_SECRET_ACCESS_KEY', AWS_SECRET_ACCESS_KEY)

USE_S3_TO_STORE_MODELS = ENV_TOKENS.get('USE_S3_TO_STORE_MODELS', USE_S3_TO_STORE_MODELS)
S3_BUCKETNAME = ENV_TOKENS.get('S3_BUCKETNAME', S3_BUCKETNAME)

BROKER_URL = AUTH_TOKENS.get('BROKER_URL', BROKER_URL)
CELERY_RESULT_BACKEND = AUTH_TOKENS.get('CELERY_RESULT_BACKEND', CELERY_RESULT_BACKEND)

ELB_HOSTNAME = ENV_TOKENS.get('ELB_HOSTNAME', None)

DNS_HOSTNAME = ENV_TOKENS.get('DNS_HOSTNAME', None)

if ELB_HOSTNAME is not None:
    ALLOWED_HOSTS += [ELB_HOSTNAME]

if DNS_HOSTNAME is not None:
    ALLOWED_HOSTS += [DNS_HOSTNAME]

EMAIL_BACKEND = ENV_TOKENS.get('EMAIL_BACKEND', EMAIL_BACKEND)

DEFAULT_FROM_EMAIL = ENV_TOKENS.get('DEFAULT_FROM_EMAIL')

ACCOUNT_EMAIL_VERIFICATION = ENV_TOKENS.get('ACCOUNT_EMAIL_VERIFICATION', ACCOUNT_EMAIL_VERIFICATION)

AWS_SES_REGION_NAME = ENV_TOKENS.get('AWS_SES_REGION_NAME', 'us-east-1')
if AWS_SES_REGION_NAME is not None:
    AWS_SES_REGION_ENDPOINT = 'email.{0}.amazonaws.com'.format(AWS_SES_REGION_NAME)

#Try to set the domain for the current site
#Needed to get the right site name for email activation
#Comment out, as this is causing issues in deployment.
#TODO: Move to a fixture
"""
try:
    if DNS_HOSTNAME is not None:
        from django.contrib.sites.models import Site
        current_site = Site.objects.get(id=SITE_ID)

        current_site.domain = DNS_HOSTNAME
        current_site.name = DNS_HOSTNAME
        current_site.save()
except:
    log.info("Could not set site name and domain.  Not a problem if this is a dev/sandbox environment.  May cause confusion with email activation in production.")
"""<|endoftext|>


<commit_before>"""List the moles in a mole catalog."""

import datetime

import mel.micro.fs

def setup_parser(parser):
    parser.add_argument(
        '--only-no-micro',
        action='store_true',
        help="Only list moles that have no microscope images.")

    assistance = parser.add_mutually_exclusive_group()
    assistance.add_argument(
        '--without-assistance',
        action='store_true',
        help="Only list moles that don't require assistance to capture.")
    assistance.add_argument(
        '--with-assistance',
        action='store_true',
        help="Only list moles that require assistance to capture.")

    parser.add_argument(
        '--sort',
        default='unsorted',
        nargs='?',
        choices=['lastmicro'],
        help='Sort the moles by the date of their last micro image. '
             'This means moles with no images are first then the oldest '
             'images.')

    format_args = parser.add_mutually_exclusive_group()

    format_args.add_argument(
        '--format',
        default="{relpath}",
        help="Print the results with the specified format. Defaults to "
             "'{relpath}'. Available keys: relpath, lastmicro, "
             "lastmicro_age_days, id.")

    format_args.add_argument(
        '--ages',
        dest='format',
        action='store_const',
        const='{relpath} {lastmicro_age_days}',
        help="Print the relative paths of the moles and their ages.")

    format_args.add_argument(
        '--ids',
        dest='format',
        action='store_const',
        const='{relpath} {id}',
        help="Print the relative paths of the moles and their ids.")

def process_args(args):
    if args.sort == 'lastmicro' or args.sort is None:
        def keyfunc(x):
            if not x.micro_image_names:
                return str()
            return sorted(x.micro_image_names)[-1]
    else:
        assert(args.sort == 'unsorted')

        def keyfunc(x):
            return 0

    now = datetime.datetime.now()

    for mole in sorted(_yield_mole_dirs('.', args), key=keyfunc):
        mole_data = {
            'relpath': mole.refrelpath,
            'lastmicro': '',
            'lastmicro_age_days': '',
            'id': mole.id,
        }
        if mole.micro_image_names:
            lastmicro = sorted(mole.micro_image_names)[-1]
            mole_data['lastmicro'] = lastmicro
            lastmicrodtstring = lastmicro.split('.', 1)[0]
            lastmicrodt = datetime.datetime.strptime(
                lastmicrodtstring, '%Y%m%dT%H%M%S')
            age = now - lastmicrodt
            mole_data['lastmicro_age_days'] = age.days

        print(args.format.format(**mole_data))

def _yield_mole_dirs(rootpath, args):
    for mole in mel.micro.fs.yield_moles(rootpath):

        if args.only_no_micro and mole.micro_image_names:
            continue

        if args.without_assistance and mole.need_assistance:
            continue

        if args.with_assistance and not mole.need_assistance:
            continue

        yield mole
<commit_msg>Refactor mel.cmd.list, sort in _yield_mole_dirs<commit_after>"""List the moles in a mole catalog."""

import datetime

import mel.micro.fs

def setup_parser(parser):
    parser.add_argument(
        '--only-no-micro',
        action='store_true',
        help="Only list moles that have no microscope images.")

    assistance = parser.add_mutually_exclusive_group()
    assistance.add_argument(
        '--without-assistance',
        action='store_true',
        help="Only list moles that don't require assistance to capture.")
    assistance.add_argument(
        '--with-assistance',
        action='store_true',
        help="Only list moles that require assistance to capture.")

    parser.add_argument(
        '--sort',
        default='unsorted',
        nargs='?',
        choices=['lastmicro'],
        help='Sort the moles by the date of their last micro image. '
             'This means moles with no images are first then the oldest '
             'images.')

    format_args = parser.add_mutually_exclusive_group()

    format_args.add_argument(
        '--format',
        default="{relpath}",
        help="Print the results with the specified format. Defaults to "
             "'{relpath}'. Available keys: relpath, lastmicro, "
             "lastmicro_age_days, id.")

    format_args.add_argument(
        '--ages',
        dest='format',
        action='store_const',
        const='{relpath} {lastmicro_age_days}',
        help="Print the relative paths of the moles and their ages.")

    format_args.add_argument(
        '--ids',
        dest='format',
        action='store_const',
        const='{relpath} {id}',
        help="Print the relative paths of the moles and their ids.")

def process_args(args):

    now = datetime.datetime.now()

    for mole in _yield_mole_dirs('.', args):
        mole_data = {
            'relpath': mole.refrelpath,
            'lastmicro': '',
            'lastmicro_age_days': '',
            'id': mole.id,
        }
        if mole.micro_image_names:
            lastmicro = sorted(mole.micro_image_names)[-1]
            mole_data['lastmicro'] = lastmicro
            lastmicrodtstring = lastmicro.split('.', 1)[0]
            lastmicrodt = datetime.datetime.strptime(
                lastmicrodtstring, '%Y%m%dT%H%M%S')
            age = now - lastmicrodt
            mole_data['lastmicro_age_days'] = age.days

        print(args.format.format(**mole_data))

def _yield_mole_dirs(rootpath, args):

    mole_iter = mel.micro.fs.yield_moles(rootpath)

    if args.sort == 'lastmicro' or args.sort is None:
        def keyfunc(x):
            if not x.micro_image_names:
                return str()
            return sorted(x.micro_image_names)[-1]
        mole_iter = sorted(mole_iter, key=keyfunc)

    for mole in mole_iter:

        if args.only_no_micro and mole.micro_image_names:
            continue

        if args.without_assistance and mole.need_assistance:
            continue

        if args.with_assistance and not mole.need_assistance:
            continue

        yield mole
<|endoftext|>


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

import os
import json
from sys import argv
from time import sleep
from subprocess import getstatusoutput
import logging
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
from scapy.all import ARP
from scapy.layers.l2 import arping
from getpass import getuser
import platform
from pwd import getpwall
import argparse

# Color class used to print colors
class Colors:
    PURPLE = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    END = '\033[0m'

class Whoshome:

    def __init__(self, args):
        self._interface = args[0]
        self._output_file_mode = args[1]
        self._output_filename = args[2]
        self._max_cycles = args[3]
        self._people = self._make_people_list(self._open_people_file())

    def _open_people_file(self):
        # Try to open .people.json in every user's home directory. This should
        # work since there should be just one config file system-wide.
        for p in getpwall():
            home_path = os.path.expanduser('~' + p[0]) + '/'
            try:
                people_json = None
                people_file = open(home_path + '.people.json', 'r')
                people_json = json.load(people_file)
                people_file.close()
                break
            except:
                pass
        if people_json == None:
            print(Colors.RED + 'ERROR opening .people.json' + Colors.END)
            exit(1)
        else:
            return people_json

    def _make_people_list(self, people_json):
        people = list()
        allowed = '1234567890abcdef:'
        for person_dict in people_json:
            person_dict['target'] = person_dict['target'].lower()
            for c in person_dict['target']:
                if c not in allowed:
                    print(Colors.RED + 'ERROR: invalid character found in one or more MAC addresses' + Colors.END)
                    exit(1)
            if len(person_dict['target']) == 17:
                person_dict['target'] = person_dict['target'][9:]
            people.append(person_dict)
        for person in people:
            person['last_seen'] = self._max_cycles
        return people

    def _create_json(self, file):
        json_obj = list()
        for person in self._people:
            temp_dict = dict()
            temp_dict['name'] = person['name']
            temp_dict['target'] = person['target']
            temp_dict['home'] = bool(person['last_seen'] < self._max_cycles)
            json_obj.append(temp_dict)
        json.dump(json_obj, file, indent=4)

    def _get_ip_from_interface(self):
        output = getstatusoutput('ip a | grep ' + self._interface + ' | grep inet')
        if output[0] == 0:
            return output[1][output[1].find('inet') + 5: output[1].find('brd') - 1]
        else:
            print(Colors.RED + 'ERROR: invalid interface' + Colors.END)
            exit(1)

    def cycle(self):
        while True:
            try:
                results, unanswered = arping(self._get_ip_from_interface(), verbose=False)
                if self._output_file_mode != 'no':
                    if self._output_file_mode != 'both':
                        file = open(self._output_filename, 'w')
                    else:
                        file_txt = open(self._output_filename + '.txt', 'w')
                        file_json = open(self._output_filename + '.json', 'w')
                for result in results:
                    # A MAC address is 17 characters long. Only the last 3 bytes of the MAC
                    # address are taken into account, to ensure compatibility with some
                    # network devices which may change the vendor part of MAC addresses
                    mac = result[1][ARP].hwsrc[9:]
                    for person in self._people:
                        if mac == person['target']:
                            # The counter is set to -1 because every counter will be incremented in
                            # the next 'for' cycle
                            person['last_seen'] = -1
                for person in self._people:
                    if person['last_seen'] < self._max_cycles:
                        person['last_seen'] += 1
                        print(Colors.GREEN + '🏡  ' + person['name'] + ' is home' + Colors.END)
                        if self._output_file_mode == 'txt':
                            file.write('🏡 ' + person['name'] + ' is home\n')
                        elif self._output_file_mode == 'both':
                            file_txt.write('🏡 ' + person['name'] + ' is home\n')
                    else:
                        print(Colors.PURPLE + '🌍  ' + person['name'] + ' is away' + Colors.END)
                        if self._output_file_mode == 'txt':
                            file.write('🌍 ' + person['name'] + ' is away\n')
                        elif self._output_file_mode == 'both':
                            file_txt.write('🌍 ' + person['name'] + ' is away\n')
                print()
                if self._output_file_mode != 'no':
                    if self._output_file_mode == 'json':
                        self._create_json(file)
                    elif self._output_file_mode == 'both':
                        self._create_json(file_json)
                    try:
                        file_txt.close()
                        file_json.close()
                    except:
                        file.close()
                sleep(30)
            except KeyboardInterrupt:
                print('\nQuit')
                exit(0)

def print_help():
    print(Colors.YELLOW + 'Usage:    ' + argv[0] + ' interface [option | filename] [max_cycles]')
    print('Options:')
    print('    -t | txt file output')
    print('    -j | json file output')
    print('    -o | txt and json file output (don\'t write file extension in filename)')
    print('The developer declines every responsibility in case of malfunction due to non observance of this tiny guide' + Colors.END)

def check_environment():
    if platform.system() != 'Linux':
        print(Colors.RED + 'ERROR: system not supported' + Colors.END)
        exit(1)
    if getuser() != 'root':
        print(Colors.RED + 'ERROR: must be run as root' + Colors.END)
        exit(1)

def parse_argv(passed_args=None):
    parser = argparse.ArgumentParser(
        description='Who\'s Home  -  Find out who\'s home based on Wi-Fi connection')
    parser.add_argument('interface', type=str, help='Interface used to send ARP-Requests.')
    parser.add_argument('-o', '--output', type=str,
                        help='Send results to a file. Available file extensions are \'.txt\' and \'.json\'. The file formatting will be inferred from the file extension. If you want to have both file formatting, omit the file extension.')
    parser.add_argument('-c', '--max-cycles', type=int, default=30,
                        help='Alter `max_cycles` variable to modify the period of time in which a person is considered at home.')
    args = parser.parse_args(passed_args)
    interface = args.interface
    if args.output == None:
        output_file_mode = 'no'
        output_filename = str()
    else:
        allowed_extensions = ['txt', 'json']
        file_split = args.output.split('.')
        file_extension = file_split[-1]
        if file_extension in allowed_extensions:
            output_file_mode = file_extension
            output_filename = args.output
        else:
            output_file_mode = 'both'
            output_filename = args.output
    max_cycles = args.max_cycles
    return (interface, output_file_mode, output_filename, max_cycles)

"""
def parse_argv():
    argc = len(argv)
    interface = str()
    output_file_mode = 'no'
    output_filename = str()
    max_cycles = 30
    if argc < 2 or argc > 5:
        print(Colors.RED + 'ERROR: wrong arguments' + Colors.END)
        print_help()
        exit(1)
    else:
        if argv[1] == '-h' or argv[1] == '--help':
            print_help()
            exit(1)
        interface = argv[1]
        if argc % 2:
            max_cycles = int(argv[argc - 1])
        if argc >= 4:
            if argv[2] == '-j':
                output_file_mode = 'json'
                output_filename = argv[3]
                if output_filename[-4:] != output_file_mode:
                    print(Colors.RED + 'ERROR: file extension' + Colors.END)
                    exit(1)
            elif argv[2] == '-t':
                output_file_mode = 'txt'
                output_filename = argv[3]
                if output_filename[-3:] != output_file_mode:
                    print(Colors.RED + 'ERROR: file extension' + Colors.END)
                    exit(1)
            elif argv[2] == '-o':
                output_file_mode = 'both'
                output_filename = argv[3]
            else:
                print(Colors.RED + 'ERROR: wrong arguments' + Colors.END)
                print_help()
                exit(1)
    return (interface, output_file_mode, output_filename, max_cycles)
"""

def main():
    check_environment()
    wh = Whoshome(parse_argv())
    wh.cycle()

if __name__ == '__main__':
    main()
<commit_msg>Integrate argv parser<commit_after>#!/usr/bin/env python3

import os
import json
from sys import argv
from time import sleep
from subprocess import getstatusoutput
import logging
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
from scapy.all import ARP
from scapy.layers.l2 import arping
from getpass import getuser
import platform
from pwd import getpwall
import argparse

# Color class used to print colors
class Colors:
    PURPLE = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    END = '\033[0m'

class Whoshome:

    def __init__(self, args):
        self._interface = args[0]
        self._output_file_mode = args[1]
        self._output_filename = args[2]
        self._max_cycles = args[3]
        self._people = self._make_people_list(self._open_people_file())

    def _open_people_file(self):
        # Try to open .people.json in every user's home directory. This should
        # work since there should be just one config file system-wide.
        for p in getpwall():
            home_path = os.path.expanduser('~' + p[0]) + '/'
            try:
                people_json = None
                people_file = open(home_path + '.people.json', 'r')
                people_json = json.load(people_file)
                people_file.close()
                break
            except:
                pass
        if people_json == None:
            print(Colors.RED + 'ERROR opening .people.json' + Colors.END)
            exit(1)
        else:
            return people_json

    def _make_people_list(self, people_json):
        people = list()
        allowed = '1234567890abcdef:'
        for person_dict in people_json:
            person_dict['target'] = person_dict['target'].lower()
            for c in person_dict['target']:
                if c not in allowed:
                    print(Colors.RED + 'ERROR: invalid character found in one or more MAC addresses' + Colors.END)
                    exit(1)
            if len(person_dict['target']) == 17:
                person_dict['target'] = person_dict['target'][9:]
            people.append(person_dict)
        for person in people:
            person['last_seen'] = self._max_cycles
        return people

    def _create_json(self, file):
        json_obj = list()
        for person in self._people:
            temp_dict = dict()
            temp_dict['name'] = person['name']
            temp_dict['target'] = person['target']
            temp_dict['home'] = bool(person['last_seen'] < self._max_cycles)
            json_obj.append(temp_dict)
        json.dump(json_obj, file, indent=4)

    def _get_ip_from_interface(self):
        output = getstatusoutput('ip a | grep ' + self._interface + ' | grep inet')
        if output[0] == 0:
            return output[1][output[1].find('inet') + 5: output[1].find('brd') - 1]
        else:
            print(Colors.RED + 'ERROR: invalid interface' + Colors.END)
            exit(1)

    def cycle(self):
        while True:
            try:
                results, unanswered = arping(self._get_ip_from_interface(), verbose=False)
                if self._output_file_mode != 'no':
                    if self._output_file_mode != 'both':
                        file = open(self._output_filename, 'w')
                    else:
                        file_txt = open(self._output_filename + '.txt', 'w')
                        file_json = open(self._output_filename + '.json', 'w')
                for result in results:
                    # A MAC address is 17 characters long. Only the last 3 bytes of the MAC
                    # address are taken into account, to ensure compatibility with some
                    # network devices which may change the vendor part of MAC addresses
                    mac = result[1][ARP].hwsrc[9:]
                    for person in self._people:
                        if mac == person['target']:
                            # The counter is set to -1 because every counter will be incremented in
                            # the next 'for' cycle
                            person['last_seen'] = -1
                for person in self._people:
                    if person['last_seen'] < self._max_cycles:
                        person['last_seen'] += 1
                        print(Colors.GREEN + '🏡  ' + person['name'] + ' is home' + Colors.END)
                        if self._output_file_mode == 'txt':
                            file.write('🏡 ' + person['name'] + ' is home\n')
                        elif self._output_file_mode == 'both':
                            file_txt.write('🏡 ' + person['name'] + ' is home\n')
                    else:
                        print(Colors.PURPLE + '🌍  ' + person['name'] + ' is away' + Colors.END)
                        if self._output_file_mode == 'txt':
                            file.write('🌍 ' + person['name'] + ' is away\n')
                        elif self._output_file_mode == 'both':
                            file_txt.write('🌍 ' + person['name'] + ' is away\n')
                print()
                if self._output_file_mode != 'no':
                    if self._output_file_mode == 'json':
                        self._create_json(file)
                    elif self._output_file_mode == 'both':
                        self._create_json(file_json)
                    try:
                        file_txt.close()
                        file_json.close()
                    except:
                        file.close()
                sleep(30)
            except KeyboardInterrupt:
                print('\nQuit')
                exit(0)

def print_help():
    print(Colors.YELLOW + 'Usage:    ' + argv[0] + ' interface [option | filename] [max_cycles]')
    print('Options:')
    print('    -t | txt file output')
    print('    -j | json file output')
    print('    -o | txt and json file output (don\'t write file extension in filename)')
    print('The developer declines every responsibility in case of malfunction due to non observance of this tiny guide' + Colors.END)

def check_environment():
    if platform.system() != 'Linux':
        print(Colors.RED + 'ERROR: system not supported' + Colors.END)
        exit(1)
    if getuser() != 'root':
        print(Colors.RED + 'ERROR: must be run as root' + Colors.END)
        parse_argv()
        exit(1)

def parse_argv(passed_args=None):
    parser = argparse.ArgumentParser(
        description='Who\'s Home  -  Find out who\'s home based on Wi-Fi connection')
    parser.add_argument('interface', type=str, help='Interface used to send ARP-Requests.')
    parser.add_argument('-o', '--output', type=str,
                        help='Send results to a file. Available file extensions are \'.txt\' and \'.json\'. The file formatting will be inferred from the file extension. If you want to have both file formatting, omit the file extension.')
    parser.add_argument('-c', '--max-cycles', type=int, default=30,
                        help='Alter `max_cycles` variable to modify the period of time in which a person is considered at home.')
    args = parser.parse_args(passed_args)
    interface = args.interface
    if args.output == None:
        output_file_mode = 'no'
        output_filename = str()
    else:
        allowed_extensions = ['txt', 'json']
        file_split = args.output.split('.')
        file_extension = file_split[-1]
        if file_extension in allowed_extensions:
            output_file_mode = file_extension
            output_filename = args.output
        else:
            output_file_mode = 'both'
            output_filename = args.output
    max_cycles = args.max_cycles
    return (interface, output_file_mode, output_filename, max_cycles)

"""
def parse_argv():
    argc = len(argv)
    interface = str()
    output_file_mode = 'no'
    output_filename = str()
    max_cycles = 30
    if argc < 2 or argc > 5:
        print(Colors.RED + 'ERROR: wrong arguments' + Colors.END)
        print_help()
        exit(1)
    else:
        if argv[1] == '-h' or argv[1] == '--help':
            print_help()
            exit(1)
        interface = argv[1]
        if argc % 2:
            max_cycles = int(argv[argc - 1])
        if argc >= 4:
            if argv[2] == '-j':
                output_file_mode = 'json'
                output_filename = argv[3]
                if output_filename[-4:] != output_file_mode:
                    print(Colors.RED + 'ERROR: file extension' + Colors.END)
                    exit(1)
            elif argv[2] == '-t':
                output_file_mode = 'txt'
                output_filename = argv[3]
                if output_filename[-3:] != output_file_mode:
                    print(Colors.RED + 'ERROR: file extension' + Colors.END)
                    exit(1)
            elif argv[2] == '-o':
                output_file_mode = 'both'
                output_filename = argv[3]
            else:
                print(Colors.RED + 'ERROR: wrong arguments' + Colors.END)
                print_help()
                exit(1)
    return (interface, output_file_mode, output_filename, max_cycles)
"""

def main():
    check_environment()
    wh = Whoshome(parse_argv())
    wh.cycle()

if __name__ == '__main__':
    main()
<|endoftext|>