<commit_before>from hsv_to_rgb import hsv_to_rgb as htr
from rgb_to_hsv import rgb_to_hsv as rth
from simple import MQTTClient
from neopixel import NeoPixel
from machine import Pin
from time import sleep
from sys import exit
import ubinascii
import machine
import network
# ESP8266 ESP-12 modules have blue, active-low LED on GPIO2, replace
# with something else if needed.
'''
light:
  - platform: mqtt
    name: test
    state_topic: 'led/status'
    command_topic: 'led/switch'
    brightness_state_topic: 'led/brightness/status'
    brightness_command_topic: 'led/brightness/set'
    rgb_state_topic: 'led/rgb/status'
    rgb_command_topic: 'led/rgb/set'
'''
# Wifi
SSID = "Home"
PASS = "***REMOVED***"

# MQTT server to connect to
SERVER = "172.16.0.30"
CLIENT_ID = ubinascii.hexlify(machine.unique_id())
STATE_TOPIC= b"led/status"
COMMAND_TOPIC = b"led/switch"
BRIGHTNESS_STATE_TOPIC = b'led/brightness/status'
BRIGHTNESS_COMMAND_TOPIC = b'led/brightness/set'
RGB_STATE_TOPIC = b'led/rgb/status'
RGB_COMMAND_TOPIC = b'led/rgb/set'
KILL_COMMAND_TOPIC = b'led/kill'
KILL_STATE_TOPIC = b'led/alive'

# Pin for NODEMCU and Neopixel
LED_PIN = 5
NUM_OF_LEDS = 6
PIN = Pin(LED_PIN, Pin.OUT)
NEOPIXEL = NeoPixel(PIN, NUM_OF_LEDS)
BRIGHTNESS_STATE = 0    # last known brightness state (DO NOT NEED)
STATE = "OFF"           # last known state (ON, OFF)
RGB_STATE = (0, 0, 0)   # last known set RGB state
LIGHT_STATE = (0, 0, 0) # last known light state

def publish_payload(topic, payload):
    '''publish to a topic with a payload'''
    c.publish(topic, str(payload))
    print("[Publish]\t%s => %s" % (topic, str(payload)))

def get_hsv_value(colors):
    '''returns value in hsv'''
    r, g, b = [int(color)/255.0 for color in colors]
    h, s, v = rth(r, g, b)
    return h, s, v

def get_rgb_value(color, new_value):
    '''returns value in rgb'''
    h, s, _ = get_hsv_value(color)
    r, g, b = [int(round(color*255)) for color in htr(h, s, new_value)]
    return r, g, b

def set_led(on=False, color=(0,0,0)):
    '''
    on is True or False
    color is list [red, green, blue]
    '''
    if not on:
        smooth(color, 50)
    elif on:
        if color:
            smooth(color, 75)

def smooth(new, smooth, transition=0):
    '''
    new is the new color (RED, GREEN, BLUE)
    smooth is how many "frames" of color to create and display
    transition is seconds to transit from one color to the next
    transition is 0 right now, but will be implimented with MQTT JSON
    '''
    global LIGHT_STATE
    red = list(linspace(LIGHT_STATE[0], new[0], smooth))
    green = list(linspace(LIGHT_STATE[1], new[1], smooth))
    blue = list(linspace(LIGHT_STATE[2], new[2], smooth))
    for color in zip(red,green,blue):
        for led in range(NUM_OF_LEDS):
            NEOPIXEL[led] = color
        NEOPIXEL.write()
        sleep(transition/smooth)
    LIGHT_STATE = new

def linspace(a,b, n=100):
    if n < 2:
        yield b
    diff = (float(b) - a)/(n - 1)
    for i in range(n):
        yield int(round(diff * i + a))

def sub_cb(topic, msg):
    global STATE
    global RGB_STATE
    if topic == COMMAND_TOPIC:
        if msg == b"ON" and STATE == "OFF":
            STATE = "ON"
            RGB_STATE = (255, 255, 255)
            publish_payload(STATE_TOPIC, "ON")
            publish_payload(BRIGHTNESS_STATE_TOPIC, 255)
            set_led(on=True, color=RGB_STATE)
        elif msg == b"OFF" and STATE == "ON":
            STATE = "OFF"
            RGB_STATE = (0, 0, 0)
            publish_payload(STATE_TOPIC, "OFF")
            publish_payload(BRIGHTNESS_STATE_TOPIC, 0)
            set_led(on=False, color=RGB_STATE)
    elif topic == BRIGHTNESS_COMMAND_TOPIC:
        brightness = int(msg)
        if brightness < 0 or brightness > 255:
            return
        else:
            colors = tuple(get_rgb_value(RGB_STATE, brightness/255.0))
            publish_payload(BRIGHTNESS_STATE_TOPIC, brightness)
            publish_payload(RGB_STATE_TOPIC, ",".join([str(color) for color in colors]))
            set_led(on=True, color=colors)
    elif topic == RGB_COMMAND_TOPIC:
        colors = msg.decode("utf-8").split(',')
        RGB_STATE = colors
        _, _, v = get_hsv_value(colors)
        publish_payload(BRIGHTNESS_STATE_TOPIC, int(v*255))
        publish_payload(RGB_STATE_TOPIC, ",".join(colors))
        set_led(on=True, color=tuple([int(color) for color in colors]))
    elif topic == KILL_COMMAND_TOPIC:
        if msg == b'dead':
            publish_payload(STATE_TOPIC, "OFF")
            publish_payload(BRIGHTNESS_STATE_TOPIC, 0)
            publish_payload(KILL_STATE_TOPIC, "dead")
            set_led()
            exit()
        else:
            publish_payload(KILL_STATE_TOPIC, "alive")

def reconnect():
    c.set_callback(sub_cb)
    c.connect()
    publish_payload(STATE_TOPIC, STATE)
    publish_payload(BRIGHTNESS_STATE_TOPIC, 0)
    publish_payload(RGB_STATE_TOPIC, LIGHT_STATE)
    publish_payload(KILL_STATE_TOPIC, "alive")
    c.subscribe(COMMAND_TOPIC)
    c.subscribe(BRIGHTNESS_COMMAND_TOPIC)
    c.subscribe(RGB_COMMAND_TOPIC)
    c.subscribe(KILL_COMMAND_TOPIC)

def main(server=SERVER):
    global c
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        print('connecting to network...')
        wlan.connect(SSID, PASS)
        while not wlan.isconnected():
            pass
    print('[Network Config]\n ', wlan.ifconfig())
    c = MQTTClient(CLIENT_ID, server)
    reconnect()
    set_led()       # sets led to off
    try:
        while True:
            if not wlan.isconnected():
                reconnect()
            c.wait_msg()
    finally:
        c.disconnect()
if __name__ == "__main__":
    main()<commit_msg>add kill switch config<commit_after>from hsv_to_rgb import hsv_to_rgb as htr
from rgb_to_hsv import rgb_to_hsv as rth
from simple import MQTTClient
from neopixel import NeoPixel
from machine import Pin
from time import sleep
from sys import exit
import ubinascii
import machine
import network
# ESP8266 ESP-12 modules have blue, active-low LED on GPIO2, replace
# with something else if needed.
'''
light:
  - platform: mqtt
    name: test
    state_topic: 'led/status'
    command_topic: 'led/switch'
    brightness_state_topic: 'led/brightness/status'
    brightness_command_topic: 'led/brightness/set'
    rgb_state_topic: 'led/rgb/status'
    rgb_command_topic: 'led/rgb/set'

switch:
  - platform: mqtt
    name: test mqtt kill switch
    command_topic: 'led/kill'
    state_topic: 'led/alive'
    payload_on: "alive"
    payload_off: "dead"
'''
# Wifi
SSID = ""
PASS = ""

# MQTT server to connect to
SERVER = "172.16.0.30"
CLIENT_ID = ubinascii.hexlify(machine.unique_id())
STATE_TOPIC= b"led/status"
COMMAND_TOPIC = b"led/switch"
BRIGHTNESS_STATE_TOPIC = b'led/brightness/status'
BRIGHTNESS_COMMAND_TOPIC = b'led/brightness/set'
RGB_STATE_TOPIC = b'led/rgb/status'
RGB_COMMAND_TOPIC = b'led/rgb/set'
KILL_COMMAND_TOPIC = b'led/kill'
KILL_STATE_TOPIC = b'led/alive'

# Pin for NODEMCU and Neopixel
LED_PIN = 5
NUM_OF_LEDS = 6
PIN = Pin(LED_PIN, Pin.OUT)
NEOPIXEL = NeoPixel(PIN, NUM_OF_LEDS)
BRIGHTNESS_STATE = 0    # last known brightness state (DO NOT NEED)
STATE = "OFF"           # last known state (ON, OFF)
RGB_STATE = (0, 0, 0)   # last known set RGB state
LIGHT_STATE = (0, 0, 0) # last known light state

def publish_payload(topic, payload):
    '''publish to a topic with a payload'''
    c.publish(topic, str(payload))
    print("[Publish]\t%s => %s" % (topic, str(payload)))

def get_hsv_value(colors):
    '''returns value in hsv'''
    r, g, b = [int(color)/255.0 for color in colors]
    h, s, v = rth(r, g, b)
    return h, s, v

def get_rgb_value(color, new_value):
    '''returns value in rgb'''
    h, s, _ = get_hsv_value(color)
    r, g, b = [int(round(color*255)) for color in htr(h, s, new_value)]
    return r, g, b

def set_led(on=False, color=(0,0,0)):
    '''
    on is True or False
    color is list [red, green, blue]
    '''
    if not on:
        smooth(color, 50)
    elif on:
        if color:
            smooth(color, 75)

def smooth(new, smooth, transition=0):
    '''
    new is the new color (RED, GREEN, BLUE)
    smooth is how many "frames" of color to create and display
    transition is seconds to transit from one color to the next
    transition is 0 right now, but will be implimented with MQTT JSON
    '''
    global LIGHT_STATE
    red = list(linspace(LIGHT_STATE[0], new[0], smooth))
    green = list(linspace(LIGHT_STATE[1], new[1], smooth))
    blue = list(linspace(LIGHT_STATE[2], new[2], smooth))
    for color in zip(red,green,blue):
        for led in range(NUM_OF_LEDS):
            NEOPIXEL[led] = color
        NEOPIXEL.write()
        sleep(transition/smooth)
    LIGHT_STATE = new

def linspace(a,b, n=100):
    if n < 2:
        yield b
    diff = (float(b) - a)/(n - 1)
    for i in range(n):
        yield int(round(diff * i + a))

def sub_cb(topic, msg):
    global STATE
    global RGB_STATE
    if topic == COMMAND_TOPIC:
        if msg == b"ON" and STATE == "OFF":
            STATE = "ON"
            RGB_STATE = (255, 255, 255)
            publish_payload(STATE_TOPIC, "ON")
            publish_payload(BRIGHTNESS_STATE_TOPIC, 255)
            set_led(on=True, color=RGB_STATE)
        elif msg == b"OFF" and STATE == "ON":
            STATE = "OFF"
            RGB_STATE = (0, 0, 0)
            publish_payload(STATE_TOPIC, "OFF")
            publish_payload(BRIGHTNESS_STATE_TOPIC, 0)
            set_led(on=False, color=RGB_STATE)
    elif topic == BRIGHTNESS_COMMAND_TOPIC:
        brightness = int(msg)
        if brightness < 0 or brightness > 255:
            return
        else:
            colors = tuple(get_rgb_value(RGB_STATE, brightness/255.0))
            publish_payload(BRIGHTNESS_STATE_TOPIC, brightness)
            publish_payload(RGB_STATE_TOPIC, ",".join([str(color) for color in colors]))
            set_led(on=True, color=colors)
    elif topic == RGB_COMMAND_TOPIC:
        colors = msg.decode("utf-8").split(',')
        RGB_STATE = colors
        _, _, v = get_hsv_value(colors)
        publish_payload(BRIGHTNESS_STATE_TOPIC, int(v*255))
        publish_payload(RGB_STATE_TOPIC, ",".join(colors))
        set_led(on=True, color=tuple([int(color) for color in colors]))
    elif topic == KILL_COMMAND_TOPIC:
        if msg == b'dead':
            publish_payload(STATE_TOPIC, "OFF")
            publish_payload(BRIGHTNESS_STATE_TOPIC, 0)
            publish_payload(KILL_STATE_TOPIC, "dead")
            set_led()
            exit()
        else:
            publish_payload(KILL_STATE_TOPIC, "alive")

def reconnect():
    c.set_callback(sub_cb)
    c.connect()
    publish_payload(STATE_TOPIC, STATE)
    publish_payload(BRIGHTNESS_STATE_TOPIC, 0)
    publish_payload(RGB_STATE_TOPIC, LIGHT_STATE)
    publish_payload(KILL_STATE_TOPIC, "alive")
    c.subscribe(COMMAND_TOPIC)
    c.subscribe(BRIGHTNESS_COMMAND_TOPIC)
    c.subscribe(RGB_COMMAND_TOPIC)
    c.subscribe(KILL_COMMAND_TOPIC)

def main(server=SERVER):
    global c
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        print('connecting to network...')
        wlan.connect(SSID, PASS)
        while not wlan.isconnected():
            pass
    print('[Network Config]\n ', wlan.ifconfig())
    c = MQTTClient(CLIENT_ID, server)
    reconnect()
    set_led()       # sets led to off
    try:
        while True:
            if not wlan.isconnected():
                reconnect()
            c.wait_msg()
    finally:
        c.disconnect()
if __name__ == "__main__":
    main()<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##
## $Id$
##
## This file is part of CDS Invenio.
## Copyright (C) 2002, 2003, 2004, 2005, 2006, 2007, 2008 CERN.
##
## CDS Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## CDS Invenio 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
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with CDS Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.

"""WebSearch Admin Regression Test Suite."""

__revision__ = "$Id$"

import unittest

from invenio.config import weburl
from invenio.testutils import make_test_suite, warn_user_about_tests_and_run, \
                              test_web_page_content, merge_error_messages

class WebSearchAdminWebPagesAvailabilityTest(unittest.TestCase):
    """Check WebSearch Admin web pages whether they are up or not."""

    def test_websearch_admin_interface_pages_availability(self):
        """websearchadmin - availability of WebSearch Admin interface pages"""

        baseurl = weburl + '/admin/websearch/websearchadmin.py'

        _exports = ['',
                    '?mtype=perform_showall',
                    '?mtype=perform_addcollection',
                    '?mtype=perform_addcollectiontotree',
                    '?mtype=perform_modifycollectiontree',
                    '?mtype=perform_checkwebcollstatus',
                    '?mtype=perform_checkcollectionstatus',]

        error_messages = []
        for url in [baseurl + page for page in _exports]:
            # first try as guest:
            error_messages.extend(test_web_page_content(url,
                                                        username='guest',
                                                        expected_text=
                                                        'Authorization failure'))
            # then try as admin:
            error_messages.extend(test_web_page_content(url,
                                                        username='admin'))
        if error_messages:
            self.fail(merge_error_messages(error_messages))
        return

    def test_websearch_admin_guide_availability(self):
        """websearchadmin - availability of WebSearch Admin guide pages"""

        url = weburl + '/help/admin/websearch-admin-guide'
        error_messages = test_web_page_content(url)
        if error_messages:
            self.fail(merge_error_messages(error_messages))
        return

test_suite = make_test_suite(WebSearchAdminWebPagesAvailabilityTest)

if __name__ == "__main__":
    warn_user_about_tests_and_run(test_suite)
<commit_msg>Better checking of admin guides availability.<commit_after># -*- coding: utf-8 -*-
##
## $Id$
##
## This file is part of CDS Invenio.
## Copyright (C) 2002, 2003, 2004, 2005, 2006, 2007, 2008 CERN.
##
## CDS Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## CDS Invenio 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
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with CDS Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.

"""WebSearch Admin Regression Test Suite."""

__revision__ = "$Id$"

import unittest

from invenio.config import weburl
from invenio.testutils import make_test_suite, warn_user_about_tests_and_run, \
                              test_web_page_content, merge_error_messages

class WebSearchAdminWebPagesAvailabilityTest(unittest.TestCase):
    """Check WebSearch Admin web pages whether they are up or not."""

    def test_websearch_admin_interface_pages_availability(self):
        """websearchadmin - availability of WebSearch Admin interface pages"""

        baseurl = weburl + '/admin/websearch/websearchadmin.py'

        _exports = ['',
                    '?mtype=perform_showall',
                    '?mtype=perform_addcollection',
                    '?mtype=perform_addcollectiontotree',
                    '?mtype=perform_modifycollectiontree',
                    '?mtype=perform_checkwebcollstatus',
                    '?mtype=perform_checkcollectionstatus',]

        error_messages = []
        for url in [baseurl + page for page in _exports]:
            # first try as guest:
            error_messages.extend(test_web_page_content(url,
                                                        username='guest',
                                                        expected_text=
                                                        'Authorization failure'))
            # then try as admin:
            error_messages.extend(test_web_page_content(url,
                                                        username='admin'))
        if error_messages:
            self.fail(merge_error_messages(error_messages))
        return

    def test_websearch_admin_guide_availability(self):
        """websearchadmin - availability of WebSearch Admin guide pages"""

        url = weburl + '/help/admin/websearch-admin-guide'
        error_messages = test_web_page_content(url,
                                               expected_text="WebSearch Admin Guide")
        if error_messages:
            self.fail(merge_error_messages(error_messages))
        return

test_suite = make_test_suite(WebSearchAdminWebPagesAvailabilityTest)

if __name__ == "__main__":
    warn_user_about_tests_and_run(test_suite)
<|endoftext|>


<commit_before>import numpy as np
import subprocess

def read_vector(filename):
    v = []
    with open(filename, "r") as fd:
        for line in fd:
            v.append(float(line))
    return np.array(v)

def check_vectors(filea="elimination.vector",
                  fileb="schur.vector"):
    veca = read_vector(filea)
    vecb = read_vector(fileb)
    return np.linalg.norm(veca - vecb) / np.linalg.norm(vecb)

def run(fd, t, mesh="../data/square-disc-p3.mesh", diffusion=False):
    commandline = ["./ex-normal",
                   "--reltol", "1.e-10",
                   "--mesh", mesh]
    if diffusion:
        commandline.append("--diffusion")
    if t[0] == "p":
        commandline = commandline + ["--penalty", t[1:]]
    elif t == "elimination":
        commandline = commandline + ["--elimination"]
    p = subprocess.Popen(commandline, stdout=fd, stderr=fd)
    p.communicate()

def main():
    with open("normalcheck.log", "w") as fd:
        for mesh in ["../data/square-disc-p3.mesh",
                     "icf.mesh",
                     "sphere_hex27.mesh"]:
            print(mesh)
            for d in [False]:
                print("  diffusion:", d)
                run(fd, "schur", mesh=mesh, diffusion=d)
                run(fd, "elimination", mesh=mesh, diffusion=d)
                print("    ", check_vectors())

def penalty(diffusion=False):
    print("Check penalty sanity, error should scale like 1/penalty:")
    with open("penalty.log", "w") as fd:
        run(fd, "schur")
        # run(fd, "elimination")
        for penalty in [1.0, 1.e+2, 1.e+4, 1.e+6]:
            run(fd, "p" + str(penalty), diffusion=diffusion)
            print("  penalty parameter:", penalty, ", error:",
                  check_vectors(filea="penalty.vector"))

if __name__ == "__main__":
    main()
    penalty()
    # penalty(True)
<commit_msg>Test --diffusion with boundary conditions, works reasonably well<commit_after>import numpy as np
import subprocess

def read_vector(filename):
    v = []
    with open(filename, "r") as fd:
        for line in fd:
            v.append(float(line))
    return np.array(v)

def check_vectors(filea="elimination.vector",
                  fileb="schur.vector"):
    veca = read_vector(filea)
    vecb = read_vector(fileb)
    return np.linalg.norm(veca - vecb) / np.linalg.norm(vecb)

def run(fd, t, mesh="../data/square-disc-p3.mesh", diffusion=False):
    commandline = ["./ex-normal",
                   "--reltol", "1.e-10",
                   "--mesh", mesh]
    if diffusion:
        commandline.append("--diffusion")
        commandline = commandline + ["--boundary-attribute", "1"]
    if t[0] == "p":
        commandline = commandline + ["--penalty", t[1:]]
    elif t == "elimination":
        commandline = commandline + ["--elimination"]
    p = subprocess.Popen(commandline, stdout=fd, stderr=fd)
    p.communicate()

def main():
    with open("normalcheck.log", "w") as fd:
        for mesh in ["../data/square-disc-p3.mesh",
                     "icf.mesh",
                     "sphere_hex27.mesh"]:
            print(mesh)
            ds = [False]
            if mesh == "../data/square-disc-p3.mesh":
                ds = [False, True]
            for d in ds:
                print("  diffusion:", d)
                run(fd, "schur", mesh=mesh, diffusion=d)
                run(fd, "elimination", mesh=mesh, diffusion=d)
                print("    ", check_vectors())

def penalty(diffusion=False):
    print("Check penalty sanity, error should scale like 1/penalty:")
    with open("penalty.log", "w") as fd:
        run(fd, "schur", diffusion=diffusion)
        for penalty in [1.0, 1.e+2, 1.e+4, 1.e+6]:
            run(fd, "p" + str(penalty), diffusion=diffusion)
            print("  penalty parameter:", penalty, ", error:",
                  check_vectors(filea="penalty.vector"))

if __name__ == "__main__":
    main()
    penalty()
    penalty(True)
<|endoftext|>


<commit_before>"""
==============================================
Conditional Interactions on the Snakes Dataset
==============================================
This example uses the snake dataset introduced in
Nowozin, Rother, Bagon, Sharp, Yao, Kohli: Decision Tree Fields ICCV 2011

This dataset is specifically designed to require the pairwise interaction terms
to be conditioned on the input, in other words to use non-trival edge-features.

The task is as following: a "snake" of length ten wandered over a grid. For
each cell, it had the option to go up, down, left or right (unless it came from
there). The input consists of these decisions, while the desired output is an
annotation of the snake from 0 (head) to 9 (tail).  See the plots for an
example.

As input features we use a 3x3 window around each pixel (and pad with background
where necessary). We code the five different input colors (for up, down, left, right,
background) using a one-hot encoding. This is a rather naive approach, not using any
information about the dataset (other than that it is a 2d grid).

The task can not be solved using the simple DirectionalGridCRF - which can only
infer head and tail (which are also possible to infer just from the unary
features). If we add edge-features that contain the features of the nodes that are
connected by the edge, the CRF can solve the task.

From an inference point of view, this task is very hard.  QPBO move-making is
not able to solve it alone, so we use the relaxed AD3 inference for learning.
"""
import numpy as np

from sklearn.preprocessing import label_binarize
from sklearn.metrics import confusion_matrix, accuracy_score

from pystruct.learners import OneSlackSSVM
from pystruct.datasets import load_snakes
from pystruct.utils import make_grid_edges, edge_list_to_features
from pystruct.models import EdgeFeatureGraphCRF

def one_hot_colors(x):
    x = x / 255
    flat = np.dot(x.reshape(-1, 3),  2 **  np.arange(3))
    one_hot = label_binarize(flat, classes=[1, 2, 3, 4, 6])
    return one_hot.reshape(x.shape[0], x.shape[1], 5)

def neighborhood_feature(x):
    """Add a 3x3 neighborhood around each pixel as a feature."""
    # we could also use a four neighborhood, that would work even better
    # but one might argue then we are using domain knowledge ;)
    features = np.zeros((x.shape[0], x.shape[1], 5, 9))
    # position 3 is background.
    features[:, :, 3, :] = 1
    features[1:, 1:, :, 0] = x[:-1, :-1, :]
    features[:, 1:, :, 1] = x[:, :-1, :]
    features[:-1, 1:, :, 2] = x[1:, :-1, :]
    features[1:, :, :, 3] = x[:-1, :, :]
    features[:-1, :-1, :, 4] = x[1:, 1:, :]
    features[:-1, :, :, 5] = x[1:, :, :]
    features[1:, :-1, :, 6] = x[:-1, 1:, :]
    features[:, :-1, :, 7] = x[:, 1:, :]
    features[:, :, :, 8] = x[:, :, :]
    return features.reshape(x.shape[0] * x.shape[1], -1)

def prepare_data(X):
    X_directions = []
    X_edge_features = []
    for x in X:
        # get edges in grid
        right, down = make_grid_edges(x, return_lists=True)
        edges = np.vstack([right, down])
        # use 3x3 patch around each point
        features = neighborhood_feature(x)
        # simple edge feature that encodes just if an edge is horizontal or
        # vertical
        edge_features_directions = edge_list_to_features([right, down])
        # edge feature that contains features from the nodes that the edge connects
        edge_features = np.zeros((edges.shape[0], features.shape[1], 4))
        edge_features[:len(right), :, 0] = features[right[:, 0]]
        edge_features[:len(right), :, 1] = features[right[:, 1]]
        edge_features[len(right):, :, 0] = features[down[:, 0]]
        edge_features[len(right):, :, 1] = features[down[:, 1]]
        edge_features = edge_features.reshape(edges.shape[0], -1)
        X_directions.append((features, edges, edge_features_directions))
        X_edge_features.append((features, edges, edge_features))
    return X_directions, X_edge_features

def main():
    snakes = load_snakes()
    X_train, Y_train = snakes['X_train'], snakes['Y_train']

    X_train = [one_hot_colors(x) for x in X_train]
    Y_train_flat = [y.ravel() for y in Y_train]

    X_train_directions, X_train_edge_features = prepare_data(X_train)

    # first, train on X with directions only:
    crf = EdgeFeatureGraphCRF(inference_method='qpbo')
    ssvm = OneSlackSSVM(crf, inference_cache=50, C=.1, tol=.3, switch_to='ad3',
                        n_jobs=-1)
    ssvm.fit(X_train_directions, Y_train_flat)

    # Evaluate using confusion matrix.
    # Clearly the middel of the snake is the hardest part.
    X_test, Y_test = snakes['X_test'], snakes['Y_test']
    X_test = [one_hot_colors(x) for x in X_test]
    Y_test_flat = [y.ravel() for y in Y_test]
    X_test_directions, X_test_edge_features = prepare_data(X_test)
    Y_pred = ssvm.predict(X_test_directions)
    print("Results using only directional features for edges")
    print("Test accuracy: %.3f"
          % accuracy_score(np.hstack(Y_test_flat), np.hstack(Y_pred)))
    print(confusion_matrix(np.hstack(Y_test_flat), np.hstack(Y_pred)))

    # now, use more informative edge features:
    crf = EdgeFeatureGraphCRF(inference_method='qpbo')
    ssvm = OneSlackSSVM(crf, inference_cache=50, C=.1, tol=.3, switch_to='ad3',
                        n_jobs=-1)
    ssvm.fit(X_train_edge_features, Y_train_flat)
    Y_pred = ssvm.predict(X_test_edge_features)
    print("Results using also input features for edges")
    print("Test accuracy: %.3f"
          % accuracy_score(np.hstack(Y_test_flat), np.hstack(Y_pred)))
    print(confusion_matrix(np.hstack(Y_test_flat), np.hstack(Y_pred)))

if __name__ == "__main__":
    main()
<commit_msg>EXAMPLE Add plotting code for snakes :)<commit_after>"""
==============================================
Conditional Interactions on the Snakes Dataset
==============================================
This example uses the snake dataset introduced in
Nowozin, Rother, Bagon, Sharp, Yao, Kohli: Decision Tree Fields ICCV 2011

This dataset is specifically designed to require the pairwise interaction terms
to be conditioned on the input, in other words to use non-trival edge-features.

The task is as following: a "snake" of length ten wandered over a grid. For
each cell, it had the option to go up, down, left or right (unless it came from
there). The input consists of these decisions, while the desired output is an
annotation of the snake from 0 (head) to 9 (tail).  See the plots for an
example.

As input features we use a 3x3 window around each pixel (and pad with background
where necessary). We code the five different input colors (for up, down, left, right,
background) using a one-hot encoding. This is a rather naive approach, not using any
information about the dataset (other than that it is a 2d grid).

The task can not be solved using the simple DirectionalGridCRF - which can only
infer head and tail (which are also possible to infer just from the unary
features). If we add edge-features that contain the features of the nodes that are
connected by the edge, the CRF can solve the task.

From an inference point of view, this task is very hard.  QPBO move-making is
not able to solve it alone, so we use the relaxed AD3 inference for learning.

PS: This example runs a bit (5 minutes on 12 cores, 20 minutes on one core for me).
But it does work as well as Decision Tree Fields ;)
"""
import numpy as np
import matplotlib.pyplot as plt

from sklearn.preprocessing import label_binarize
from sklearn.metrics import confusion_matrix, accuracy_score

from pystruct.learners import OneSlackSSVM
from pystruct.datasets import load_snakes
from pystruct.utils import make_grid_edges, edge_list_to_features
from pystruct.models import EdgeFeatureGraphCRF

def one_hot_colors(x):
    x = x / 255
    flat = np.dot(x.reshape(-1, 3),  2 **  np.arange(3))
    one_hot = label_binarize(flat, classes=[1, 2, 3, 4, 6])
    return one_hot.reshape(x.shape[0], x.shape[1], 5)

def neighborhood_feature(x):
    """Add a 3x3 neighborhood around each pixel as a feature."""
    # we could also use a four neighborhood, that would work even better
    # but one might argue then we are using domain knowledge ;)
    features = np.zeros((x.shape[0], x.shape[1], 5, 9))
    # position 3 is background.
    features[:, :, 3, :] = 1
    features[1:, 1:, :, 0] = x[:-1, :-1, :]
    features[:, 1:, :, 1] = x[:, :-1, :]
    features[:-1, 1:, :, 2] = x[1:, :-1, :]
    features[1:, :, :, 3] = x[:-1, :, :]
    features[:-1, :-1, :, 4] = x[1:, 1:, :]
    features[:-1, :, :, 5] = x[1:, :, :]
    features[1:, :-1, :, 6] = x[:-1, 1:, :]
    features[:, :-1, :, 7] = x[:, 1:, :]
    features[:, :, :, 8] = x[:, :, :]
    return features.reshape(x.shape[0] * x.shape[1], -1)

def prepare_data(X):
    X_directions = []
    X_edge_features = []
    for x in X:
        # get edges in grid
        right, down = make_grid_edges(x, return_lists=True)
        edges = np.vstack([right, down])
        # use 3x3 patch around each point
        features = neighborhood_feature(x)
        # simple edge feature that encodes just if an edge is horizontal or
        # vertical
        edge_features_directions = edge_list_to_features([right, down])
        # edge feature that contains features from the nodes that the edge connects
        edge_features = np.zeros((edges.shape[0], features.shape[1], 4))
        edge_features[:len(right), :, 0] = features[right[:, 0]]
        edge_features[:len(right), :, 1] = features[right[:, 1]]
        edge_features[len(right):, :, 0] = features[down[:, 0]]
        edge_features[len(right):, :, 1] = features[down[:, 1]]
        edge_features = edge_features.reshape(edges.shape[0], -1)
        X_directions.append((features, edges, edge_features_directions))
        X_edge_features.append((features, edges, edge_features))
    return X_directions, X_edge_features

def main():
    print("Please be patient. Will take 5-20 minutes.")
    snakes = load_snakes()
    X_train, Y_train = snakes['X_train'], snakes['Y_train']

    X_train = [one_hot_colors(x) for x in X_train]
    Y_train_flat = [y_.ravel() for y_ in Y_train]

    X_train_directions, X_train_edge_features = prepare_data(X_train)

    # first, train on X with directions only:
    crf = EdgeFeatureGraphCRF(inference_method='qpbo')
    ssvm = OneSlackSSVM(crf, inference_cache=50, C=.1, tol=.1, switch_to='ad3',
                        n_jobs=-1)
    ssvm.fit(X_train_directions, Y_train_flat)

    # Evaluate using confusion matrix.
    # Clearly the middel of the snake is the hardest part.
    X_test, Y_test = snakes['X_test'], snakes['Y_test']
    X_test = [one_hot_colors(x) for x in X_test]
    Y_test_flat = [y_.ravel() for y_ in Y_test]
    X_test_directions, X_test_edge_features = prepare_data(X_test)
    Y_pred = ssvm.predict(X_test_directions)
    print("Results using only directional features for edges")
    print("Test accuracy: %.3f"
          % accuracy_score(np.hstack(Y_test_flat), np.hstack(Y_pred)))
    print(confusion_matrix(np.hstack(Y_test_flat), np.hstack(Y_pred)))

    # now, use more informative edge features:
    crf = EdgeFeatureGraphCRF(inference_method='qpbo')
    ssvm = OneSlackSSVM(crf, inference_cache=50, C=.1, tol=.1, switch_to='ad3',
                        n_jobs=-1)
    ssvm.fit(X_train_edge_features, Y_train_flat)
    Y_pred2 = ssvm.predict(X_test_edge_features)
    print("Results using also input features for edges")
    print("Test accuracy: %.3f"
          % accuracy_score(np.hstack(Y_test_flat), np.hstack(Y_pred2)))
    print(confusion_matrix(np.hstack(Y_test_flat), np.hstack(Y_pred2)))

    # plot stuff
    fig, axes = plt.subplots(2, 2)
    axes[0, 0].imshow(snakes['X_test'][0], interpolation='nearest')
    axes[0, 0].set_title('Input')
    y = Y_test[0].astype(np.int)
    bg = 2 * (y != 0)  # enhance contrast
    axes[0, 1].matshow(y + bg, cmap=plt.cm.Greys)
    axes[0, 1].set_title("Ground Truth")
    axes[1, 0].matshow(Y_pred[0].reshape(y.shape) + bg, cmap=plt.cm.Greys)
    axes[1, 0].set_title("Prediction w/o edge features")
    axes[1, 1].matshow(Y_pred2[0].reshape(y.shape) + bg, cmap=plt.cm.Greys)
    axes[1, 1].set_title("Prediction with edge features")
    for a in axes.ravel():
        a.set_xticks(())
        a.set_yticks(())
    plt.show()
    from IPython.core.debugger import Tracer
    Tracer()()

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


<commit_before>from bs4 import BeautifulSoup
import csv
import urllib2

year = 2014
r  = urllib2.urlopen("http://www.hockey-reference.com/leagues/NHL_" + str(year) + "_games.html")
data = r.read()
r.close()

soup = BeautifulSoup(data)

table = soup.find('table')
rows = []

for row in table.find_all('tr'):
	rows.append([val.text.encode('utf8') for val in row.find_all('td')])

with open('../data/' + str(year) + '.csv', 'w') as f:
	writer = csv.writer(f)
	writer.writerows(row for row in rows if row)
	<commit_msg>updated permissions on scrape script<commit_after>from bs4 import BeautifulSoup
import csv
import urllib2

year = 2014
r  = urllib2.urlopen("http://www.hockey-reference.com/leagues/NHL_" + str(year) + "_games.html")
data = r.read()
r.close()

soup = BeautifulSoup(data)

table = soup.find('table')
rows = []

for row in table.find_all('tr'):
	rows.append([val.text.encode('utf8') for val in row.find_all('td')])

with open('../data/' + str(year) + '.csv', 'w') as f:
	writer = csv.writer(f)
	writer.writerows(row for row in rows if row)
	<|endoftext|>


<commit_before># Copyright (C) 2014 Rémi Bèges
# For conditions of distribution and use, see copyright notice in the LICENSE file

from enum import Enum
from pubsub import pub

class RX_STATE(Enum):
    IDLE = 0
    IN_PROCESS = 1

class ESC_STATE(Enum):
    IDLE = 0
    NEXT = 1

#Robust serial protocol with bit stuffing
# ON RX ERROR, RESET PROTOCOL ?
class SerialProtocol():

    def __init__(self):
        self.rx_state = RX_STATE.IDLE;
        self.escape_state = ESC_STATE.IDLE;
        self.SOF = int('f7',16)
        self.EOF = int('7f',16)
        self.ESC = int('7d',16)
        self.payload = bytearray()
        pub.subscribe(self.new_rx_byte,"new_rx_byte")

    def new_rx_byte(self, rxbyte):
        newbyte = int.from_bytes(rxbyte,byteorder='big')
        #No frame in process
        if self.rx_state == RX_STATE.IDLE:
            if newbyte == self.SOF:
                # New frame started
                self.rx_state = RX_STATE.IN_PROCESS
            else:
                t = newbyte,
                pub.sendMessage('new_ignored_rx_byte',rxbyte=bytes(t))

        #Frame is in process
        else:
            #Next char must be data
            if self.escape_state == ESC_STATE.NEXT:
                #Byte destuffing, this char must not be interpreted as flag
                #See serial_protocols_definition.xlsx
                self.payload.append(newbyte)
                self.escape_state = ESC_STATE.IDLE

            #Next char can be data or flag (EOF, SOF,..)
            elif self.escape_state == ESC_STATE.IDLE:
                #End of frame, the payload is immediatly send to callback function
                if newbyte == self.EOF:
                    pub.sendMessage("new_rx_payload",rxpayload=self.payload)
                    self.payload = bytearray()
                    self.rx_state = RX_STATE.IDLE

                #Escaping
                elif newbyte == self.ESC:
                    self.escape_state = ESC_STATE.NEXT
                #Storing data
                else:
                    self.payload.append(newbyte)

    def process_tx_payload(self, rxpayload):
        frame = bytearray()
        frame.append(self.SOF)

        for c in rxpayload:
            if c == self.SOF or c == self.EOF or c == self.ESC:
                frame.append(self.ESC)
            frame.append(c)

        frame.append(self.EOF)

        return frame

<commit_msg>Serial protocol can now compute frame size<commit_after># Copyright (C) 2014 Rémi Bèges
# For conditions of distribution and use, see copyright notice in the LICENSE file

from enum import Enum
from pubsub import pub

class RX_STATE(Enum):
    IDLE = 0
    IN_PROCESS = 1

class ESC_STATE(Enum):
    IDLE = 0
    NEXT = 1

#Robust serial protocol with bit stuffing
# ON RX ERROR, RESET PROTOCOL ?
class SerialProtocol():

    def __init__(self):
        self.rx_state = RX_STATE.IDLE;
        self.escape_state = ESC_STATE.IDLE;
        self.SOF = int('f7',16)
        self.EOF = int('7f',16)
        self.ESC = int('7d',16)
        self.payload = bytearray()
        pub.subscribe(self.new_rx_byte,"new_rx_byte")
        self.framesize = 0;

    def new_rx_byte(self, rxbyte):
        newbyte = int.from_bytes(rxbyte,byteorder='big')
        #No frame in process
        if self.rx_state == RX_STATE.IDLE:
            if newbyte == self.SOF:
                # New frame started
                self.rx_state = RX_STATE.IN_PROCESS
                self.framesize = 0;
            else:
                t = newbyte,
                pub.sendMessage('new_ignored_rx_byte',rxbyte=bytes(t))

        #Frame is in process
        else:
            #Next char must be data
            if self.escape_state == ESC_STATE.NEXT:
                #Byte destuffing, this char must not be interpreted as flag
                #See serial_protocols_definition.xlsx
                self.payload.append(newbyte)
                self.escape_state = ESC_STATE.IDLE
                self.framesize += 1;

            #Next char can be data or flag (EOF, SOF,..)
            elif self.escape_state == ESC_STATE.IDLE:
                #End of frame, the payload is immediatly send to callback function
                if newbyte == self.EOF:
                    pub.sendMessage("new_rx_payload",rxpayload=self.payload)
                    self.payload = bytearray()
                    self.rx_state = RX_STATE.IDLE

                #Escaping
                elif newbyte == self.ESC:
                    self.escape_state = ESC_STATE.NEXT
                #Storing data
                else:
                    self.payload.append(newbyte)
                    self.framesize += 1;

    def process_tx_payload(self, rxpayload):
        frame = bytearray()
        frame.append(self.SOF)

        for c in rxpayload:
            if c == self.SOF or c == self.EOF or c == self.ESC:
                frame.append(self.ESC)
            frame.append(c)

        frame.append(self.EOF)

        return frame

<|endoftext|>


<commit_before>from klampt import *
from klampt.model import trajectory
from klampt import robotsim
import json

_title_id = '__TITLE__'
_scene_id = '__SCENE_JSON__'
_path_id = '__PATH_JSON__'
_rpc_id = '__RPC_JSON__'
_compressed_id = '__COMPRESSED__'
_dt_id = '__TIMESTEP__'

def make_fixed_precision(obj,digits):
	if isinstance(obj,float):
		return round(obj,digits)
	elif isinstance(obj,list):
		for i in xrange(len(obj)):
			obj[i] = make_fixed_precision(obj[i],digits)
	elif isinstance(obj,tuple):
		return [make_fixed_precision(val,digits) for val in obj]
	elif isinstance(obj,dict):
		for i,v in obj.iteritems():
			obj[i] = make_fixed_precision(v,digits)
	return obj

class HTMLSharePath:
	"""An exporter that converts animations to shareable HTML files.

	Example code:
	sharer = HTMLSharePath("mypath.html",name="My spiffy path")
	sharer.start(sim)  #or world
	while [simulation is running]:
		#do whatever control you wish to do here
		sim.simulate(...)
		sharer.animate()
	sharer.end() #this saves to the given filename
	"""
	def __init__(self,filename="path.html",name="Klamp't Three.js app",boilerplate='auto'):
		"""Initializes this to save to the given filename, with 'name' as the title of the HTML page.
		boilerplate, if not 'auto', is the boilerplate HTML file.  Otherwise it it automatically found
		in the klampt folder
		"""
		self.name = name
		if boilerplate == 'auto':
			import pkg_resources
			boilerplate = pkg_resources.resource_filename('klampt','data/share_path_boilerplate.html')
		f = open(boilerplate,'r')
		self.boilerplate_file = ''.join(f.readlines())
		f.close()
		if any(v not in self.boilerplate_file for v in [_title_id,_scene_id,_path_id,_rpc_id,_compressed_id,_dt_id]):
			raise RuntimeError("Boilerplate file does not contain the right tags")
		self.fn = filename
		self.scene = []
		self.transforms = {}
		self.rpc = []
		self.dt = 0
		self.last_t = 0
	def start(self,world):
		"""Begins the path saving with the given WorldModel or Simulator"""
		if isinstance(world,Simulator):
			self.sim = world
			self.world = world.world
		else:
			self.sim = None
			self.world = world
		self.scene = robotsim.ThreeJSGetScene(self.world)
	def animate(self,time=None):
		"""Updates the path from the world.  If the world wasn't a simulator, the time
		argument needs to be provided"""
		if self.sim != None and time == None:
			time = self.sim.getTime()
			self.sim.updateWorld()
		dt = time - self.last_t
		if self.dt == 0:
			self.dt = dt
		if self.dt == 0:
			return
		if abs(dt - self.dt) <= 1e-6:
			dt = self.dt
		numadd = 0
		while dt >= self.dt:
			numadd += 1
			transforms = json.loads(robotsim.ThreeJSGetTransforms(self.world))
			for update in transforms['object']:
				n = update['name']
				mat = make_fixed_precision(update['matrix'],4)
				matpath = self.transforms.setdefault(n,[])
				assert len(matpath) == len(self.rpc)
				lastmat = None
				for m in matpath[::-1]:
					if m != None:
						lastmat = m
						break
				if lastmat != mat:
					matpath.append(mat)
				else:
					matpath.append(None)
			self.rpc.append('null')
			dt -= self.dt
			self.last_t += self.dt
		if numadd > 1:
			print "Uneven time spacing, duplicating frame",numadd,"times"
	def end(self):
		data = self.boilerplate_file.replace(_title_id,self.name)
		data = data.replace(_scene_id,self.scene)
		data = data.replace(_path_id,json.dumps(self.transforms))
		data = data.replace(_rpc_id,'['+','.join(self.rpc)+']')
		data = data.replace(_compressed_id,'true')
		data = data.replace(_dt_id,str(self.dt))
		print "Path with",len(self.rpc),"frames saved to",self.fn
		f = open(self.fn,'w')
		f.write(data)
		f.close()

if __name__ == '__main__':
	import sys
	import os
	from klampt import trajectory
	world = WorldModel()
	if len(sys.argv) == 1:
		world.readFile("../../data/athlete_plane.xml")
		q = world.robot(0).getConfig()
		q[2] = 2
		world.robot(0).setConfig(q)
		sim = Simulator(world)
		share = HTMLSharePath(name="Klamp't simulation path")
		share.start(sim)
		for i in range(100):
			sim.simulate(0.02)
			share.animate()
		share.end()
	else:
		assert len(sys.argv) == 3,"Usage: sharepath.py world.xml robot_path"
		world.readFile(sys.argv[1])
		traj = trajectory.RobotTrajectory(world.robot(0))
		traj.load(sys.argv[2])
		world.robot(0).setConfig(traj.milestones[0])

		dt = 0.02
		excess = 1.0

		share = HTMLSharePath(name="Klamp't path "+os.path.split(sys.argv[2])[1])
		share.start(world)
		share.dt = dt
		t = traj.times[0]
		while t < traj.times[-1] + excess:
			world.robot(0).setConfig(traj.eval(t))
			share.animate(t)
			t += dt
		share.end()
<commit_msg>HTML saving is started at the current simulation time<commit_after>from klampt import *
from klampt.model import trajectory
from klampt import robotsim
import json

_title_id = '__TITLE__'
_scene_id = '__SCENE_JSON__'
_path_id = '__PATH_JSON__'
_rpc_id = '__RPC_JSON__'
_compressed_id = '__COMPRESSED__'
_dt_id = '__TIMESTEP__'

def make_fixed_precision(obj,digits):
	if isinstance(obj,float):
		return round(obj,digits)
	elif isinstance(obj,list):
		for i in xrange(len(obj)):
			obj[i] = make_fixed_precision(obj[i],digits)
	elif isinstance(obj,tuple):
		return [make_fixed_precision(val,digits) for val in obj]
	elif isinstance(obj,dict):
		for i,v in obj.iteritems():
			obj[i] = make_fixed_precision(v,digits)
	return obj

class HTMLSharePath:
	"""An exporter that converts animations to shareable HTML files.

	Example code:
	sharer = HTMLSharePath("mypath.html",name="My spiffy path")
	sharer.start(sim)  #or world
	while [simulation is running]:
		#do whatever control you wish to do here
		sim.simulate(...)
		sharer.animate()
	sharer.end() #this saves to the given filename
	"""
	def __init__(self,filename="path.html",name="Klamp't Three.js app",boilerplate='auto'):
		"""Initializes this to save to the given filename, with 'name' as the title of the HTML page.
		boilerplate, if not 'auto', is the boilerplate HTML file.  Otherwise it it automatically found
		in the klampt folder
		"""
		self.name = name
		if boilerplate == 'auto':
			import pkg_resources
			boilerplate = pkg_resources.resource_filename('klampt','data/share_path_boilerplate.html')
		f = open(boilerplate,'r')
		self.boilerplate_file = ''.join(f.readlines())
		f.close()
		if any(v not in self.boilerplate_file for v in [_title_id,_scene_id,_path_id,_rpc_id,_compressed_id,_dt_id]):
			raise RuntimeError("Boilerplate file does not contain the right tags")
		self.fn = filename
		self.scene = []
		self.transforms = {}
		self.rpc = []
		self.dt = 0
		self.last_t = 0
	def start(self,world):
		"""Begins the path saving with the given WorldModel or Simulator"""
		if isinstance(world,Simulator):
			self.sim = world
			self.world = world.world
			self.last_t = world.getTime()
		else:
			self.sim = None
			self.world = world
		self.scene = robotsim.ThreeJSGetScene(self.world)
	def animate(self,time=None):
		"""Updates the path from the world.  If the world wasn't a simulator, the time
		argument needs to be provided"""
		if self.sim != None and time == None:
			time = self.sim.getTime()
			self.sim.updateWorld()
		dt = time - self.last_t
		if self.dt == 0:
			self.dt = dt
		if self.dt == 0:
			return
		if abs(dt - self.dt) <= 1e-6:
			dt = self.dt
		numadd = 0
		while dt >= self.dt:
			numadd += 1
			transforms = json.loads(robotsim.ThreeJSGetTransforms(self.world))
			for update in transforms['object']:
				n = update['name']
				mat = make_fixed_precision(update['matrix'],4)
				matpath = self.transforms.setdefault(n,[])
				assert len(matpath) == len(self.rpc)
				lastmat = None
				for m in matpath[::-1]:
					if m != None:
						lastmat = m
						break
				if lastmat != mat:
					matpath.append(mat)
				else:
					matpath.append(None)
			self.rpc.append('null')
			dt -= self.dt
			self.last_t += self.dt
		if numadd > 1:
			print "Uneven time spacing, duplicating frame",numadd,"times"
	def end(self):
		data = self.boilerplate_file.replace(_title_id,self.name)
		data = data.replace(_scene_id,self.scene)
		data = data.replace(_path_id,json.dumps(self.transforms))
		data = data.replace(_rpc_id,'['+','.join(self.rpc)+']')
		data = data.replace(_compressed_id,'true')
		data = data.replace(_dt_id,str(self.dt))
		print "Path with",len(self.rpc),"frames saved to",self.fn
		f = open(self.fn,'w')
		f.write(data)
		f.close()

if __name__ == '__main__':
	import sys
	import os
	from klampt import trajectory
	world = WorldModel()
	if len(sys.argv) == 1:
		world.readFile("../../data/athlete_plane.xml")
		q = world.robot(0).getConfig()
		q[2] = 2
		world.robot(0).setConfig(q)
		sim = Simulator(world)
		share = HTMLSharePath(name="Klamp't simulation path")
		share.start(sim)
		for i in range(100):
			sim.simulate(0.02)
			share.animate()
		share.end()
	else:
		assert len(sys.argv) == 3,"Usage: sharepath.py world.xml robot_path"
		world.readFile(sys.argv[1])
		traj = trajectory.RobotTrajectory(world.robot(0))
		traj.load(sys.argv[2])
		world.robot(0).setConfig(traj.milestones[0])

		dt = 0.02
		excess = 1.0

		share = HTMLSharePath(name="Klamp't path "+os.path.split(sys.argv[2])[1])
		share.start(world)
		share.dt = dt
		t = traj.times[0]
		while t < traj.times[-1] + excess:
			world.robot(0).setConfig(traj.eval(t))
			share.animate(t)
			t += dt
		share.end()
<|endoftext|>


<commit_before># coding: utf-8
import json
import os
import sys
import urlparse
import webbrowser
import clipboard
import dialogs

# import the funimation file
ROOT_PATH = os.path.dirname(__file__)
sys.path.append(os.path.join(ROOT_PATH, '..'))
import funimation

f = funimation.Funimation()
def dumps(dictionary):
    return json.dumps(dictionary, sort_keys=True, indent=4, separators=(',', ': '))

config_file = '../config/ios-settings.json'
if os.path.exists(config_file):
    try:
        with open(config_file, 'r') as in_file:
            opener = json.load(in_file)['opener']
    except:
        with open(config_file, 'w') as out_file:
            out_file.write(dumps({'opener': 'Safari'}))
        opener = 'Safari'
else:
    with open(config_file, 'w') as out_file:
        out_file.write(dumps({'opener': 'Safari'}))
    opener = 'Safari'

url_dict = {'Safari': 'safari-http://',
            'nPlayer': 'nplayer-http://',
            'Pythonista': 'http://',
            'Chrome': 'googlechrome://',
            'Clipboard': 'clipboard'}
url_scheme = url_dict.get(opener, 'http://')
openwith = [{'urlscheme': value, 'title': key}
           for key, value in url_dict.items()]

def showpicker():
    return [{'title': 'Opener'}] + [{'title': show.series_name,
             'asset_id': show.asset_id} for show in f.get_shows()]

def videos_list(item_list):
    return [{'title': '{} : {}-{}'.format(parts.info['episode'], parts.title,
              parts.dub_sub), 'url': parts.video_url} for parts in item_list]

choice = dialogs.list_dialog('shows', showpicker()) or sys.exit('Quit')
if choice['title'] is 'Opener':
    choice = dialogs.list_dialog('Open With', openwith)
    with open('ios-settings.json', 'w') as out_file:
        out_file.write(dumps({'opener': choice['title']}))
    sys.exit('Settings changed')
vtable = f.get_videos(choice['asset_id'])
videos = videos_list(vtable)
item = dialogs.list_dialog(choice['title'], videos) or sys.exit('Quit')
url = urlparse.urlparse(item['url'])
if urlscheme is 'clipboard':
    clipboard.set(item['url'])
else:
    open_url = urlscheme + url.hostname + url.path + '?' + url.query
    webbrowser.open(open_url)
<commit_msg>Update pythonista.py<commit_after># coding: utf-8
import json
import os
import sys
import urlparse
import webbrowser
import clipboard
import dialogs

# import the funimation file
ROOT_PATH = os.path.dirname(__file__)
sys.path.append(os.path.join(ROOT_PATH, '..'))
import funimation

f = funimation.Funimation()
def dumps(dictionary):
    return json.dumps(dictionary, sort_keys=True, indent=4, separators=(',', ': '))

config_file = '../config/ios-settings.json'
if os.path.exists(config_file):
    try:
        with open(config_file, 'r') as in_file:
            opener = json.load(in_file)['opener']
    except:
        with open(config_file, 'w') as out_file:
            out_file.write(dumps({'opener': 'Safari'}))
        opener = 'Safari'
else:
    with open(config_file, 'w') as out_file:
        out_file.write(dumps({'opener': 'Safari'}))
    opener = 'Safari'

url_dict = {'Safari': 'safari-http://',
            'nPlayer': 'nplayer-http://',
            'Pythonista': 'http://',
            'Chrome': 'googlechrome://',
            'Clipboard': 'clipboard'}
url_scheme = url_dict.get(opener, 'http://')
openwith = [{'urlscheme': value, 'title': key} for key, value in url_dict.items()]

def showpicker():
    return [{'title': 'Opener'}] + [{'title': show.series_name,
             'asset_id': show.asset_id} for show in f.get_shows()]

def videos_list(item_list):
    return [{'title': '{} : {}-{}'.format(parts.info['episode'], parts.title,
              parts.dub_sub), 'url': parts.video_url} for parts in item_list]

choice = dialogs.list_dialog('shows', showpicker()) or sys.exit('Quit')
if choice['title'] is 'Opener':
    choice = dialogs.list_dialog('Open With', openwith)
    with open('ios-settings.json', 'w') as out_file:
        out_file.write(dumps({'opener': choice['title']}))
    sys.exit('Settings changed')
vtable = f.get_videos(choice['asset_id'])
videos = videos_list(vtable)
item = dialogs.list_dialog(choice['title'], videos) or sys.exit('Quit')
url = urlparse.urlparse(item['url'])
if urlscheme is 'clipboard':
    clipboard.set(item['url'])
else:
    open_url = urlscheme + url.hostname + url.path + '?' + url.query
    webbrowser.open(open_url)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

###############################################################################
#  Copyright Kitware Inc.
#
#  Licensed under the Apache License, Version 2.0 ( the "License" );
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
###############################################################################

import json
import os
import sys
import time
#import zipfile

#import geojson

# Need to set the environment variable before importing girder
girder_port = os.environ.get('GIRDER_TEST_PORT', '20200')
os.environ['GIRDER_PORT'] = girder_port# noqa

from tests import base

from girder_client import GirderClient

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../utility')))
import import_analyses
import romanesco

def setUpModule():
    """
    Enable the minerva plugin and start the server.
    """
    base.enabledPlugins.append('jobs')
    base.enabledPlugins.append('romanesco')
    base.enabledPlugins.append('gravatar')
    base.enabledPlugins.append('minerva')
    base.startServer(False)

def tearDownModule():
    """
    Stop the server.
    """
    base.stopServer()

class MeanContourAnalysisTestCase(base.TestCase):
    """
    Tests the server side of the mean_contour_analysis.
    """

    def setUp(self):
        """
        Set up the test case with  a user
        """
        super(MeanContourAnalysisTestCase, self).setUp()

        self._user = self.model('user').createUser(
            'minervauser', 'password', 'minerva', 'user',
            'minervauser@example.com')

        # create a dataset folder for the user

        path = '/minerva_dataset/folder'
        params = {
            'userId': self._user['_id'],
        }
        response = self.request(path=path, method='POST', params=params, user=self._user)
        self.assertStatusOk(response)
        self._datasetFolder = response.json['folder']

        # create the analysis folder

        path = '/minerva_analysis/folder'
        response = self.request(path=path, method='POST', user=self._user)
        self.assertStatusOk(response)
        self.assertNotEquals(response.json['folder'], None, 'An analysis folder should exist')
        self._analysisFolder = response.json['folder']

        # import the mean_contour_analysis
        client = GirderClient('localhost', girder_port)
        client.authenticate('minervauser', 'password')

        analysis_path = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../analyses/NEX'))
        import_analyses.import_analyses(client, analysis_path)

        path = '/item'
        params = {
            'folderId': self._analysisFolder['_id']
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        self.assertStatusOk(response)
        self._meanContourAnalysis = None
        for analysis in response.json:
            if analysis['name'] == 'mean contour':
                self._meanContourAnalysis = analysis
        self.assertNotEquals(self._meanContourAnalysis, None, 'mean contour analysis not found')

        # give this user the ability to run romanesco jobs
        path = '/system/setting'
        params = {
                'list': json.dumps([
                    {"key": "romanesco.broker", "value": ""},
                    {"key": "romanesco.backend", "value": ""},
                    {"key": "romanesco.full_access_users", "value": json.dumps(['minervauser'])},
                    {"key": "romanesco.full_access_groups", "value": "[]"},
                    {"key": "romanesco.safe_folders", "value": "[]"},
                    {"key": "romanesco.require_auth", "value": True}])
        }
        response = self.request(path=path, method='PUT', params=params, user=self._user)
        self.assertTrue(response.json)

    def testAnalysis(self):
        """
        Test the mean_contour_analysis workflow by issuing the API
        calls the inerva dataset API enppoints.
        """

        # create the S3 Nasanex dataset item
        path = '/item'
        params = {
            'userId': self._user['_id'],
            'folderId': self._datasetFolder['_id'],
            'name': 'nasanex'
        }
        response = self.request(path=path, method='POST', params=params, user=self._user)
        self.assertStatusOk(response)
        nasanexItem = response.json

        # create an s3 dataset
        path = '/minerva_dataset_s3/{}/dataset'.format(nasanexItem['_id'])
        params = {
            'bucket': 'nasanex',
            'name': 'nasanex',
            'prefix': '/CMIP5/CommonGrid/gfdl-esm2g/rcp45/mon/r1i1p1/pr/',
            'readOnly': True
        }
        response = self.request(path=path, method='POST', params=params, user=self._user)
        self.assertStatusOk(response)
        nasanexMinervaMetadata = response.json

        # get the job tied in with dataset creation
        path = '/job'
        response = self.request(path=path, method='GET', params=params, user=self._user)
        self.assertStatusOk(response)
        # first job should have type='s3.import'
        job = response.json[0]
        self.assertEquals(job['type'], 's3.import')

        # import jobs after server has started and jobs pseudo-module exists
        from girder.plugins.jobs.constants import JobStatus
        #wait for the async job to complete
        jobFinished = int(job['status']) == JobStatus.SUCCESS
        count = 0
        while not jobFinished and count < 5:
            time.sleep(10)
            count += 1
            # get the job and check if it has been updated
            path = '/job/{}'.format(job['_id'])
            response = self.request(
                path=path,
                method='GET',
                user=self._user
            )
            self.assertStatusOk(response)
            job = response.json
            jobFinished = int(job['status']) == JobStatus.SUCCESS
        self.assertTrue(jobFinished, 'Dataset job never finished')

        # get the item now imported into the S3 dataset folder
        s3FolderId = nasanexMinervaMetadata['folderId']
        path = '/item'
        params = {
            'folderId': s3FolderId
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        self.assertStatusOk(response)
        s3ImportedItem = response.json[0]

        # select this imported item in the S3 dataset
        path = '/minerva_dataset_s3/{}/dataset'.format(nasanexItem['_id'])
        nasanexMinervaMetadata['selectedItems'] = '["'+str(s3ImportedItem['_id'])+'"]'
        response = self.request(path=path, method='PUT', params=nasanexMinervaMetadata, user=self._user)
        self.assertStatusOk(response)
        nasanexMinervaMetadata = response.json

        # get the file id of the file in the imported item
        path = '/item/{}/files'.format(s3ImportedItem['_id'])
        params = {
            'limit': 1
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        fileId = response.json[0]['_id']

        # Run the analysis in romanseco, calling romanesco.run directly so that
        # the celery/romanesco worker doesn't need to be running.

        # Requires SPARK_HOME to be set in the running ENV

        parameter = 'pr'
        data = {
            'inputs': {
                'host': {
                    'format': 'json',
                    'data': 'localhost'
                },
                'port': {
                    'format': 'json',
                    'data': girder_port
                },
                'token': {
                    'format': 'json',
                    'data': self._genToken(self._user)
                },
                'fileId': {
                    'format': 'json',
                    'name': 'fileId',
                    'data': fileId
                },
                'variable': {
                    'format': 'json',
                    'data': parameter
                }

            },
            'outputs': {
                'result': {
                    'format': 'json'
                }
            }
        }

        romAnalysis = self._meanContourAnalysis['meta']['analysis']
        outputs = romanesco.run(romAnalysis, data['inputs'], data['outputs'])
        outputItemId = outputs['output_item_id']['data']

        # check the sha512 of the result file
        path = '/item/{}/files'.format(outputItemId)
        params = {
            'limit': 1
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        uploadedSha512 = response.json[0]['sha512']
        expectedSha512 = 'af37c2e693b1b9f45dac8d3baff2b7017403be1455e572ce6aa2d324cc526444f7d50b43825740707f1b2e8d6150029eb65ac23e3e979ca56c5d4e2f5182daf1'
        self.assertEquals(uploadedSha512, expectedSha512, 'mean_contour_analysis output file has unexpected sha512')
<commit_msg>Remove dependence of test on sha512 key in file<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

###############################################################################
#  Copyright Kitware Inc.
#
#  Licensed under the Apache License, Version 2.0 ( the "License" );
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
###############################################################################

import json
import os
import sys
import time
#import zipfile

#import geojson

# Need to set the environment variable before importing girder
girder_port = os.environ.get('GIRDER_TEST_PORT', '20200')
os.environ['GIRDER_PORT'] = girder_port# noqa

from tests import base

from girder_client import GirderClient

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../utility')))
import import_analyses
import romanesco

def setUpModule():
    """
    Enable the minerva plugin and start the server.
    """
    base.enabledPlugins.append('jobs')
    base.enabledPlugins.append('romanesco')
    base.enabledPlugins.append('gravatar')
    base.enabledPlugins.append('minerva')
    base.startServer(False)

def tearDownModule():
    """
    Stop the server.
    """
    base.stopServer()

class MeanContourAnalysisTestCase(base.TestCase):
    """
    Tests the server side of the mean_contour_analysis.
    """

    def setUp(self):
        """
        Set up the test case with  a user
        """
        super(MeanContourAnalysisTestCase, self).setUp()

        self._user = self.model('user').createUser(
            'minervauser', 'password', 'minerva', 'user',
            'minervauser@example.com')

        # create a dataset folder for the user

        path = '/minerva_dataset/folder'
        params = {
            'userId': self._user['_id'],
        }
        response = self.request(path=path, method='POST', params=params, user=self._user)
        self.assertStatusOk(response)
        self._datasetFolder = response.json['folder']

        # create the analysis folder

        path = '/minerva_analysis/folder'
        response = self.request(path=path, method='POST', user=self._user)
        self.assertStatusOk(response)
        self.assertNotEquals(response.json['folder'], None, 'An analysis folder should exist')
        self._analysisFolder = response.json['folder']

        # import the mean_contour_analysis
        client = GirderClient('localhost', girder_port)
        client.authenticate('minervauser', 'password')

        analysis_path = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../analyses/NEX'))
        import_analyses.import_analyses(client, analysis_path)

        path = '/item'
        params = {
            'folderId': self._analysisFolder['_id']
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        self.assertStatusOk(response)
        self._meanContourAnalysis = None
        for analysis in response.json:
            if analysis['name'] == 'mean contour':
                self._meanContourAnalysis = analysis
        self.assertNotEquals(self._meanContourAnalysis, None, 'mean contour analysis not found')

        # give this user the ability to run romanesco jobs
        path = '/system/setting'
        params = {
                'list': json.dumps([
                    {"key": "romanesco.broker", "value": ""},
                    {"key": "romanesco.backend", "value": ""},
                    {"key": "romanesco.full_access_users", "value": json.dumps(['minervauser'])},
                    {"key": "romanesco.full_access_groups", "value": "[]"},
                    {"key": "romanesco.safe_folders", "value": "[]"},
                    {"key": "romanesco.require_auth", "value": True}])
        }
        response = self.request(path=path, method='PUT', params=params, user=self._user)
        self.assertTrue(response.json)

    def testAnalysis(self):
        """
        Test the mean_contour_analysis workflow by issuing the API
        calls the inerva dataset API enppoints.
        """

        # create the S3 Nasanex dataset item
        path = '/item'
        params = {
            'userId': self._user['_id'],
            'folderId': self._datasetFolder['_id'],
            'name': 'nasanex'
        }
        response = self.request(path=path, method='POST', params=params, user=self._user)
        self.assertStatusOk(response)
        nasanexItem = response.json

        # create an s3 dataset
        path = '/minerva_dataset_s3/{}/dataset'.format(nasanexItem['_id'])
        params = {
            'bucket': 'nasanex',
            'name': 'nasanex',
            'prefix': '/CMIP5/CommonGrid/gfdl-esm2g/rcp45/mon/r1i1p1/pr/',
            'readOnly': True
        }
        response = self.request(path=path, method='POST', params=params, user=self._user)
        self.assertStatusOk(response)
        nasanexMinervaMetadata = response.json

        # get the job tied in with dataset creation
        path = '/job'
        response = self.request(path=path, method='GET', params=params, user=self._user)
        self.assertStatusOk(response)
        # first job should have type='s3.import'
        job = response.json[0]
        self.assertEquals(job['type'], 's3.import')

        # import jobs after server has started and jobs pseudo-module exists
        from girder.plugins.jobs.constants import JobStatus
        #wait for the async job to complete
        jobFinished = int(job['status']) == JobStatus.SUCCESS
        count = 0
        while not jobFinished and count < 5:
            time.sleep(10)
            count += 1
            # get the job and check if it has been updated
            path = '/job/{}'.format(job['_id'])
            response = self.request(
                path=path,
                method='GET',
                user=self._user
            )
            self.assertStatusOk(response)
            job = response.json
            jobFinished = int(job['status']) == JobStatus.SUCCESS
        self.assertTrue(jobFinished, 'Dataset job never finished')

        # get the item now imported into the S3 dataset folder
        s3FolderId = nasanexMinervaMetadata['folderId']
        path = '/item'
        params = {
            'folderId': s3FolderId
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        self.assertStatusOk(response)
        s3ImportedItem = response.json[0]

        # select this imported item in the S3 dataset
        path = '/minerva_dataset_s3/{}/dataset'.format(nasanexItem['_id'])
        nasanexMinervaMetadata['selectedItems'] = '["'+str(s3ImportedItem['_id'])+'"]'
        response = self.request(path=path, method='PUT', params=nasanexMinervaMetadata, user=self._user)
        self.assertStatusOk(response)
        nasanexMinervaMetadata = response.json

        # get the file id of the file in the imported item
        path = '/item/{}/files'.format(s3ImportedItem['_id'])
        params = {
            'limit': 1
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        fileId = response.json[0]['_id']

        # Run the analysis in romanseco, calling romanesco.run directly so that
        # the celery/romanesco worker doesn't need to be running.

        # Requires SPARK_HOME to be set in the running ENV

        parameter = 'pr'
        data = {
            'inputs': {
                'host': {
                    'format': 'json',
                    'data': 'localhost'
                },
                'port': {
                    'format': 'json',
                    'data': girder_port
                },
                'token': {
                    'format': 'json',
                    'data': self._genToken(self._user)
                },
                'fileId': {
                    'format': 'json',
                    'name': 'fileId',
                    'data': fileId
                },
                'variable': {
                    'format': 'json',
                    'data': parameter
                }

            },
            'outputs': {
                'result': {
                    'format': 'json'
                }
            }
        }

        romAnalysis = self._meanContourAnalysis['meta']['analysis']
        outputs = romanesco.run(romAnalysis, data['inputs'], data['outputs'])
        outputItemId = outputs['output_item_id']['data']

        # check the size of the result file
        path = '/item/{}/files'.format(outputItemId)
        params = {
            'limit': 1
        }
        response = self.request(path=path, method='GET', params=params, user=self._user)
        uploadedSize = response.json[0]['size']
        expectedSize = 1265241
        self.assertEquals(uploadedSize, expectedSize, 'mean_contour_analysis output file has unexpected size')
<|endoftext|>


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

__author__     = "Jan Kogut"
__copyright__  = "Jan Kogut"
__license__    = "MIT"
__version__    = "0.0.1"
__maintainer__ = "Jan Kogut"
__status__     = "Beta"

import json
from optparse import OptionParser
from kazoo.client import KazooClient

## START of config section
##################################################

class OurConfig(object):
    ''' Our config section '''

    pass

cfg = OurConfig()

cfg.zkServers  = 'zoo1.dmz:2181,zoo2.dmz:2181,zoo3.dmz:2181'
cfg.aPath      = '/ansible-test'

#################################################
## END of config section

def oParser():
    '''
    Commandline options parsing function.

    Return dict (parsed options).
    '''

    parser = OptionParser(usage="usage: %prog [opts] <args>",
                          version="%prog 0.0.1")
    parser.add_option("-A", nargs = 1,
                      help="add hostname with hostvars: <groupname1:newhostname1,var1:value1,var2:value2,var3:value3>")
    parser.add_option("-G", nargs = 1,
                      help="add hostname to hostgroup: <groupname:hostname>")
    parser.add_option("-D", nargs = 1,
                      help="delete hostname or groupname recursively: <groupname1:newhostname1> or <groupname1>")
    parser.add_option("-U", nargs = 1,
                      help="update hostname with comma separated hostvars: <groupname1:newhostname1,var1:newvalue1,var2:newvalue2>")
    parser.add_option("-S", nargs = 1,
                      help="show hostvars for a given hostname or groupname: <groupname1:newhostname1> or <groupname1>")
    parser.add_option("-I", nargs = 1,
                      help="inventory mode: true|ansible dumps inventory in json format from zookeeper")

    (opts, args) = parser.parse_args()

    if (opts.A or opts.G or opts.D or opts.U or opts.I or opts.S) == None:

        parser.print_help()
        exit(-1)

    return {'addMode':opts.A, 'groupMode':opts.G, 'deleteMode':opts.D, 'updateMode':opts.U,
            'showMode':opts.S, 'inventoryMode':opts.I}

def zkStartRo():
    '''
    Start a zookeeper client connection in read-only mode.

    Return zookeeper read-only connection object.
    '''

    zk = KazooClient(hosts=cfg.zkServers, read_only = True)
    zk.start()

    return zk

def zkStartRw():
    '''
    Start a zookeeper client connection in read-write mode.

    Return zookeeper read-write connection object.
    '''

    zk = KazooClient(hosts=cfg.zkServers)
    zk.start()

    return zk

def splitZnodeVarString(znodeVarString):
    '''
    Parse string for commandline opts: <-A|-U>.

    Return dict.
    '''

    ## spliting example string into dictionary:
    ## example string: groupname:hostname1,var1:val1,var2:val2,var3:val3
    ## desired dict  : {"groupname":{"hostname1":{"var1":"val1", "var2":"val2", "var3":"val3"}}}

    varList = znodeVarString.split(',')
    varDict = {}

    for var in varList[1:]:
        varDict[var.split(':')[0]] = var.split(':')[1]

    groupName, hostName = varList[0].split(':')[0], varList[0].split(':')[1]

    return { groupName : { hostName : varDict }}

def splitZnodeString(znodeString):
    '''
    Splits znodeString into groupName, hostName, groupPath, hostPath.

    Return list of tuples.
    '''

    ## spliting example string into list of tuples or tuple :
    ## example string: groupname:hostname1
    ## example output: [("groupname","/ansible_zk/groups/groupname"),("hostname1","/ansible_zk/groups/groupname/hostname1")]

    if ':' in znodeString:
        groupName = znodeString.split(':')[0]
        hostName  = znodeString.split(':')[1]
        groupPath = "{0}/groups/{1}".format(cfg.aPath, groupName)
        hostPath  = "{0}/{1}".format(groupPath, hostName)

        return [(groupName, groupPath),(hostName, hostPath)]

    else:
        groupName = znodeString
        groupPath = "{0}/groups/{1}".format(cfg.aPath, groupName)

        return [(groupName, groupPath)]

def addHostWithHostvars(znodeDict):
    '''
    Add existing znode to new group.

    Return a string (ADDED    ==> host: <hostname> to group: <groupname>).
    '''

    zk = zkStartRw()

    groupName = znodeDict.keys()[0]
    hostName  = znodeDict[groupName].keys()[0]
    groupPath = "{0}/groups/{1}".format(cfg.aPath, groupName)
    hostPath  = "{0}/{1}".format(groupPath, hostName)

    if zk.exists(hostPath):
        zk.stop()
        return "ERROR  ==> host: {0} in group {1} exist !!!".format(hostName, groupName)

    if zk.exists(groupPath):
        zk.create(hostPath)
    else:
        zk.ensure_path(hostPath)

    for key in znodeDict[groupName][hostName]:
        varPath = "{0}/{1}".format(hostPath, key)
        varVal  = znodeDict[groupName][hostName][key]
        zk.create(varPath, varVal)

    zk.stop()

    return "ADDED   ==> host: {0} to group: {1}".format(hostName, groupName)

def addHostToGroup(znodeStringSplited):
    '''
    Add new znode with hostvars.

    Return a string (ADDED    ==> host: <hostname> to group: <groupname>).
    '''

    zk = zkStartRw()

    groupName = znodeStringSplited[0][0]
    groupPath = znodeStringSplited[0][1]
    hostName  = znodeStringSplited[1][0]
    hostPath  = znodeStringSplited[1][1]

    if zk.exists(hostPath):
        zk.stop()
        return "ERROR  ==> host: {0} in group {1} exist !!!".format(hostName, groupName)

    if zk.exists(groupPath):
        zk.create(hostPath)
    else:
        zk.ensure_path(hostPath)

    zk.stop()

    return "ADDED   ==> host: {0} to group: {1}".format(hostName, groupName)

def deleteZnodeRecur(znodeStringSplited):
    '''
    Delete znode with hostvars for a given list of tuple <groupname:hostname>.

    Return a string (DELETED||ERROR  ==> host: <hostname> in group: <groupname>||group: <groupname>).
    '''

    zk = zkStartRw()

    if len(znodeStringSplited) > 1:

        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]
        hostName  = znodeStringSplited[1][0]
        hostPath  = znodeStringSplited[1][1]

        if zk.exists(hostPath) is None:
            zk.stop()
            return "ERROR  ==> could not delete host: {0} that does not exist !!!".format(hostName)

        if len(zk.get_children(groupPath)) == 1:
            zk.delete(groupPath, recursive=True)
        else:
            zk.delete(hostPath, recursive=True)
            zk.stop()
            return "DELETED ==> host: {0} in group: {1}".format(hostName, groupName)

    elif len(znodeStringSplited) == 1:

        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]

        if zk.exists(groupPath) is None:
            zk.stop()
            return "ERROR  ==> could not delete group: {0} that does not exist !!!".format(groupName)

        else:
            zk.delete(groupPath, recursive=True)
    else:
        return "ERROR with processing znodeStrings !!!"

    zk.stop()
    return "DELETED ==> group: {0}".format(groupName)

# def updateZnode(znodeDict):
#     '''
#     Update znode with hostvars.

#     Return a string (UPDATED   ==> host: <hostname> to group: <groupname>).
#     '''

#     zk = zkStartRw()

#     groupName   = znodeDict.keys()[0]
#     hostName    = znodeDict[groupName].keys()[0]
#     groupPath   = "{0}/groups/{1}".format(cfg.aPath, groupName)
#     hostPath    = "{0}/{1}".format(groupPath, hostName)
#     hostVarList = zk.get_children(hostPath)

#     if zk.exists(hostPath) is None:
#         zk.stop()
#         return "ERROR  ==> host: {0} in group {1} does not exist !!!".format(hostName, groupName)

#     for hostVar in hostVarList:
#         if zk.exists("{0}/{1}".format(hostPath, hostVar)) is None:
#             zk.stop()
#             return "ERROR  ==> hostvar: {0} for host {1} in group {2} does not exist !!!".format(hostVar, hostName, groupName)

#     for key in znodeDict[groupName][hostName]:
#         varPath = "{0}/{1}".format(hostPath, key)
#         varVal  = znodeDict[groupName][hostName][key]
#         if CONDITION: #FIXME!!!
#            zk.create(varPath, varVal)

#     zk.stop()
#     return "ADDED   ==> host: {0} to group: {1}".format(hostName, groupName)

def hostVarsShow(znodeStringSplited):
    '''
    Show hostvars for a given <groupname:hostname> or <groupname>.

    Return dict or string (in case of ERROR).
    '''

    zk = zkStartRo()

    if len(znodeStringSplited) > 1:

        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]
        hostName  = znodeStringSplited[1][0]
        hostPath  = znodeStringSplited[1][1]

        if zk.exists(hostPath) is None:
            zk.stop()
            return "ERROR  ==> no such hostname: {0} in group {1} !!!".format(hostName, groupName)

        else:
            hostVarList = zk.get_children(hostPath)
            valDict     = {}

            for var in hostVarList:
                valDict[var] = zk.get('{0}/{1}'.format(hostPath, var))[0]

    elif len(znodeStringSplited) == 1:
        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]

        if zk.exists(groupPath) is None:
            zk.stop()
            return "ERROR  ==> no such groupname: {0} !!!".format(groupName)

        else:
            hostList = zk.get_children(groupPath)
            valDict  = {}

            for host in hostList:
                tmpHostPath    = '{0}/{1}'.format(groupPath, host)
                valDict[host]  = zk.get_children('{0}'.format(tmpHostPath))

            for host in valDict.keys():
                varDict = {}

                for var in valDict[host]:
                    tmpHostPath   = '{0}/{1}'.format(groupPath, host)
                    varDict[var]  = zk.get('{0}/{1}'.format(tmpHostPath, var))[0]
                    valDict[host] = varDict

            return valDict

    else:
        return "ERROR with processing znodeStrings !!!"

    zk.stop()
    return valDict

def inventoryDump():
    '''
    Inventory dump for a given list of zookeeper servers and ansible-keeper path.

    Return dict.
    '''

    zk = zkStartRo()

    groupList = zk.get_children("{}/groups".format(cfg.aPath))
    groupDict = {}

    for group in groupList:
        path     = "{0}/groups/{1}".format(cfg.aPath, group)
        children = zk.get_children(path)
        groupDict[group] = children

    zk.stop()
    return groupDict

def ansibleInventoryDump():
    '''
    Ansible compliant inventory dump for a given list of zookeeper servers and ansible-keeper path.

    Return dict.
    '''

    zk = zkStartRo()

    groupList = zk.get_children("{}/groups".format(cfg.aPath))
    groupDict = {}

    for group in groupList:
        path     = "{0}/groups/{1}".format(cfg.aPath, group)
        children = zk.get_children(path)
        tmpDict  = {}
        tmpDict['hosts'] = children
        tmpDict['vars']  = {"a":"b"}
        groupDict[group] = tmpDict

    zk.stop()
    return groupDict

def main():
    '''
    Main logic
    '''

    if oParser()['inventoryMode'] == 'true':
        print json.dumps(inventoryDump())

    if oParser()['inventoryMode'] == 'ansible':
        print json.dumps(ansibleInventoryDump())

    if oParser()['addMode'] is not None:
        znodeDict = splitZnodeVarString(oParser()['addMode'])
        print addHostWithHostvars(znodeDict)

    if oParser()['groupMode'] is not None:
        znodeStringSplited = splitZnodeString(oParser()['groupMode'])
        print addHostToGroup(znodeStringSplited)

    if oParser()['deleteMode'] is not None:
        znodeStringSplited = splitZnodeString(oParser()['deleteMode'])
        print deleteZnodeRecur(znodeStringSplited)

    if oParser()['showMode'] is not None:
        znodeStringSplited = splitZnodeString(oParser()['showMode'])
        print json.dumps(hostVarsShow(znodeStringSplited))

if __name__ == "__main__":
    main()
<commit_msg>addHostWithHostvars: moved hostvars to hosts/hostname1 from groups/groupname1/hostname1<commit_after>#!/usr/bin/env python

__author__     = "Jan Kogut"
__copyright__  = "Jan Kogut"
__license__    = "MIT"
__version__    = "0.0.1"
__maintainer__ = "Jan Kogut"
__status__     = "Beta"

import json
from optparse import OptionParser
from kazoo.client import KazooClient

## START of config section
##################################################

class OurConfig(object):
    ''' Our config section '''

    pass

cfg = OurConfig()

cfg.zkServers  = 'zoo1.dmz:2181,zoo2.dmz:2181,zoo3.dmz:2181'
cfg.aPath      = '/ansible-test'

#################################################
## END of config section

def oParser():
    '''
    Commandline options parsing function.

    Return dict (parsed options).
    '''

    parser = OptionParser(usage="usage: %prog [opts] <args>",
                          version="%prog 0.0.1")
    parser.add_option("-A", nargs = 1,
                      help="add hostname with hostvars: <groupname1:newhostname1,var1:value1,var2:value2,var3:value3>")
    parser.add_option("-G", nargs = 1,
                      help="add hostname to hostgroup: <groupname:hostname>")
    parser.add_option("-D", nargs = 1,
                      help="delete hostname or groupname recursively: <groupname1:newhostname1> or <groupname1>")
    parser.add_option("-U", nargs = 1,
                      help="update hostname with comma separated hostvars: <groupname1:newhostname1,var1:newvalue1,var2:newvalue2>")
    parser.add_option("-S", nargs = 1,
                      help="show hostvars for a given hostname or groupname: <groupname1:newhostname1> or <groupname1>")
    parser.add_option("-I", nargs = 1,
                      help="inventory mode: true|ansible dumps inventory in json format from zookeeper")

    (opts, args) = parser.parse_args()

    if (opts.A or opts.G or opts.D or opts.U or opts.I or opts.S) == None:

        parser.print_help()
        exit(-1)

    return {'addMode':opts.A, 'groupMode':opts.G, 'deleteMode':opts.D, 'updateMode':opts.U,
            'showMode':opts.S, 'inventoryMode':opts.I}

def zkStartRo():
    '''
    Start a zookeeper client connection in read-only mode.

    Return zookeeper read-only connection object.
    '''

    zk = KazooClient(hosts=cfg.zkServers, read_only = True)
    zk.start()

    return zk

def zkStartRw():
    '''
    Start a zookeeper client connection in read-write mode.

    Return zookeeper read-write connection object.
    '''

    zk = KazooClient(hosts=cfg.zkServers)
    zk.start()

    return zk

def splitZnodeVarString(znodeVarString):
    '''
    Parse string for commandline opts: <-A|-U>.

    Return dict.
    '''

    ## spliting example string into dictionary:
    ## example string: groupname:hostname1,var1:val1,var2:val2,var3:val3
    ## desired dict  : {"groupname":{"hostname1":{"var1":"val1", "var2":"val2", "var3":"val3"}}}

    varList = znodeVarString.split(',')
    varDict = {}

    for var in varList[1:]:
        varDict[var.split(':')[0]] = var.split(':')[1]

    groupName, hostName = varList[0].split(':')[0], varList[0].split(':')[1]

    return { groupName : { hostName : varDict }}

def splitZnodeString(znodeString):
    '''
    Splits znodeString into groupName, hostName, groupPath, hostPath.

    Return list of tuples.
    '''

    ## spliting example string into list of tuples or tuple :
    ## example string: groupname:hostname1
    ## example output: [("groupname","/ansible_zk/groups/groupname"),("hostname1","/ansible_zk/groups/groupname/hostname1")]

    if ':' in znodeString:
        groupName = znodeString.split(':')[0]
        hostName  = znodeString.split(':')[1]
        groupPath = "{0}/groups/{1}".format(cfg.aPath, groupName)
        hostPath  = "{0}/{1}".format(groupPath, hostName)

        return [(groupName, groupPath),(hostName, hostPath)]

    else:
        groupName = znodeString
        groupPath = "{0}/groups/{1}".format(cfg.aPath, groupName)

        return [(groupName, groupPath)]

def addHostWithHostvars(znodeDict):
    '''
    Add existing znode to new group.

    Return a string (ADDED    ==> host: <hostname> to group: <groupname>).
    '''

    zk = zkStartRw()

    groupName      = znodeDict.keys()[0]
    hostName       = znodeDict[groupName].keys()[0]
    groupPath      = "{0}/groups/{1}".format(cfg.aPath, groupName)
    hostPath       = "{0}/hosts/{1}".format(cfg.aPath, hostName)
    hostGroupPath  = "{0}/{1}".format(groupPath, hostName)

    if zk.exists(hostPath):
        zk.stop()
        return "ERROR  ==> host: {0} exists !!!".format(hostName, groupName)

    elif zk.exists(hostGroupPath):
        zk.stop()
        return "ERROR  ==> host: {0} in group {1} exists !!!".format(hostName, groupName)

    else:
        zk.ensure_path(hostPath)
        zk.ensure_path(hostGroupPath)

        for key in znodeDict[groupName][hostName]:
            varPath = "{0}/{1}".format(hostPath, key)
            varVal  = znodeDict[groupName][hostName][key]
            zk.create(varPath, varVal)

        zk.stop()

        return "ADDED   ==> host: {0} to group: {1}".format(hostName, groupName)

def addHostToGroup(znodeStringSplited):
    '''
    Add host to group.

    Return a string (ADDED    ==> host: <hostname> to group: <groupname>).
    '''

    zk = zkStartRw()

    groupName = znodeStringSplited[0][0]
    groupPath = znodeStringSplited[0][1]
    hostName  = znodeStringSplited[1][0]
    hostPath  = znodeStringSplited[1][1]

    if zk.exists(hostPath):
        zk.stop()
        return "ERROR  ==> host: {0} in group {1} exist !!!".format(hostName, groupName)

    if zk.exists(groupPath):
        zk.create(hostPath)
    else:
        zk.ensure_path(hostPath)

    zk.stop()

    return "ADDED   ==> host: {0} to group: {1}".format(hostName, groupName)

def deleteZnodeRecur(znodeStringSplited):
    '''
    Delete znode with hostvars for a given list of tuple <groupname:hostname>.

    Return a string (DELETED||ERROR  ==> host: <hostname> in group: <groupname>||group: <groupname>).
    '''

    zk = zkStartRw()

    if len(znodeStringSplited) > 1:

        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]
        hostName  = znodeStringSplited[1][0]
        hostPath  = znodeStringSplited[1][1]

        if zk.exists(hostPath) is None:
            zk.stop()
            return "ERROR  ==> could not delete host: {0} that does not exist !!!".format(hostName)

        if len(zk.get_children(groupPath)) == 1:
            zk.delete(groupPath, recursive=True)
        else:
            zk.delete(hostPath, recursive=True)
            zk.stop()
            return "DELETED ==> host: {0} in group: {1}".format(hostName, groupName)

    elif len(znodeStringSplited) == 1:

        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]

        if zk.exists(groupPath) is None:
            zk.stop()
            return "ERROR  ==> could not delete group: {0} that does not exist !!!".format(groupName)

        else:
            zk.delete(groupPath, recursive=True)
    else:
        return "ERROR with processing znodeStrings !!!"

    zk.stop()
    return "DELETED ==> group: {0}".format(groupName)

# def updateZnode(znodeDict):
#     '''
#     Update znode with hostvars.

#     Return a string (UPDATED   ==> host: <hostname> to group: <groupname>).
#     '''

#     zk = zkStartRw()

#     groupName   = znodeDict.keys()[0]
#     hostName    = znodeDict[groupName].keys()[0]
#     groupPath   = "{0}/groups/{1}".format(cfg.aPath, groupName)
#     hostPath    = "{0}/{1}".format(groupPath, hostName)
#     hostVarList = zk.get_children(hostPath)

#     if zk.exists(hostPath) is None:
#         zk.stop()
#         return "ERROR  ==> host: {0} in group {1} does not exist !!!".format(hostName, groupName)

#     for hostVar in hostVarList:
#         if zk.exists("{0}/{1}".format(hostPath, hostVar)) is None:
#             zk.stop()
#             return "ERROR  ==> hostvar: {0} for host {1} in group {2} does not exist !!!".format(hostVar, hostName, groupName)

#     for key in znodeDict[groupName][hostName]:
#         varPath = "{0}/{1}".format(hostPath, key)
#         varVal  = znodeDict[groupName][hostName][key]
#         if CONDITION: #FIXME!!!
#            zk.create(varPath, varVal)

#     zk.stop()
#     return "ADDED   ==> host: {0} to group: {1}".format(hostName, groupName)

def hostVarsShow(znodeStringSplited):
    '''
    Show hostvars for a given <groupname:hostname> or <groupname>.

    Return dict or string (in case of ERROR).
    '''

    zk = zkStartRo()

    if len(znodeStringSplited) > 1:

        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]
        hostName  = znodeStringSplited[1][0]
        hostPath  = znodeStringSplited[1][1]

        if zk.exists(hostPath) is None:
            zk.stop()
            return "ERROR  ==> no such hostname: {0} in group {1} !!!".format(hostName, groupName)

        else:
            hostVarList = zk.get_children(hostPath)
            valDict     = {}

            for var in hostVarList:
                valDict[var] = zk.get('{0}/{1}'.format(hostPath, var))[0]

    elif len(znodeStringSplited) == 1:
        groupName = znodeStringSplited[0][0]
        groupPath = znodeStringSplited[0][1]

        if zk.exists(groupPath) is None:
            zk.stop()
            return "ERROR  ==> no such groupname: {0} !!!".format(groupName)

        else:
            hostList = zk.get_children(groupPath)
            valDict  = {}

            for host in hostList:
                tmpHostPath    = '{0}/{1}'.format(groupPath, host)
                valDict[host]  = zk.get_children('{0}'.format(tmpHostPath))

            for host in valDict.keys():
                varDict = {}

                for var in valDict[host]:
                    tmpHostPath   = '{0}/{1}'.format(groupPath, host)
                    varDict[var]  = zk.get('{0}/{1}'.format(tmpHostPath, var))[0]
                    valDict[host] = varDict

            return valDict

    else:
        return "ERROR with processing znodeStrings !!!"

    zk.stop()
    return valDict

def inventoryDump():
    '''
    Inventory dump for a given list of zookeeper servers and ansible-keeper path.

    Return dict.
    '''

    zk = zkStartRo()

    groupList = zk.get_children("{}/groups".format(cfg.aPath))
    groupDict = {}

    for group in groupList:
        path     = "{0}/groups/{1}".format(cfg.aPath, group)
        children = zk.get_children(path)
        groupDict[group] = children

    zk.stop()
    return groupDict

def ansibleInventoryDump():
    '''
    Ansible compliant inventory dump for a given list of zookeeper servers and ansible-keeper path.

    Return dict.
    '''

    zk = zkStartRo()

    groupList = zk.get_children("{}/groups".format(cfg.aPath))
    groupDict = {}

    for group in groupList:
        path     = "{0}/groups/{1}".format(cfg.aPath, group)
        children = zk.get_children(path)
        tmpDict  = {}
        tmpDict['hosts'] = children
        tmpDict['vars']  = {"a":"b"}
        groupDict[group] = tmpDict

    zk.stop()
    return groupDict

def main():
    '''
    Main logic
    '''

    if oParser()['inventoryMode'] == 'true':
        print json.dumps(inventoryDump())

    if oParser()['inventoryMode'] == 'ansible':
        print json.dumps(ansibleInventoryDump())

    if oParser()['addMode'] is not None:
        znodeDict = splitZnodeVarString(oParser()['addMode'])
        print addHostWithHostvars(znodeDict)

    if oParser()['groupMode'] is not None:
        znodeStringSplited = splitZnodeString(oParser()['groupMode'])
        print addHostToGroup(znodeStringSplited)

    if oParser()['deleteMode'] is not None:
        znodeStringSplited = splitZnodeString(oParser()['deleteMode'])
        print deleteZnodeRecur(znodeStringSplited)

    if oParser()['showMode'] is not None:
        znodeStringSplited = splitZnodeString(oParser()['showMode'])
        print json.dumps(hostVarsShow(znodeStringSplited))

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