<commit_before>import cv2
import numpy as np
import random
from sklearn.cluster import k_means
from sklearn.preprocessing import normalize

def imageSIFT(img, n_clusters = 20):
    s = cv2.SURF()
    pic = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    keypoints, descriptors = s.detectAndCompute(pic, None)
    descriptors = normalize(descriptors, norm = 'l2', axis = 1)
    centroid, l, i = k_means(descriptors, n_clusters = n_clusters)
    return np.reshape(centroid, newshape = (1, n_clusters * 64))

# TODO: Rewrite this function to construct a bag-of-word feature
def extractSIFT(images):
    return np.asarray([imageSIFT(images[index]) for index in range(len(images))])

# TODO: SaliencyELD Call - Linrong Jin
# img1 is the original image, img2 is the image processed by SaliencyELD
# the return is the image without background
def removeBackground(img1, img2):
    ret, mask = cv2.threshold(img2, 10, 255, cv2.THRESH_BINARY)
    margin_x, margin_y = np.sum(mask, axis = 0), np.sum(mask, axis = 1)
    left, right, top, bottom = -1, -1, -1, -1
    for i in range(margin_x.shape[0]):
        if margin_x[i] > 0:
            left = i
            break
    for i in range(margin_x.shape[0]):
        if margin_x[margin_x.shape[0] - i - 1] > 0:
            right = margin_x.shape[0] - i - 1
            break
    for i in range(margin_y.shape[0]):
        if margin_y[i] > 0:
            top = i
            break
    for i in range(margin_y.shape[0]):
        if margin_y[margin_y.shape[0] - i - 1] > 0:
            bottom = margin_y.shape[0] - i - 1
            break
    img1_bg = cv2.bitwise_and(img1, img1, mask = mask)
    return img1_bg[top:bottom, left:right]

def resizeImages(images, size = (64, 64)):
    return np.asarray([cv2.resize(images[index], size) for index in range(images.shape[0])])

def rotateImage(image, angle):
    height = image.shape[0]
    width = image.shape[1]
    height_big = height
    width_big = width
    image_big = cv2.resize(image, (width_big, height_big))
    image_center = (width_big / 2, height_big / 2)
    rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
    result = cv2.warpAffine(image_big, rot_mat, (width_big, height_big), flags = cv2.INTER_LINEAR)
    return result

# flag = 0 / 1 / -1
def flipImage(image, flag):
    return cv2.flip(image, flag)

def scale(image, multiple, interpolation = None):
    rows, cols, channels = image.shape
    return cv2.resize(image, (int(cols * multiple), int(rows * multiple)), interpolation = interpolation)

def zoom(image, dsize, interpolation = None):
    return cv2.resize(image, dsize, interpolation = interpolation)

def shift(image, x, y):
    M = np.array([[1, 0, x], [0, 1, y]], np.float32)
    rows, cols, channels = image.shape
    return cv2.warpAffine(image, M, (cols, rows))

def contrast(image, alpha, beta):
    image = image * alpha + beta
    return image

def noise(image, size):
    for i in range(0, size):
        xi = int(np.random.uniform(0, image.shape[1]))
        xj = int(np.random.uniform(0, image.shape[0]))
        image[xj, xi] = 255
    return image

def shuffle(image, label, features = None):
    if features is not None:
        _images, _labels, _features = image.tolist(), label.tolist(), features.tolist()
        table = []
        for i in range(len(_images)):
            table.append([_images[i], _labels[i], _features[i]])
        random.shuffle(table)
        return np.asarray([item[0] for item in table]), np.asarray([item[1] for item in table]), np.asarray([item[2] for item in table])
    else:
        _images, _labels = image.tolist(), label.tolist()
        table = []
        for i in range(len(_images)):
            table.append([_images[i], _labels[i]])
        random.shuffle(table)
        return np.asarray([item[0] for item in table]), np.asarray([item[1] for item in table]), None
<commit_msg>完成重构之后的SIFT特征提取（转换成bag-of-words）<commit_after>import cv2
import numpy as np
import random
from sklearn.cluster import k_means
from sklearn.preprocessing import normalize

def imageSIFT(img):
    s = cv2.SURF()
    pic = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    keypoints, descriptors = s.detectAndCompute(pic, None)
    return descriptors

def extractSIFT(images, n_clusters = 100):
    sifts = [imageSIFT(images[index]) for index in range(len(images))]
    _map = [[], [], np.zeros(shape = (len(images), n_clusters))]
    cnt = 0
    for patch in sifts:
        for point in patch:
            _map[0].append(point)
            _map[1].append(cnt)
        cnt += 1
    _map[0] = normalize(np.asarray(_map[0]), norm = "l1", axis = 0)
    centroid, labels = k_means(_map[0], n_clusters = n_clusters)
    for index in labels.shape[0]:
        _map[2][_map[1][index], labels[index]] += 1
    return normalize(_map[2], axis = 0)

def removeBackground(img1, img2):
    ret, mask = cv2.threshold(img2, 10, 255, cv2.THRESH_BINARY)
    margin_x, margin_y = np.sum(mask, axis = 0), np.sum(mask, axis = 1)
    left, right, top, bottom = -1, -1, -1, -1
    for i in range(margin_x.shape[0]):
        if margin_x[i] > 0:
            left = i
            break
    for i in range(margin_x.shape[0]):
        if margin_x[margin_x.shape[0] - i - 1] > 0:
            right = margin_x.shape[0] - i - 1
            break
    for i in range(margin_y.shape[0]):
        if margin_y[i] > 0:
            top = i
            break
    for i in range(margin_y.shape[0]):
        if margin_y[margin_y.shape[0] - i - 1] > 0:
            bottom = margin_y.shape[0] - i - 1
            break
    img1_bg = cv2.bitwise_and(img1, img1, mask = mask)
    return img1_bg[top:bottom, left:right]

def resizeImages(images, size = (64, 64)):
    return np.asarray([cv2.resize(images[index], size) for index in range(images.shape[0])])

def rotateImage(image, angle):
    height = image.shape[0]
    width = image.shape[1]
    height_big = height
    width_big = width
    image_big = cv2.resize(image, (width_big, height_big))
    image_center = (width_big / 2, height_big / 2)
    rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
    result = cv2.warpAffine(image_big, rot_mat, (width_big, height_big), flags = cv2.INTER_LINEAR)
    return result

def flipImage(image, flag):
    return cv2.flip(image, flag)

def scale(image, multiple, interpolation = None):
    rows, cols, channels = image.shape
    return cv2.resize(image, (int(cols * multiple), int(rows * multiple)), interpolation = interpolation)

def zoom(image, dsize, interpolation = None):
    return cv2.resize(image, dsize, interpolation = interpolation)

def shift(image, x, y):
    M = np.array([[1, 0, x], [0, 1, y]], np.float32)
    rows, cols, channels = image.shape
    return cv2.warpAffine(image, M, (cols, rows))

def contrast(image, alpha, beta):
    image = image * alpha + beta
    return image

def noise(image, size):
    for i in range(0, size):
        xi = int(np.random.uniform(0, image.shape[1]))
        xj = int(np.random.uniform(0, image.shape[0]))
        image[xj, xi] = 255
    return image

def shuffle(image, label, features = None):
    if features is not None:
        _images, _labels, _features = image.tolist(), label.tolist(), features.tolist()
        table = []
        for i in range(len(_images)):
            table.append([_images[i], _labels[i], _features[i]])
        random.shuffle(table)
        return np.asarray([item[0] for item in table]), np.asarray([item[1] for item in table]), np.asarray([item[2] for item in table])
    else:
        _images, _labels = image.tolist(), label.tolist()
        table = []
        for i in range(len(_images)):
            table.append([_images[i], _labels[i]])
        random.shuffle(table)
        return np.asarray([item[0] for item in table]), np.asarray([item[1] for item in table]), None
<|endoftext|>


<commit_before>'''
sympy helper functions
'''

import sympy
from sympy import *

th,ph,psi,thd,phd,psid,thdd,phdd,psidd = symbols('theta,phi,psi,thetadot,phidot,psidot,thetaddot,phiddot,psiddot')
w1,w2,w3 = symbols('omega_1,omega_2,omega_3')
t,g,m,h = symbols('t,g,m,h')
circmat = Matrix([eye(3)[2-j,:] for j in range(3)]) #define circulant matrix
polarframe = ['\mathbf{e_r}','\mathbf{e_\\theta}','\mathbf{e_z}']
sphericalframe = ['\mathbf{e_\\phi}','\mathbf{e_\\theta}','\mathbf{e_\\rho}']

def difftotal(expr, diffby, diffmap):
    """Take the total derivative with respect to a variable.
    Originally by Chris Wagner (http://robotfantastic.org/total-derivatives-in-sympy.html)

    Example:

        theta, t, theta_dot = symbols("theta t theta_dot")
        difftotal(cos(theta), t, {theta: theta_dot})

    returns

        -theta_dot*sin(theta)
    """
    # Replace all symbols in the diffmap by a functional form
    fnexpr = expr.subs({s:s(diffby) for s in diffmap})
    # Do the differentiation
    diffexpr = diff(fnexpr, diffby)
    # Replace the Derivatives with the variables in diffmap
    derivmap = {Derivative(v(diffby), diffby):dv
                for v,dv in diffmap.iteritems()}
    finaldiff = diffexpr.subs(derivmap)
    # Replace the functional forms with their original form
    return finaldiff.subs({s(diffby):s for s in diffmap})

def difftotalmat(mat,diffby,diffmap):
    """Apply total derivative element by element to a matrix"""
    return Matrix([difftotal(x,diffby,diffmap) for x in mat]).reshape(*mat.shape)

def transportEq(vec, diffby, diffmap, omega):
    """Apply the transport equation to a vector:
    Frame I d/dt(vec) = Frame B d/dt(vec) + omega \cross vec
    Assumes that vector vec and angular velocity vector omega
    are given in terms of Frame B unit vectors (in the same,
    right-handed order, i.e. b_1, b_2, b_3). diffby is the
    variable with respect to which we differentiate (typically
    time) and diffmap is a dictionary of derivatives of dependent
    variables. See difftotal for details."""

    return difftotalmat(vec,diffby,diffmap) + omega.cross(vec)

def rotMat(axis,angle):
    """ Returns the DCM ({}^B C^A) for a frame rotation of angle about
    the specified axis
    """
    if axis == 1:
        return Matrix(([1,0,0],[0,cos(angle),sin(angle)],[0,-sin(angle),cos(angle)]))
    elif axis == 2:
        return Matrix(([cos(angle),0,-sin(angle)],[0,1,0],[sin(angle), 0, cos(angle)]))
    elif axis == 3:
        return Matrix(([cos(angle),sin(angle),0],[-sin(angle),cos(angle),0],[0,0,1]))
    else:
        return -1

def parallelAxis(I_G,r_QG,m):
    """ Applies the parallel axis theorem to matrix of inertia I_G
    (assumed to be about the COM) to find the matrix of inertia I_Q
    where the vector from G to Q is r_QG and the total mass of the
    body is m.  I_G and I_QG are assumed to be with respect to the
    same body frame."""

    return I_G + m*((r_QG.transpose()*r_QG)[0]*eye(3) + r_QG*r_QG.transpose())

def skew(v):
    """Given 3x1 vector v, return skew-symmetric matrix"""

    assert ((hasattr(v,'__iter__') or isinstance(v,Matrix)) and len(v)==3),\
            "v must be an iterable of length 3."

    return Matrix([
        [0,   -v[2], v[1]],
        [v[2],  0,  -v[0]],
        [-v[1],v[0],  0  ]
    ])

def calcDCM(n,th):
    """Calculated the DCM  ({}^A C^B) for a rotation of angle theta about an
    axis n"""

    return eye(3)*cos(th)+skew(n)*sin(th)+(1-cos(th))*n*n.T

def DCM2angVel(dcm,diffmap):
    """ Given a direction cosine matrix (DCM) transforming from
    frame A to frame B ({}^B C^A) returns the angular velocity of
    frame B in frame A ({}^A\omega^B). The returned vector will be
    the components in frame B.

    Assumes that differentiation is with respect to symbol t.
    diffmap is defined as in difftotal
    """

    #s = solve(dcm*difftotalmat(dcm,t,diffmap).T-skew([w1,w2,w3]),(w1,w2,w3))
    #return Matrix([s[w1],s[w2],s[w3]])

    tmp = dcm*difftotalmat(dcm,t,diffmap).T
    return simplify(Matrix([tmp[2,1],tmp[0,2],tmp[1,0]]))

def mat2vec(mat,basis='e'):
    """ Transform matrix representation of a vector to the vector equation
    for a given basis.

    basis can be either a string (i.e. 'e' becomes basis e_1,e_2,e3) or an
    iterable of length 3.

    mat is assumed to be sympy matrix representing a column vector
    """

    assert isinstance(basis,str) or (hasattr(basis,'__iter__') and len(basis)==3),\
            "basis input must be a string or iterable of length 3."

    if isinstance(basis,str):
        basis = ['\mathbf{'+basis+'}_'+str(j) for j in range(1,4)]

    basissyms = symbols(basis,commutative=False)
    basisvec = Matrix(basissyms)

    return (mat.T*basisvec)[0]

def fancyMat(prefix,shape):
    """ Create an indexed matrix akin to symarray using the prefix
    and with dimensions given by shape (this must be a 2 element iterable)

    Indexing is 1-based.
    """

    M = []
    for r in range(1,shape[0]+1):
        row = []
        for c in range(1,shape[1]+1):
            row.append(prefix+'_{'+str(r)+str(c)+'}')
        M.append(row)
    M = Matrix(symbols(M))

    return M

def EulerAngSet(rots, angs):
    """ Calculate the total DCM for an Euler Angle set defined by
    ordered rotations about body-fixed axes rots of angles angs.

    Final DCM is {}^B C^A where the first rotation and angle in rots
    and angs refer to rotation starting from frame A (and the final
    entry in rots and angs is rotation into frame B).
    """

    assert (hasattr(rots,'__iter__') and len(rots)==3),\
            "rots must be an iterable of length 3."

    assert (hasattr(angs,'__iter__') and len(angs)==3),\
            "v must be an iterable of length 3."

    DCM = eye(3)
    for rot,ang in zip(rots,angs):
        DCM = rotMat(rot,ang)*DCM

    return simplify(DCM)

def EulerLagrange(L, qs, diffmap):
    """ Apply the Euler-Lagrange equations to Lagrangian L, with
    generalized coordinates qs, with their derivatives identified by the
    differentation map diffmap (see difftotal).

    Note that qs must be iterable even if this is only a one dof system.

    diffmap must map all qs to their second derivatives (via the first
    derivative).  i.e., for any q = \theta, diffmap must contain \dot\theta
    and \ddot\theta.

    The returned value will be a system of equations for the second
    time derivatives of the generalized coordinates.
    """

    assert hasattr(qs,'__iter__'),\
            "qs must be an iterable."

    eqs = [simplify(difftotal(L.diff(diffmap[q]),t,diffmap) - L.diff(q)) \
            for q in qs]

    return simplify(solve(eqs,[diffmap[diffmap[q]] for q in qs]))

<commit_msg>fixed parallel axis theorem and added docstring<commit_after>'''
sympy helper functions
'''

import sympy
from sympy import *

th,ph,psi,thd,phd,psid,thdd,phdd,psidd = symbols('theta,phi,psi,thetadot,phidot,psidot,thetaddot,phiddot,psiddot')
w1,w2,w3 = symbols('omega_1,omega_2,omega_3')
t,g,m,h = symbols('t,g,m,h')
circmat = Matrix([eye(3)[2-j,:] for j in range(3)]) #define circulant matrix
polarframe = ['\mathbf{e_r}','\mathbf{e_\\theta}','\mathbf{e_z}']
sphericalframe = ['\mathbf{e_\\phi}','\mathbf{e_\\theta}','\mathbf{e_\\rho}']

def difftotal(expr, diffby, diffmap):
    """Take the total derivative with respect to a variable.
    Originally by Chris Wagner (http://robotfantastic.org/total-derivatives-in-sympy.html)

    Example:

        theta, t, theta_dot = symbols("theta t theta_dot")
        difftotal(cos(theta), t, {theta: theta_dot})

    returns

        -theta_dot*sin(theta)
    """
    # Replace all symbols in the diffmap by a functional form
    fnexpr = expr.subs({s:s(diffby) for s in diffmap})
    # Do the differentiation
    diffexpr = diff(fnexpr, diffby)
    # Replace the Derivatives with the variables in diffmap
    derivmap = {Derivative(v(diffby), diffby):dv
                for v,dv in diffmap.iteritems()}
    finaldiff = diffexpr.subs(derivmap)
    # Replace the functional forms with their original form
    return finaldiff.subs({s(diffby):s for s in diffmap})

def difftotalmat(mat,diffby,diffmap):
    """Apply total derivative element by element to a matrix"""
    return Matrix([difftotal(x,diffby,diffmap) for x in mat]).reshape(*mat.shape)

def transportEq(vec, diffby, diffmap, omega):
    """Apply the transport equation to a vector:
    Frame I d/dt(vec) = Frame B d/dt(vec) + omega \cross vec
    Assumes that vector vec and angular velocity vector omega
    are given in terms of Frame B unit vectors (in the same,
    right-handed order, i.e. b_1, b_2, b_3). diffby is the
    variable with respect to which we differentiate (typically
    time) and diffmap is a dictionary of derivatives of dependent
    variables. See difftotal for details."""

    return difftotalmat(vec,diffby,diffmap) + omega.cross(vec)

def rotMat(axis,angle):
    """ Returns the DCM ({}^B C^A) for a frame rotation of angle about
    the specified axis
    """
    if axis == 1:
        return Matrix(([1,0,0],[0,cos(angle),sin(angle)],[0,-sin(angle),cos(angle)]))
    elif axis == 2:
        return Matrix(([cos(angle),0,-sin(angle)],[0,1,0],[sin(angle), 0, cos(angle)]))
    elif axis == 3:
        return Matrix(([cos(angle),sin(angle),0],[-sin(angle),cos(angle),0],[0,0,1]))
    else:
        return -1

def parallelAxis(I_G,r_QG,m):
    """ Applies the parallel axis theorem to matrix of inertia I_G
    (assumed to be about the COM) to find the matrix of inertia I_Q
    where the vector from G to Q is r_QG and the total mass of the
    body is m.  I_G and I_QG are assumed to be with respect to the
    same body frame."""

    return I_G + m*((r_QG.transpose()*r_QG)[0]*eye(3) - r_QG*r_QG.transpose())

def skew(v):
    """Given 3x1 vector v, return skew-symmetric matrix"""

    assert ((hasattr(v,'__iter__') or isinstance(v,Matrix)) and len(v)==3),\
            "v must be an iterable of length 3."

    return Matrix([
        [0,   -v[2], v[1]],
        [v[2],  0,  -v[0]],
        [-v[1],v[0],  0  ]
    ])

def calcDCM(n,th):
    """Calculated the DCM  ({}^A C^B) for a rotation of angle theta about an
    axis n"""

    return eye(3)*cos(th)+skew(n)*sin(th)+(1-cos(th))*n*n.T

def DCM2angVel(dcm,diffmap):
    """ Given a direction cosine matrix (DCM) transforming from
    frame A to frame B ({}^B C^A) returns the angular velocity of
    frame B in frame A ({}^A\omega^B). The returned vector will be
    the components in frame B.

    Assumes that differentiation is with respect to symbol t.
    diffmap is defined as in difftotal
    """

    #s = solve(dcm*difftotalmat(dcm,t,diffmap).T-skew([w1,w2,w3]),(w1,w2,w3))
    #return Matrix([s[w1],s[w2],s[w3]])

    tmp = dcm*difftotalmat(dcm,t,diffmap).T
    return simplify(Matrix([tmp[2,1],tmp[0,2],tmp[1,0]]))

def mat2vec(mat,basis='e'):
    """ Transform matrix representation of a vector to the vector equation
    for a given basis.

    basis can be either a string (i.e. 'e' becomes basis e_1,e_2,e3) or an
    iterable of length 3.

    mat is assumed to be sympy matrix representing a column vector
    """

    assert isinstance(basis,str) or (hasattr(basis,'__iter__') and len(basis)==3),\
            "basis input must be a string or iterable of length 3."

    if isinstance(basis,str):
        basis = ['\mathbf{'+basis+'}_'+str(j) for j in range(1,4)]

    basissyms = symbols(basis,commutative=False)
    basisvec = Matrix(basissyms)

    return (mat.T*basisvec)[0]

def fancyMat(prefix,shape):
    """ Create an indexed matrix akin to symarray using the prefix
    and with dimensions given by shape (this must be a 2 element iterable)

    Indexing is 1-based.

    Example:
        fancyMat('{}^\mathcal{B}C^{\mathcal{A}}',(3,3))
    """

    M = []
    for r in range(1,shape[0]+1):
        row = []
        for c in range(1,shape[1]+1):
            row.append(prefix+'_{'+str(r)+str(c)+'}')
        M.append(row)
    M = Matrix(symbols(M))

    return M

def EulerAngSet(rots, angs):
    """ Calculate the total DCM for an Euler Angle set defined by
    ordered rotations about body-fixed axes rots of angles angs.

    Final DCM is {}^B C^A where the first rotation and angle in rots
    and angs refer to rotation starting from frame A (and the final
    entry in rots and angs is rotation into frame B).
    """

    assert (hasattr(rots,'__iter__') and len(rots)==3),\
            "rots must be an iterable of length 3."

    assert (hasattr(angs,'__iter__') and len(angs)==3),\
            "v must be an iterable of length 3."

    DCM = eye(3)
    for rot,ang in zip(rots,angs):
        DCM = rotMat(rot,ang)*DCM

    return simplify(DCM)

def EulerLagrange(L, qs, diffmap):
    """ Apply the Euler-Lagrange equations to Lagrangian L, with
    generalized coordinates qs, with their derivatives identified by the
    differentation map diffmap (see difftotal).

    Note that qs must be iterable even if this is only a one dof system.

    diffmap must map all qs to their second derivatives (via the first
    derivative).  i.e., for any q = \theta, diffmap must contain \dot\theta
    and \ddot\theta.

    The returned value will be a system of equations for the second
    time derivatives of the generalized coordinates.
    """

    assert hasattr(qs,'__iter__'),\
            "qs must be an iterable."

    eqs = [simplify(difftotal(L.diff(diffmap[q]),t,diffmap) - L.diff(q)) \
            for q in qs]

    return simplify(solve(eqs,[diffmap[diffmap[q]] for q in qs]))

<|endoftext|>


<commit_before>import numpy as np
import pytest

import lattice

def test_first_difference():
    assert lattice.first_difference([], [], reverse=False) == 0
    assert lattice.first_difference([], [], reverse=True) == 0
    assert lattice.first_difference([0], [9], reverse=False) == 0
    assert lattice.first_difference([0, 2, 4], [0, 2, 4], reverse=False) == 3
    assert lattice.first_difference([9], [9], reverse=False) == 1
    assert lattice.first_difference(['ant', 'dog', 'cat'],
                                    ['ant', 'dog', 'bird'], reverse=False) == 2
    assert lattice.first_difference(['ant', 'dog', 'cat'],
                                    ['ant', 'dog', 'bird'], reverse=True) == 0

def test_parse_lattice_spec():
    """The lattice specification should be a string such as "-5,10,5/0.5,2,2",
    which defines a two dimensional lattice where the x coordinate goes
    from -5 to 10 in 5 steps, while the y coordinate goes from 0.5 to 2 in 2
    steps. Another example is "-5,10,5/1.23,1.23,1", which defines a one
    dimensional lattice (since the y component is fixed to 1.23).
    The same can be expressed also as "-5,10,5/1.23".
    """
    s = lattice.parse_lattice_spec

    # 1d tests:
    retval = s("0,10,11")

    # data types
    assert retval == [[0., 10., 11]]
    assert isinstance(retval[0][0], float)
    assert isinstance(retval[0][1], float)
    assert isinstance(retval[0][2], int)

    # 1d data with constant
    assert s("3") == [[3, 3, 1]]

    # from docstring
    assert s("-5,10,5/1.23,1.23,1") == s("-5,10,5/1.23")

    # 2d tests
    assert s("0,3,2/-3,3,7") == [[0., 3, 2], [-3, 3, 7]]

    # from docstring
    assert s("-5,10,5/0.5,2,2") == [[-5, 10, 5], [0.5, 2, 2]]

    # 3d tests
    assert s("0,3,2/-3,3,7/-0.5,0.7,17") == [[0., 3, 2],
                                             [-3, 3, 7],
                                             [-0.5, 0.7, 17]]

    # spaces
    assert s("0,   10,   11") == [[0., 10., 11]]

    # check error
    with pytest.raises(ValueError):
        assert s("0, 10, 11, 42")

def test_lattice_object():
    # can call with spec as in min, max, num:
    l1 = lattice.Lattice([[0., 10., 6]])
    # or with string
    l2 = lattice.Lattice("0.,10.,6")

    # should result in the same object
    assert np.allclose(l1.get_positions(), l2.get_positions())

    # only order 'C' and 'F' allowed
    l2 = lattice.Lattice("0.,10.,6", order='F')
    l2 = lattice.Lattice("0.,10.,6", order='C')
    with pytest.raises(ValueError):
        l2 = lattice.Lattice("0.,10.,6", order='B')

def test_lattice_object_1d():
    l = lattice.Lattice([[0, 10, 6]])
    assert l.dim == 1

    check_points_index = []
    check_points_value = []

    def f(idx, x):
        print("idx = {} x={}".format(idx, x))
        check_points_index.append(idx[0])
        check_points_value.append(x[0])

    l.foreach(f)

    assert check_points_index == list(range(6))
    assert np.allclose(check_points_value, np.linspace(0, 10, 6))

    assert l.get_closest([1.01]) == [1]
    assert l.get_closest([0.99]) == [0]

    assert l.get_num_points() == 6

    assert l.get_pos_from_idx([4]) == [8]

    assert l.get_shape() == [6]

    assert l.max_node_pos == [10]
    assert l.min_node_pos == [0]

    assert l.min_max_num_list == [[0, 10, 6]]
    assert l.order == 'F'
    assert l.nodes == [6]
    assert str(l) == "Lattice([[0, 10, 6]])"
    assert l.stepsizes == [2]

    assert np.allclose(l.get_positions()[:, 0], np.linspace(0, 10, 6))

def test_lattice_object_1d_scales():
    l = lattice.Lattice([[0, 10, 6]])
    assert l.dim == 1

    l.scale(0.5)

    assert l.max_node_pos == [5]
    assert l.stepsizes == [1]

    l.scale(1 / 0.5)
    assert l.max_node_pos == [10]
    assert l.stepsizes == [2]

    # scale differently in different directions (silly for 1d)
    l.scale([0.5])

    assert l.max_node_pos == [5]
    assert l.stepsizes == [1]

    l.scale(1 / 0.5)
    assert l.max_node_pos == [10]
    assert l.stepsizes == [2]

def test_lattice_object_2d():
    l = lattice.Lattice([[0, 10, 6], [-3, 1, 2]])
    assert l.dim == 2

    check_points_index = []
    check_points_value = []

    def f(idx, x):
        print("idx = {} x={}".format(idx, x))
        check_points_index.append(idx[:])
        check_points_value.append(x[:])

    l.foreach(f)

    assert check_points_index == [[0, 0],
                                  [1, 0],
                                  [2, 0],
                                  [3, 0],
                                  [4, 0],
                                  [5, 0],
                                  [0, 1],
                                  [1, 1],
                                  [2, 1],
                                  [3, 1],
                                  [4, 1],
                                  [5, 1]]
    assert np.allclose(check_points_value, [[0.0, -3.0],
                                            [2.0, -3.0],
                                            [4.0, -3.0],
                                            [6.0, -3.0],
                                            [8.0, -3.0],
                                            [10.0, -3.0],
                                            [0.0, 1.0],
                                            [2.0, 1.0],
                                            [4.0, 1.0],
                                            [6.0, 1.0],
                                            [8.0, 1.0],
                                            [10.0, 1.0]])

    assert l.get_closest([1.01, -0.9]) == [1, 1]
    assert l.get_closest([1.01, -1.0001]) == [1, 0]
    assert l.get_num_points() == 12

    assert l.dim == 2

    with pytest.raises(IndexError):
        l.get_pos_from_idx([1])  # should raise IndexError for 2d mesh

    l.get_pos_from_idx([0, 0]) == [0., -3.]
    l.get_pos_from_idx([0, 1]) == [0., 1.]
    l.get_pos_from_idx([4, 1]) == [8., 1.]

    assert l.get_shape() == [6, 2]

    l.max_node_pos == [10, 1]

    assert l.min_max_num_list == [[0, 10, 6], [-3, 1, 2]]
    assert l.min_node_pos == [0, -3]
    assert l.nodes == [6, 2]
    assert l.order == 'F'
    assert str(l) == "Lattice([[0, 10, 6], [-3, 1, 2]])"

    assert np.allclose(l.get_positions(),
                       np.array([[[0., -3.],
                                  [2., -3.],
                                  [4., -3.],
                                  [6., -3.],
                                  [8., -3.],
                                  [10., -3.]],
                                 [[0., 1.],
                                  [2., 1.],
                                  [4., 1.],
                                  [6., 1.],
                                  [8., 1.],
                                  [10., 1.]]]))

def test_lattice_object_2d_scale():

    l = lattice.Lattice([[0, 10, 6], [-3, 1, 2]])
    assert l.dim == 2

    l.scale(2)
    assert l.max_node_pos == [20, 2]
    assert l.stepsizes == [4, 8]

    l.scale(0.5)
    assert l.max_node_pos == [10, 1]
    assert l.stepsizes == [2, 4]

def test_lattice_add():
    l1 = lattice.Lattice([[0, 10, 6]])
    l2 = lattice.Lattice([[0, 5, 3]])
    l_add = l1 + l2
    assert l1.dim == 1
    assert l2.dim == 1
    assert l_add.dim == 2

def test_field_lattice_init():
    fl = lattice.FieldLattice(lattice.Lattice([[0, 10, 6], [-3, 1, 2]]))

    # we want to see numpy.ndarray for the data
    assert type(fl.field_data) == np.ndarray
    # and contiguous data
    assert fl.field_data.flags['F_CONTIGUOUS'] is True
    # and the right shape
    assert fl.field_data.shape[1:] == tuple(fl.lattice.get_shape())
    # and 3d data vectors
    assert fl.field_data.shape[0] == 3 == fl.field_dim
    # but the lattice is in 2d only
    assert fl.lattice.dim == 2

def test_field_lattice_set():
    z_counter = [0]  # cheat to see in which order field vaules are set

    def mydata(pos):
        print("type of pos = {}".format(type(pos)))
        print("pos = {}".format(pos))
        # use position as field data (for our testing here)
        x, y = pos
        z = z_counter[0]
        z_counter[0] += 0.1
        return [x, y, z]

    fl = lattice.FieldLattice(lattice.Lattice([[0, 5, 6], [-2, -1, 2]]))
    fl.set(mydata)
    print(fl.field_data.shape)
    print(fl.field_data[:, :, :])
    print("==============")
    print(fl.field_data[2, :, :])
    print(np.array([np.linspace(0, 1.1, 12)]))
    print("==============")

    # check x-component
    assert np.allclose(fl.field_data[0, :, :],
                       np.array([[0., 0.],
                                 [1., 1.],
                                 [2., 2.],
                                 [3., 3.],
                                 [4., 4.],
                                 [5., 5.]]))

    # check y-component
    assert np.allclose(fl.field_data[1, :, :],
                       np.array([[-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.]]))

    # check z-component in Fortran ordering
    assert np.allclose(fl.field_data[2, :, :],
                       np.array([[0.0, 0.6],
                                 [0.1, 0.7],
                                 [0.2, 0.8],
                                 [0.3, 0.9],
                                 [0.4, 1.0],
                                 [0.5, 1.1]]))

    # How does C order differ?
    # Okay, simplify: establish only that C-order is different from F-order:
    flf = fl  # FieldLattice Fortran

    # Create FieldLattice C with identical data
    flc = lattice.FieldLattice(lattice.Lattice([[0, 5, 6], [-2, -1, 2]]),
                               order="C")
    z_counter[0] = 0
    fl.set(mydata)

    # compare values (should be the same)
    assert np.allclose(flc.field_data, flf.field_data)

    # check that the strides are different (check help(numpy.ndarray))
    assert flc.field_data.strides != flf.field_data.strides
<commit_msg>record that test fails on Osiris, but passes elsewhere<commit_after>import numpy as np
import pytest

import lattice

def test_first_difference():
    assert lattice.first_difference([], [], reverse=False) == 0
    assert lattice.first_difference([], [], reverse=True) == 0
    assert lattice.first_difference([0], [9], reverse=False) == 0
    assert lattice.first_difference([0, 2, 4], [0, 2, 4], reverse=False) == 3
    assert lattice.first_difference([9], [9], reverse=False) == 1
    assert lattice.first_difference(['ant', 'dog', 'cat'],
                                    ['ant', 'dog', 'bird'], reverse=False) == 2
    assert lattice.first_difference(['ant', 'dog', 'cat'],
                                    ['ant', 'dog', 'bird'], reverse=True) == 0

def test_parse_lattice_spec():
    """The lattice specification should be a string such as "-5,10,5/0.5,2,2",
    which defines a two dimensional lattice where the x coordinate goes
    from -5 to 10 in 5 steps, while the y coordinate goes from 0.5 to 2 in 2
    steps. Another example is "-5,10,5/1.23,1.23,1", which defines a one
    dimensional lattice (since the y component is fixed to 1.23).
    The same can be expressed also as "-5,10,5/1.23".
    """
    s = lattice.parse_lattice_spec

    # 1d tests:
    retval = s("0,10,11")

    # data types
    assert retval == [[0., 10., 11]]
    assert isinstance(retval[0][0], float)
    assert isinstance(retval[0][1], float)
    assert isinstance(retval[0][2], int)

    # 1d data with constant
    assert s("3") == [[3, 3, 1]]

    # from docstring
    assert s("-5,10,5/1.23,1.23,1") == s("-5,10,5/1.23")

    # 2d tests
    assert s("0,3,2/-3,3,7") == [[0., 3, 2], [-3, 3, 7]]

    # from docstring
    assert s("-5,10,5/0.5,2,2") == [[-5, 10, 5], [0.5, 2, 2]]

    # 3d tests
    assert s("0,3,2/-3,3,7/-0.5,0.7,17") == [[0., 3, 2],
                                             [-3, 3, 7],
                                             [-0.5, 0.7, 17]]

    # spaces
    assert s("0,   10,   11") == [[0., 10., 11]]

    # check error
    with pytest.raises(ValueError):
        assert s("0, 10, 11, 42")

def test_lattice_object():
    # can call with spec as in min, max, num:
    l1 = lattice.Lattice([[0., 10., 6]])
    # or with string
    l2 = lattice.Lattice("0.,10.,6")

    # should result in the same object
    assert np.allclose(l1.get_positions(), l2.get_positions())

    # only order 'C' and 'F' allowed
    l2 = lattice.Lattice("0.,10.,6", order='F')
    l2 = lattice.Lattice("0.,10.,6", order='C')
    with pytest.raises(ValueError):
        l2 = lattice.Lattice("0.,10.,6", order='B')

def test_lattice_object_1d():
    l = lattice.Lattice([[0, 10, 6]])
    assert l.dim == 1

    check_points_index = []
    check_points_value = []

    def f(idx, x):
        print("idx = {} x={}".format(idx, x))
        check_points_index.append(idx[0])
        check_points_value.append(x[0])

    l.foreach(f)

    assert check_points_index == list(range(6))
    assert np.allclose(check_points_value, np.linspace(0, 10, 6))

    assert l.get_closest([1.01]) == [1]
    assert l.get_closest([0.99]) == [0]

    assert l.get_num_points() == 6

    assert l.get_pos_from_idx([4]) == [8]

    assert l.get_shape() == [6]

    assert l.max_node_pos == [10]
    assert l.min_node_pos == [0]

    assert l.min_max_num_list == [[0, 10, 6]]
    assert l.order == 'F'
    assert l.nodes == [6]
    assert str(l) == "Lattice([[0, 10, 6]])"
    assert l.stepsizes == [2]

    assert np.allclose(l.get_positions()[:, 0], np.linspace(0, 10, 6))

def test_lattice_object_1d_scales():
    l = lattice.Lattice([[0, 10, 6]])
    assert l.dim == 1

    l.scale(0.5)

    assert l.max_node_pos == [5]
    assert l.stepsizes == [1]

    l.scale(1 / 0.5)
    assert l.max_node_pos == [10]
    assert l.stepsizes == [2]

    # scale differently in different directions (silly for 1d)
    l.scale([0.5])

    assert l.max_node_pos == [5]
    assert l.stepsizes == [1]

    l.scale(1 / 0.5)
    assert l.max_node_pos == [10]
    assert l.stepsizes == [2]

def test_lattice_object_2d():
    l = lattice.Lattice([[0, 10, 6], [-3, 1, 2]])
    assert l.dim == 2

    check_points_index = []
    check_points_value = []

    def f(idx, x):
        print("idx = {} x={}".format(idx, x))
        check_points_index.append(idx[:])
        check_points_value.append(x[:])

    l.foreach(f)

    assert check_points_index == [[0, 0],
                                  [1, 0],
                                  [2, 0],
                                  [3, 0],
                                  [4, 0],
                                  [5, 0],
                                  [0, 1],
                                  [1, 1],
                                  [2, 1],
                                  [3, 1],
                                  [4, 1],
                                  [5, 1]]
    assert np.allclose(check_points_value, [[0.0, -3.0],
                                            [2.0, -3.0],
                                            [4.0, -3.0],
                                            [6.0, -3.0],
                                            [8.0, -3.0],
                                            [10.0, -3.0],
                                            [0.0, 1.0],
                                            [2.0, 1.0],
                                            [4.0, 1.0],
                                            [6.0, 1.0],
                                            [8.0, 1.0],
                                            [10.0, 1.0]])

    assert l.get_closest([1.01, -0.9]) == [1, 1]
    assert l.get_closest([1.01, -1.0001]) == [1, 0]
    assert l.get_num_points() == 12

    assert l.dim == 2

    with pytest.raises(IndexError):
        l.get_pos_from_idx([1])  # should raise IndexError for 2d mesh

    l.get_pos_from_idx([0, 0]) == [0., -3.]
    l.get_pos_from_idx([0, 1]) == [0., 1.]
    l.get_pos_from_idx([4, 1]) == [8., 1.]

    assert l.get_shape() == [6, 2]

    l.max_node_pos == [10, 1]

    assert l.min_max_num_list == [[0, 10, 6], [-3, 1, 2]]
    assert l.min_node_pos == [0, -3]
    assert l.nodes == [6, 2]
    assert l.order == 'F'
    assert str(l) == "Lattice([[0, 10, 6], [-3, 1, 2]])"

    assert np.allclose(l.get_positions(),
                       np.array([[[0., -3.],
                                  [2., -3.],
                                  [4., -3.],
                                  [6., -3.],
                                  [8., -3.],
                                  [10., -3.]],
                                 [[0., 1.],
                                  [2., 1.],
                                  [4., 1.],
                                  [6., 1.],
                                  [8., 1.],
                                  [10., 1.]]]))

def test_lattice_object_2d_scale():

    l = lattice.Lattice([[0, 10, 6], [-3, 1, 2]])
    assert l.dim == 2

    l.scale(2)
    assert l.max_node_pos == [20, 2]
    assert l.stepsizes == [4, 8]

    l.scale(0.5)
    assert l.max_node_pos == [10, 1]
    assert l.stepsizes == [2, 4]

def test_lattice_add():
    l1 = lattice.Lattice([[0, 10, 6]])
    l2 = lattice.Lattice([[0, 5, 3]])
    l_add = l1 + l2
    assert l1.dim == 1
    assert l2.dim == 1
    assert l_add.dim == 2

def test_field_lattice_init():
    fl = lattice.FieldLattice(lattice.Lattice([[0, 10, 6], [-3, 1, 2]]))

    # we want to see numpy.ndarray for the data
    assert type(fl.field_data) == np.ndarray
    # and contiguous data
    assert fl.field_data.flags['F_CONTIGUOUS'] is True
    # and the right shape
    assert fl.field_data.shape[1:] == tuple(fl.lattice.get_shape())
    # and 3d data vectors
    assert fl.field_data.shape[0] == 3 == fl.field_dim
    # but the lattice is in 2d only
    assert fl.lattice.dim == 2

def test_field_lattice_set():
    z_counter = [0]  # cheat to see in which order field vaules are set

    def mydata(pos):
        print("type of pos = {}".format(type(pos)))
        print("pos = {}".format(pos))
        # use position as field data (for our testing here)
        x, y = pos
        z = z_counter[0]
        z_counter[0] += 0.1
        return [x, y, z]

    fl = lattice.FieldLattice(lattice.Lattice([[0, 5, 6], [-2, -1, 2]]))
    fl.set(mydata)
    print(fl.field_data.shape)
    print(fl.field_data[:, :, :])
    print("==============")
    print(fl.field_data[2, :, :])
    print(np.array([np.linspace(0, 1.1, 12)]))
    print("==============")

    # check x-component
    assert np.allclose(fl.field_data[0, :, :],
                       np.array([[0., 0.],
                                 [1., 1.],
                                 [2., 2.],
                                 [3., 3.],
                                 [4., 4.],
                                 [5., 5.]]))

    # check y-component
    assert np.allclose(fl.field_data[1, :, :],
                       np.array([[-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.],
                                 [-2., -1.]]))

    # check z-component in Fortran ordering
    assert np.allclose(fl.field_data[2, :, :],
                       np.array([[0.0, 0.6],
                                 [0.1, 0.7],
                                 [0.2, 0.8],
                                 [0.3, 0.9],
                                 [0.4, 1.0],
                                 [0.5, 1.1]]))

    # How does C order differ?
    # Okay, simplify: establish only that C-order is different from F-order:
    flf = fl  # FieldLattice Fortran

    # Create FieldLattice C with identical data
    flc = lattice.FieldLattice(lattice.Lattice([[0, 5, 6], [-2, -1, 2]]),
                               order="C")
    z_counter[0] = 0
    fl.set(mydata)

    print("Fortran field_data:  {}".format(flf.field_data))
    print("C field_data:  {}".format(flc.field_data))

    # compare values (should be the same)

    # Hans: 2015: this passes on OS X (and travis?) but not on Osiris (Ubuntu 64)
    assert np.allclose(flc.field_data, flf.field_data)

    # check that the strides are different (check help(numpy.ndarray))
    assert flc.field_data.strides != flf.field_data.strides

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
This module provides methods for interacting with Hashicorp Vault via the HVAC
library.
"""
from __future__ import absolute_import

import logging
import inspect
import time
from functools import wraps
from datetime import datetime, timedelta

log = logging.getLogger(__name__)
EXCLUDED_HVAC_FUNCTIONS = ['initialize']

try:
    import hvac
    import requests
    DEPS_INSTALLED = True
except ImportError:
    log.debug('Unable to import the HVAC library.')
    DEPS_INSTALLED = False

__all__ = ['initialize', 'is_initialized']

class InsufficientParameters(Exception):
    pass

def __virtual__():
    return DEPS_INSTALLED

def _cache_client(client_builder):
    _client = []
    @wraps(client_builder)
    def get_client(*args, **kwargs):
        if not _client:
            _client.append(client_builder(*args, **kwargs))
        return _client[0]
    return get_client

@_cache_client
def _build_client(url='https://localhost:8200', token=None, cert=None,
                  verify=True, timeout=30, proxies=None, allow_redirects=True,
                  session=None):
    client_kwargs = locals()
    for k, v in client_kwargs.items():
        if k.startswith('_'):
            continue
        arg_val = __salt__['config.get']('vault.{key}'.format(key=k), v)
        log.debug('Setting {0} parameter for HVAC client to {1}.'
                  .format(k, arg_val))
        client_kwargs[k] = arg_val
    return hvac.Client(**client_kwargs)

def _bind_client(unbound_function):
    @wraps(unbound_function)
    def bound_function(*args, **kwargs):
        filtered_kwargs = {k: v for k, v in kwargs.items() if not k.startswith('_')}
        ignore_invalid = filtered_kwargs.pop('ignore_invalid', None)
        client = _build_client()
        try:
            return unbound_function(client, *args, **filtered_kwargs)
        except hvac.exceptions.InvalidRequest:
            if ignore_invalid:
                return None
            else:
                raise
    return bound_function

def _get_keybase_pubkey(username):
    """
    Return the base64 encoded public PGP key for a keybase user.
    """
    # Retrieve the text of the public key stored in Keybase
    user = requests.get('https://keybase.io/{username}/key.asc'.format(
        username=username))
    # Explicitly raise an exception if there is an HTTP error. No-op if no error
    user.raise_for_status()
    # Process the key to only include the contents and not the wrapping
    # contents (e.g. ----BEGIN PGP KEY---)
    user_key = user.text
    key_lines = user.text.strip('\n').split('\n')
    key_lines = key_lines[key_lines.index(''):-2]
    return ''.join(key_lines)

def _unseal(sealing_keys):
    client = _build_client()
    client.unseal_multi(sealing_keys)

def _rekey(secret_shares, secret_threshold, sealing_keys, pgp_keys, root_token):
    client = _build_client(token=root_token)
    rekey = client.start_rekey(secret_shares, secret_threshold, pgp_keys,
                               backup=True)
    client.rekey_multi(sealing_keys, nonce=rekey['nonce'])

def _wait_after_init(client, retries=5):
    '''This function will allow for a configurable delay before attempting
    to issue requests after an initialization. This is necessary because when
    running on an HA backend there is a short period where the Vault instance
    will be on standby while it acquires the lock.'''
    ready = False
    while retries > 0 and not ready:
        try:
            status = client.read('sys/health')
            ready = (status.get('initialized') and not status.get('sealed')
                     and not status.get('standby'))
        except hvac.exceptions.VaultError:
            pass
        if ready:
            break
        retries -= 1
        time.sleep(1)

def initialize(secret_shares=5, secret_threshold=3, pgp_keys=None,
               keybase_users=None, unseal=True):
    success = True
    if keybase_users and isinstance(keybase_users, list):
        keybase_keys = []
        for user in keybase_users:
            log.debug('Retrieving public keys for Keybase user {}.'
                      .format(user))
            keybase_keys.append(_get_keybase_pubkey(user))
        pgp_keys = pgp_keys or []
        pgp_keys.extend(keybase_keys)
    if pgp_keys and len(pgp_keys) < secret_shares:
        raise InsufficientParameters('The number of PGP keys does not match'
                                     ' the number of secret shares.')
    client = _build_client()
    try:
        if pgp_keys and not unseal:
            secrets = client.initialize(secret_shares, secret_threshold,
                                        pgp_keys)
        else:
            secrets = client.initialize(secret_shares, secret_threshold)
        sealing_keys = secrets['keys']
        root_token = secrets['root_token']
        if unseal:
            _wait_after_init(client)
            log.debug('Unsealing Vault with generated sealing keys.')
            _unseal(sealing_keys)
    except hvac.exceptions.VaultError as e:
        log.exception(e)
        success = False
        sealing_keys = None
    try:
        if pgp_keys and unseal:
            _wait_after_init(client)
            log.debug('Regenerating PGP encrypted keys and backing them up.')
            log.debug('PGP keys: {}'.format(pgp_keys))
            client.token = root_token
            _rekey(secret_shares, secret_threshold, sealing_keys,
                   pgp_keys, root_token)
            encrypted_sealing_keys = client.get_backed_up_keys()['keys']
            if encrypted_sealing_keys:
                sealing_keys = encrypted_sealing_keys
    except hvac.exceptions.VaultError as e:
        log.error('Vault was initialized but PGP encrypted keys were not able to'
                  ' be generated after unsealing.')
        log.debug('Failed to rekey and backup the sealing keys.')
        log.exception(e)
    client.token = root_token
    return success, sealing_keys, root_token

def scan_leases(prefix='', time_horizon=0):
    """Scan all leases and generate an event for any that are near expiration

    :param prefix: The prefix path of leases that you want to scan
    :param time_horizon: How far in advance you want to be alerted for expiring leases (seconds)
    :returns: List of lease info for leases expiring soon
    :rtype: list

    """
    client = _build_client()
    try:
        prefixes = client.list('sys/leases/lookup/{0}'.format(prefix))
    except hvac.exceptions.VaultError as e:
        log.exception('Failed to retrieve lease information for prefix %s', prefix)
        return []
    if prefixes:
        prefixes = prefixes.get('data', {}).get('keys', [])
    else:
        prefixes = []
    expiring_leases = []
    for node in prefixes:
        if node.endswith('/'):
            log.debug('Recursing into path %s for prefix %s', node, prefix)
            expiring_leases.extend(scan_leases('{0}/{1}'.format(prefix.strip('/'), node), time_horizon))
        else:
            log.debug('Retrieving lease information for %s/%s', prefix, node)
            try:
                lease_info = client.write('sys/leases/lookup', lease_id='{0}/{1}'.format(prefix.strip('/'), node))
            except hvac.exceptions.VaultError as e:
                log.exception('Failed to retrieve lease information for %s',
                              '{0}/{1}'.format(prefix.strip('/'), node))
                continue
            lease_expiry = datetime.strptime(lease_info.get('data', {}).get('expire_time')[:-4], '%Y-%m-%dT%H:%M:%S.%f')
            lease_lifetime = lease_expiry - datetime.utcnow()
            if lease_lifetime < timedelta(seconds=time_horizon):
                __salt__['event.send']('vault/lease/expiring/{0}/{1}'.format(prefix, node), data=lease_info.get('data', {}))
                expiring_leases.append(lease_info.get('data', {}))
    return expiring_leases

def clean_expired_leases(prefix='', time_horizon=0):
    """Scan all leases and delete any that have an expiration beyond the specified time horizon

    :param prefix: The prefix path of leases that you want to scan
    :param time_horizon: How far in advance you want to be alerted for expiring leases (seconds)
    :returns: List of lease info for leases that were deleted
    :rtype: list

    """
    client = _build_client()
    expired_leases = scan_leases(prefix, time_horizon)
    for index, lease in enumerate(expired_leases):
        try:
            client.write('sys/leases/revoke', lease_id=lease['id'])
        except hvac.exceptions.VaultError:
            log.exception('Failed to revoke lease %s', lease['id'])
            expired_leases.pop(index)
            continue
    return expired_leases

def _register_functions():
    method_dict = {}
    for method_name in dir(hvac.Client):
        if not method_name.startswith('_'):
            method = getattr(hvac.Client, method_name)
            if (not isinstance(method, property) and
                  method_name not in EXCLUDED_HVAC_FUNCTIONS):
                if method_name == 'list':
                    method_name = 'list_values'
                globals()[method_name] = _bind_client(method)

if DEPS_INSTALLED:
    _register_functions()
<commit_msg>Updated to prevent sending events when cleaning out expired leases<commit_after># -*- coding: utf-8 -*-
"""
This module provides methods for interacting with Hashicorp Vault via the HVAC
library.
"""
from __future__ import absolute_import

import logging
import inspect
import time
from functools import wraps
from datetime import datetime, timedelta

log = logging.getLogger(__name__)
EXCLUDED_HVAC_FUNCTIONS = ['initialize']

try:
    import hvac
    import requests
    DEPS_INSTALLED = True
except ImportError:
    log.debug('Unable to import the HVAC library.')
    DEPS_INSTALLED = False

__all__ = ['initialize', 'is_initialized']

class InsufficientParameters(Exception):
    pass

def __virtual__():
    return DEPS_INSTALLED

def _cache_client(client_builder):
    _client = []
    @wraps(client_builder)
    def get_client(*args, **kwargs):
        if not _client:
            _client.append(client_builder(*args, **kwargs))
        return _client[0]
    return get_client

@_cache_client
def _build_client(url='https://localhost:8200', token=None, cert=None,
                  verify=True, timeout=30, proxies=None, allow_redirects=True,
                  session=None):
    client_kwargs = locals()
    for k, v in client_kwargs.items():
        if k.startswith('_'):
            continue
        arg_val = __salt__['config.get']('vault.{key}'.format(key=k), v)
        log.debug('Setting {0} parameter for HVAC client to {1}.'
                  .format(k, arg_val))
        client_kwargs[k] = arg_val
    return hvac.Client(**client_kwargs)

def _bind_client(unbound_function):
    @wraps(unbound_function)
    def bound_function(*args, **kwargs):
        filtered_kwargs = {k: v for k, v in kwargs.items() if not k.startswith('_')}
        ignore_invalid = filtered_kwargs.pop('ignore_invalid', None)
        client = _build_client()
        try:
            return unbound_function(client, *args, **filtered_kwargs)
        except hvac.exceptions.InvalidRequest:
            if ignore_invalid:
                return None
            else:
                raise
    return bound_function

def _get_keybase_pubkey(username):
    """
    Return the base64 encoded public PGP key for a keybase user.
    """
    # Retrieve the text of the public key stored in Keybase
    user = requests.get('https://keybase.io/{username}/key.asc'.format(
        username=username))
    # Explicitly raise an exception if there is an HTTP error. No-op if no error
    user.raise_for_status()
    # Process the key to only include the contents and not the wrapping
    # contents (e.g. ----BEGIN PGP KEY---)
    user_key = user.text
    key_lines = user.text.strip('\n').split('\n')
    key_lines = key_lines[key_lines.index(''):-2]
    return ''.join(key_lines)

def _unseal(sealing_keys):
    client = _build_client()
    client.unseal_multi(sealing_keys)

def _rekey(secret_shares, secret_threshold, sealing_keys, pgp_keys, root_token):
    client = _build_client(token=root_token)
    rekey = client.start_rekey(secret_shares, secret_threshold, pgp_keys,
                               backup=True)
    client.rekey_multi(sealing_keys, nonce=rekey['nonce'])

def _wait_after_init(client, retries=5):
    '''This function will allow for a configurable delay before attempting
    to issue requests after an initialization. This is necessary because when
    running on an HA backend there is a short period where the Vault instance
    will be on standby while it acquires the lock.'''
    ready = False
    while retries > 0 and not ready:
        try:
            status = client.read('sys/health')
            ready = (status.get('initialized') and not status.get('sealed')
                     and not status.get('standby'))
        except hvac.exceptions.VaultError:
            pass
        if ready:
            break
        retries -= 1
        time.sleep(1)

def initialize(secret_shares=5, secret_threshold=3, pgp_keys=None,
               keybase_users=None, unseal=True):
    success = True
    if keybase_users and isinstance(keybase_users, list):
        keybase_keys = []
        for user in keybase_users:
            log.debug('Retrieving public keys for Keybase user {}.'
                      .format(user))
            keybase_keys.append(_get_keybase_pubkey(user))
        pgp_keys = pgp_keys or []
        pgp_keys.extend(keybase_keys)
    if pgp_keys and len(pgp_keys) < secret_shares:
        raise InsufficientParameters('The number of PGP keys does not match'
                                     ' the number of secret shares.')
    client = _build_client()
    try:
        if pgp_keys and not unseal:
            secrets = client.initialize(secret_shares, secret_threshold,
                                        pgp_keys)
        else:
            secrets = client.initialize(secret_shares, secret_threshold)
        sealing_keys = secrets['keys']
        root_token = secrets['root_token']
        if unseal:
            _wait_after_init(client)
            log.debug('Unsealing Vault with generated sealing keys.')
            _unseal(sealing_keys)
    except hvac.exceptions.VaultError as e:
        log.exception(e)
        success = False
        sealing_keys = None
    try:
        if pgp_keys and unseal:
            _wait_after_init(client)
            log.debug('Regenerating PGP encrypted keys and backing them up.')
            log.debug('PGP keys: {}'.format(pgp_keys))
            client.token = root_token
            _rekey(secret_shares, secret_threshold, sealing_keys,
                   pgp_keys, root_token)
            encrypted_sealing_keys = client.get_backed_up_keys()['keys']
            if encrypted_sealing_keys:
                sealing_keys = encrypted_sealing_keys
    except hvac.exceptions.VaultError as e:
        log.error('Vault was initialized but PGP encrypted keys were not able to'
                  ' be generated after unsealing.')
        log.debug('Failed to rekey and backup the sealing keys.')
        log.exception(e)
    client.token = root_token
    return success, sealing_keys, root_token

def scan_leases(prefix='', time_horizon=0, send_events=True):
    """Scan all leases and generate an event for any that are near expiration

    :param prefix: The prefix path of leases that you want to scan
    :param time_horizon: How far in advance you want to be alerted for expiring leases (seconds)
    :returns: List of lease info for leases expiring soon
    :rtype: list

    """
    client = _build_client()
    try:
        prefixes = client.list('sys/leases/lookup/{0}'.format(prefix))
    except hvac.exceptions.VaultError as e:
        log.exception('Failed to retrieve lease information for prefix %s', prefix)
        return []
    if prefixes:
        prefixes = prefixes.get('data', {}).get('keys', [])
    else:
        prefixes = []
    expiring_leases = []
    for node in prefixes:
        if node.endswith('/'):
            log.debug('Recursing into path %s for prefix %s', node, prefix)
            expiring_leases.extend(scan_leases('{0}/{1}'.format(prefix.strip('/'), node), time_horizon))
        else:
            log.debug('Retrieving lease information for %s/%s', prefix, node)
            try:
                lease_info = client.write('sys/leases/lookup', lease_id='{0}/{1}'.format(prefix.strip('/'), node))
            except hvac.exceptions.VaultError as e:
                log.exception('Failed to retrieve lease information for %s',
                              '{0}/{1}'.format(prefix.strip('/'), node))
                continue
            lease_expiry = datetime.strptime(lease_info.get('data', {}).get('expire_time')[:-4], '%Y-%m-%dT%H:%M:%S.%f')
            lease_lifetime = lease_expiry - datetime.utcnow()
            if lease_lifetime < timedelta(seconds=time_horizon):
                if send_events:
                    __salt__['event.send']('vault/lease/expiring/{0}/{1}'.format(prefix, node), data=lease_info.get('data', {}))
                expiring_leases.append(lease_info.get('data', {}))
    return expiring_leases

def clean_expired_leases(prefix='', time_horizon=0):
    """Scan all leases and delete any that have an expiration beyond the specified time horizon

    :param prefix: The prefix path of leases that you want to scan
    :param time_horizon: How far in advance you want to be alerted for expiring leases (seconds)
    :returns: List of lease info for leases that were deleted
    :rtype: list

    """
    client = _build_client()
    expired_leases = scan_leases(prefix, time_horizon, send_events=False)
    for index, lease in enumerate(expired_leases):
        try:
            client.write('sys/leases/revoke', lease_id=lease['id'])
        except hvac.exceptions.VaultError:
            log.exception('Failed to revoke lease %s', lease['id'])
            expired_leases.pop(index)
            continue
    return expired_leases

def _register_functions():
    method_dict = {}
    for method_name in dir(hvac.Client):
        if not method_name.startswith('_'):
            method = getattr(hvac.Client, method_name)
            if (not isinstance(method, property) and
                  method_name not in EXCLUDED_HVAC_FUNCTIONS):
                if method_name == 'list':
                    method_name = 'list_values'
                globals()[method_name] = _bind_client(method)

if DEPS_INSTALLED:
    _register_functions()
<|endoftext|>


<commit_before>import uuid

from nose.tools import eq_

from kazoo.testing import KazooTestCase

class KazooPartyTests(KazooTestCase):
    def setUp(self):
        super(KazooPartyTests, self).setUp()
        self.path = "/" + uuid.uuid4().hex

    def test_party(self):
        parties = [self.client.Party(self.path, "p%s" % i)
                   for i in range(5)]

        one_party = parties[0]

        eq_(list(one_party), [])
        eq_(len(one_party), 0)

        participants = set()
        for party in parties:
            party.join()
            participants.add(party.data)

            eq_(set(party), participants)
            eq_(len(party), len(participants))

        for party in parties:
            party.leave()
            participants.remove(party.data)

            eq_(set(party), participants)
            eq_(len(party), len(participants))

    def test_party_reuse_node(self):
        party = self.client.Party(self.path, "p1")
        self.client.ensure_path(self.path)
        self.client.create(party.create_path)
        party.join()
        self.assertTrue(party.participating)
        party.leave()
        self.assertFalse(party.participating)
        self.assertEqual(len(party), 0)

    def test_party_vanishing_node(self):
        party = self.client.Party(self.path, "p1")
        party.join()
        self.assertTrue(party.participating)
        self.client.delete(party.create_path)
        party.leave()
        self.assertFalse(party.participating)
        self.assertEqual(len(party), 0)

class KazooShallowPartyTests(KazooTestCase):
    def setUp(self):
        super(KazooShallowPartyTests, self).setUp()
        self.path = "/" + uuid.uuid4().hex

    def test_party(self):
        parties = [self.client.ShallowParty(self.path, "p%s" % i)
                   for i in range(5)]

        one_party = parties[0]

        eq_(list(one_party), [])
        eq_(len(one_party), 0)

        participants = set()
        for party in parties:
            party.join()
            participants.add(party.data)

            eq_(set(party), participants)
            eq_(len(party), len(participants))

        for party in parties:
            party.leave()
            participants.remove(party.data)

            eq_(set(party), participants)
            eq_(len(party), len(participants))
<commit_msg>more party.data bytes fixes in tests<commit_after>import uuid

from nose.tools import eq_

from kazoo.testing import KazooTestCase

class KazooPartyTests(KazooTestCase):
    def setUp(self):
        super(KazooPartyTests, self).setUp()
        self.path = "/" + uuid.uuid4().hex

    def test_party(self):
        parties = [self.client.Party(self.path, "p%s" % i)
                   for i in range(5)]

        one_party = parties[0]

        eq_(list(one_party), [])
        eq_(len(one_party), 0)

        participants = set()
        for party in parties:
            party.join()
            participants.add(party.data.decode('utf-8'))

            eq_(set(party), participants)
            eq_(len(party), len(participants))

        for party in parties:
            party.leave()
            participants.remove(party.data.decode('utf-8'))

            eq_(set(party), participants)
            eq_(len(party), len(participants))

    def test_party_reuse_node(self):
        party = self.client.Party(self.path, "p1")
        self.client.ensure_path(self.path)
        self.client.create(party.create_path)
        party.join()
        self.assertTrue(party.participating)
        party.leave()
        self.assertFalse(party.participating)
        self.assertEqual(len(party), 0)

    def test_party_vanishing_node(self):
        party = self.client.Party(self.path, "p1")
        party.join()
        self.assertTrue(party.participating)
        self.client.delete(party.create_path)
        party.leave()
        self.assertFalse(party.participating)
        self.assertEqual(len(party), 0)

class KazooShallowPartyTests(KazooTestCase):
    def setUp(self):
        super(KazooShallowPartyTests, self).setUp()
        self.path = "/" + uuid.uuid4().hex

    def test_party(self):
        parties = [self.client.ShallowParty(self.path, "p%s" % i)
                   for i in range(5)]

        one_party = parties[0]

        eq_(list(one_party), [])
        eq_(len(one_party), 0)

        participants = set()
        for party in parties:
            party.join()
            participants.add(party.data.decode('utf-8'))

            eq_(set(party), participants)
            eq_(len(party), len(participants))

        for party in parties:
            party.leave()
            participants.remove(party.data.decode('utf-8'))

            eq_(set(party), participants)
            eq_(len(party), len(participants))
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import print_function
import os
import json
import sys
from .common import epsilon
from .common import floatx
from .common import set_epsilon
from .common import set_floatx
from .common import get_uid
from .common import cast_to_floatx
from .common import image_dim_ordering
from .common import set_image_dim_ordering
from .common import is_keras_tensor
from .common import legacy_weight_ordering
from .common import set_legacy_weight_ordering

_keras_base_dir = os.path.expanduser('~')
if not os.access(_keras_base_dir, os.W_OK):
    _keras_base_dir = '/tmp'

_keras_dir = os.path.join(_keras_base_dir, '.keras')
if not os.path.exists(_keras_dir):
    os.makedirs(_keras_dir)

_BACKEND = 'tensorflow'
_config_path = os.path.expanduser(os.path.join(_keras_dir, 'keras.json'))
if os.path.exists(_config_path):
    _config = json.load(open(_config_path))
    _floatx = _config.get('floatx', floatx())
    assert _floatx in {'float16', 'float32', 'float64'}
    _epsilon = _config.get('epsilon', epsilon())
    assert type(_epsilon) == float
    _backend = _config.get('backend', _BACKEND)
    assert _backend in {'theano', 'tensorflow'}
    _image_dim_ordering = _config.get('image_dim_ordering', image_dim_ordering())
    assert _image_dim_ordering in {'tf', 'th'}

    set_floatx(_floatx)
    set_epsilon(_epsilon)
    set_image_dim_ordering(_image_dim_ordering)
    _BACKEND = _backend

# save config file
if not os.path.exists(_config_path):
    _config = {'floatx': floatx(),
               'epsilon': epsilon(),
               'backend': _BACKEND,
               'image_dim_ordering': image_dim_ordering()}
    with open(_config_path, 'w') as f:
        f.write(json.dumps(_config, indent=4))

if 'KERAS_BACKEND' in os.environ:
    _backend = os.environ['KERAS_BACKEND']
    assert _backend in {'theano', 'tensorflow'}
    _BACKEND = _backend

# import backend
if _BACKEND == 'theano':
    sys.stderr.write('Using Theano backend.\n')
    from .theano_backend import *
elif _BACKEND == 'tensorflow':
    sys.stderr.write('Using TensorFlow backend.\n')
    from .tensorflow_backend import *
else:
    raise Exception('Unknown backend: ' + str(_BACKEND))

def backend():
    '''Publicly accessible method
    for determining the current backend.
    '''
    return _BACKEND
<commit_msg>Set theano as default backend for windows users (#3831)<commit_after>from __future__ import absolute_import
from __future__ import print_function
import os
import json
import sys
from .common import epsilon
from .common import floatx
from .common import set_epsilon
from .common import set_floatx
from .common import get_uid
from .common import cast_to_floatx
from .common import image_dim_ordering
from .common import set_image_dim_ordering
from .common import is_keras_tensor
from .common import legacy_weight_ordering
from .common import set_legacy_weight_ordering

_keras_base_dir = os.path.expanduser('~')
if not os.access(_keras_base_dir, os.W_OK):
    _keras_base_dir = '/tmp'

_keras_dir = os.path.join(_keras_base_dir, '.keras')
if not os.path.exists(_keras_dir):
    os.makedirs(_keras_dir)

# Set theano as default backend for Windows users since tensorflow is not available for Windows yet.
if os.name == 'nt':
    _BACKEND = 'theano'
else:
    _BACKEND = 'tensorflow'

_config_path = os.path.expanduser(os.path.join(_keras_dir, 'keras.json'))
if os.path.exists(_config_path):
    _config = json.load(open(_config_path))
    _floatx = _config.get('floatx', floatx())
    assert _floatx in {'float16', 'float32', 'float64'}
    _epsilon = _config.get('epsilon', epsilon())
    assert type(_epsilon) == float
    _backend = _config.get('backend', _BACKEND)
    assert _backend in {'theano', 'tensorflow'}
    _image_dim_ordering = _config.get('image_dim_ordering', image_dim_ordering())
    assert _image_dim_ordering in {'tf', 'th'}

    set_floatx(_floatx)
    set_epsilon(_epsilon)
    set_image_dim_ordering(_image_dim_ordering)
    _BACKEND = _backend

# save config file
if not os.path.exists(_config_path):
    _config = {'floatx': floatx(),
               'epsilon': epsilon(),
               'backend': _BACKEND,
               'image_dim_ordering': image_dim_ordering()}
    with open(_config_path, 'w') as f:
        f.write(json.dumps(_config, indent=4))

if 'KERAS_BACKEND' in os.environ:
    _backend = os.environ['KERAS_BACKEND']
    assert _backend in {'theano', 'tensorflow'}
    _BACKEND = _backend

# import backend
if _BACKEND == 'theano':
    sys.stderr.write('Using Theano backend.\n')
    from .theano_backend import *
elif _BACKEND == 'tensorflow':
    sys.stderr.write('Using TensorFlow backend.\n')
    from .tensorflow_backend import *
else:
    raise Exception('Unknown backend: ' + str(_BACKEND))

def backend():
    '''Publicly accessible method
    for determining the current backend.
    '''
    return _BACKEND
<|endoftext|>


<commit_before>#! /usr/bin/env python
"""
Usage:
    rain -h | --help
    rain --version
    rain admin add [LABEL] --file=FILE
    rain admin on HOSTS
    rain admin off HOSTS
    rain admin [-i] delete HOSTS
    rain admin [-i] rm HOSTS
    rain admin list users
    rain admin list projects
    rain admin list roles
    rain admin list hosts [--user=USERS|--project=PROJECTS|--role=ROLE]
                          [--start=TIME_START]
                          [--end=TIME_END]
                          [--format=FORMAT]
    rain admin policy [--user=USERS|--project=PROJECTS|--role=ROLE]
                      (-l HOSTS|-n COUNT)
                      [--start=TIME_START]
                      [--end=TIME_END]
    rain list [--project=PROJECTS] [HOSTS]
    rain list hosts [--start=TIME_START]
                    [--end=TIME_END]
                    [--format=FORMAT]
    rain status [--short|--summary][--kind=KIND] [HOSTS]
    rain provision --profile=PROFILE HOSTS
    rain provision list (--distro=DISTRO|--kickstart=KICKSTART)
    rain provision --distro=DITRO --kickstart=KICKSTART HOSTS
    rain provision add (--distro=URL|--kickstart=KICk_CONTENT) NAME

Arguments:
    HOSTS     the list of hosts passed
    LABEL     the label of a host
    COUNT     the count of the bare metal provisioned hosts
    KIND      the kind

Options:
    -n COUNT     count of teh bare metal hosts to be provisined
    -p PROJECTS  --projects=PROJECTS
    -u USERS     --user=USERS        Specify users
    -f FILE, --file=FILE  file to be specified
    -i           interactive mode adds a yes/no
                 question for each host specified
    --role=ROLE            Specify predefined role
    --start=TIME_START     Start time of the reservation, in
                           YYYY/MM/DD HH:MM:SS format. [default: current_time]
    --end=TIME_END         End time of the reservation, in
                           YYYY/MM/DD HH:MM:SS format. In addition a duration
                           can be specified if the + sign is the first sign.
                           The duration will than be added to
                           the start time. [default: +1d]
    --kind=KIND            Format of the output -png, jpg, pdf. [default:png]
    --format=FORMAT        Format of the output json, cfg. [default:json]

"""
#from __future__ import dict
from docopt import docopt
import hostlist
from datetime import datetime, timedelta
from pytimeparse.timeparse import timeparse
# from timestring import Range
# from timestring import Date
from datetime import datetime, timedelta
import sys

def not_implemented():
    print "ERROR: not yet implemented"

def parse_time_interval (time_start, time_end):
    t_end = time_end
    t_start = time_start

    if t_start is not None:
        if t_start in ["current_time","now"]:
            t_start = str(datetime.now())

    if t_end is not None:
        if t_end.startswith("+"):
            duration = t_end[1:]
            delta = timeparse(duration)
            t_start = datetime.strptime(t_start, "%Y-%m-%d %H:%M:%S.%f")
            t_end = t_start + timedelta(seconds=delta)

    return (str(t_start), str(t_end))

def rain_command(arguments):

    for list in ["HOSTS", "USERS", "PROJECTS","--project", "--user"]:
        try:
            expanded_list = hostlist.expand_hostlist(arguments[list])
            arguments[list]=expanded_list
        except:
            pass

    print(arguments)

    """
    rain admin on HOSTS
    rain admin off HOSTS
    """
    if arguments["admin"]:

        if arguments["add"]:
            print "add"

            if arguments["LABEL"] is not None:
                """admin add LABEL --file=FILE"""

                print(arguments["LABEL"])
                print(arguments["--file"])
                not_implemented()

            else:
                """admin add --file=FILE"""

                print(arguments["--file"])
                not_implemented()

        elif arguments["on"]:
            """rain admin on HOSTS"""

            print "switch on"
            print (arguments["HOSTS"])
            not_implemented()

        elif arguments["off"]:
            """rain admin off HOSTS"""

            print "switch off"
            print (arguments["HOSTS"])
            not_implemented()

        elif arguments["delete"] or arguments["rm"] :
            """rain admin [-i] delete HOSTS"""
            """rain admin [-i] rm HOSTS"""

            interactive = arguments["-i"]

            print "delete", interactive

            for host in arguments["HOSTS"]:
                if interactive:
                    answer = raw_input("Do you want to delete the host %s? (y)es/(n)o: " % host)
                    if answer in ["yes","y","Y", "YES"]:
                        print "delete %s" % host
                    else:
                        print "keeping %s" % host
            not_implemented()

        elif arguments["list"]:

            if arguments["users"]:
                print "list users"
                not_implemented()

            elif arguments["projects"]:
                print "list projects"
                not_implemented()

            elif arguments["roles"]:
                print "list roles"
                not_implemented()

            elif arguments["hosts"]:
                print "list hosts"
                not_implemented()

                (time_start, time_end) = parse_time_interval(arguments["--start"],
                                                           arguments["--end"])
                print "From:", time_start
                print "To  :", time_end

                if ["--users"] is not None:
                    not_implemented()
                elif arguments["--projects"] is not None:
                    not_implemented()
                elif arguments["--role"] is not None:
                    not_implemented()
                else:
                    print ("all users, projects, roles")
                    not_implemented()

        elif arguments["policy"]:
            print "policy"

            (time_start, time_end) = parse_time_interval(arguments["--start"],
                                                         arguments["--end"])

            print "From:", time_start
            print "To  :", time_end

            if ["--users"] is not None:
                not_implemented()
            elif arguments["--projects"] is not None:
                not_implemented()
            elif arguments["--role"] is not None:
                not_implemented()
            else:
                print ("all users, projects, roles")
                not_implemented()

        elif arguments["list"]:
            print "list"

            not_implemented()

    elif arguments["status"]:
            print "status"
            not_implemented()

    elif arguments["list"]:
            print "user list"

            (time_start, time_end) = parse_time_interval(arguments["--start"],
                                                         arguments["--end"])

            print "From:", time_start
            print "To  :", time_end

            not_implemented()
    ###
    # provisioning
    ###
    elif arguments["provision"]:
        print "provision a node..."
        if arguments["list"]:
            print "this will list distro or kickstart info"
            not_implemented()
        elif arguments["add"]:
            print "add a new distro or kickstart"
            not_implemented()
        else:
            if arguments["--profile"]:
                profile = arguments["--profile"]
                if arguments["HOSTS"]:
                    #print arguments["HOSTS"]
                    #hosts = hostlist.expand_hostlist(arguments["HOSTS"])
                    hosts = arguments["HOSTS"]
                # call Heng's rest service for implementation
                # two steps:
                # 1. check and modify the 'system' in cobbler
                # 2. do the provision
                not_implemented()
                print "fake provision in progress..."
                for host in hosts:
                    print "now provision host: %s" % host

if __name__ == '__main__':
    arguments = docopt(__doc__)

    rain_command(arguments)

    """
    timeparse('1.2 minutes')32m

2h32m
3d2h32m
1w3d2h32m
1w 3d 2h 32m
1 w 3 d 2 h 32 m
4:13
4:13:02
4:13:02.266
2:04:13:02.266
2 days,  4:13:02 (uptime format)
2 days,  4:13:02.266
5hr34m56s
5 hours, 34 minutes, 56 seconds
5 hrs, 34 mins, 56 secs
2 days, 5 hours, 34 minutes, 56 seconds
1.2 m
1.2 min
1.2 mins
1.2 minute
1.2 minutes
172 hours
172 hr
172 h
172 hrs
172 hour
1.24 days
5 d
5 day
5 days
5.6 wk
5.6 week
5.6 weeks
"""

<commit_msg>adding the function registration for the rain command<commit_after>#! /usr/bin/env python
#from __future__ import dict
from docopt import docopt
import hostlist
from datetime import datetime, timedelta
from pytimeparse.timeparse import timeparse
# from timestring import Range
# from timestring import Date
from datetime import datetime, timedelta
import sys
from cloudmesh.util.util import parse_time_interval, not_implemented

def rain_command(arguments):
    """
    Usage:
        rain -h | --help
        rain --version
        rain admin add [LABEL] --file=FILE
        rain admin on HOSTS
        rain admin off HOSTS
        rain admin [-i] delete HOSTS
        rain admin [-i] rm HOSTS
        rain admin list users
        rain admin list projects
        rain admin list roles
        rain admin list hosts [--user=USERS|--project=PROJECTS|--role=ROLE]
                              [--start=TIME_START]
                              [--end=TIME_END]
                              [--format=FORMAT]
        rain admin policy [--user=USERS|--project=PROJECTS|--role=ROLE]
                          (-l HOSTS|-n COUNT)
                          [--start=TIME_START]
                          [--end=TIME_END]
        rain list [--project=PROJECTS] [HOSTS]
        rain list hosts [--start=TIME_START]
                        [--end=TIME_END]
                        [--format=FORMAT]
        rain status [--short|--summary][--kind=KIND] [HOSTS]
        rain provision --profile=PROFILE HOSTS
        rain provision list (--distro=DISTRO|--kickstart=KICKSTART)
        rain provision --distro=DITRO --kickstart=KICKSTART HOSTS
        rain provision add (--distro=URL|--kickstart=KICk_CONTENT) NAME

    Arguments:
        HOSTS     the list of hosts passed
        LABEL     the label of a host
        COUNT     the count of the bare metal provisioned hosts
        KIND      the kind

    Options:
        -n COUNT     count of teh bare metal hosts to be provisined
        -p PROJECTS  --projects=PROJECTS
        -u USERS     --user=USERS        Specify users
        -f FILE, --file=FILE  file to be specified
        -i           interactive mode adds a yes/no
                     question for each host specified
        --role=ROLE            Specify predefined role
        --start=TIME_START     Start time of the reservation, in
                               YYYY/MM/DD HH:MM:SS format. [default: current_time]
        --end=TIME_END         End time of the reservation, in
                               YYYY/MM/DD HH:MM:SS format. In addition a duration
                               can be specified if the + sign is the first sign.
                               The duration will than be added to
                               the start time. [default: +1d]
        --kind=KIND            Format of the output -png, jpg, pdf. [default:png]
        --format=FORMAT        Format of the output json, cfg. [default:json]

    """
    arguments = docopt(self.__doc__)
    for list in ["HOSTS", "USERS", "PROJECTS","--project", "--user"]:
        try:
            expanded_list = hostlist.expand_hostlist(arguments[list])
            arguments[list]=expanded_list
        except:
            pass

    print(arguments)

    """
    rain admin on HOSTS
    rain admin off HOSTS
    """
    if arguments["admin"]:

        if arguments["add"]:
            print "add"

            if arguments["LABEL"] is not None:
                """admin add LABEL --file=FILE"""

                print(arguments["LABEL"])
                print(arguments["--file"])
                not_implemented()

            else:
                """admin add --file=FILE"""

                print(arguments["--file"])
                not_implemented()

        elif arguments["on"]:
            """rain admin on HOSTS"""

            print "switch on"
            print (arguments["HOSTS"])
            not_implemented()

        elif arguments["off"]:
            """rain admin off HOSTS"""

            print "switch off"
            print (arguments["HOSTS"])
            not_implemented()

        elif arguments["delete"] or arguments["rm"] :
            """rain admin [-i] delete HOSTS"""
            """rain admin [-i] rm HOSTS"""

            interactive = arguments["-i"]

            print "delete", interactive

            for host in arguments["HOSTS"]:
                if interactive:
                    answer = raw_input("Do you want to delete the host %s? (y)es/(n)o: " % host)
                    if answer in ["yes","y","Y", "YES"]:
                        print "delete %s" % host
                    else:
                        print "keeping %s" % host
            not_implemented()

        elif arguments["list"]:

            if arguments["users"]:
                print "list users"
                not_implemented()

            elif arguments["projects"]:
                print "list projects"
                not_implemented()

            elif arguments["roles"]:
                print "list roles"
                not_implemented()

            elif arguments["hosts"]:
                print "list hosts"
                not_implemented()

                (time_start, time_end) = parse_time_interval(arguments["--start"],
                                                           arguments["--end"])
                print "From:", time_start
                print "To  :", time_end

                if ["--users"] is not None:
                    not_implemented()
                elif arguments["--projects"] is not None:
                    not_implemented()
                elif arguments["--role"] is not None:
                    not_implemented()
                else:
                    print ("all users, projects, roles")
                    not_implemented()

        elif arguments["policy"]:
            print "policy"

            (time_start, time_end) = parse_time_interval(arguments["--start"],
                                                         arguments["--end"])

            print "From:", time_start
            print "To  :", time_end

            if ["--users"] is not None:
                not_implemented()
            elif arguments["--projects"] is not None:
                not_implemented()
            elif arguments["--role"] is not None:
                not_implemented()
            else:
                print ("all users, projects, roles")
                not_implemented()

        elif arguments["list"]:
            print "list"

            not_implemented()

    elif arguments["status"]:
            print "status"
            not_implemented()

    elif arguments["list"]:
            print "user list"

            (time_start, time_end) = parse_time_interval(arguments["--start"],
                                                         arguments["--end"])

            print "From:", time_start
            print "To  :", time_end

            not_implemented()
    ###
    # provisioning
    ###
    elif arguments["provision"]:
        print "provision a node..."
        if arguments["list"]:
            print "this will list distro or kickstart info"
            not_implemented()
        elif arguments["add"]:
            print "add a new distro or kickstart"
            not_implemented()
        else:
            if arguments["--profile"]:
                profile = arguments["--profile"]
                if arguments["HOSTS"]:
                    #print arguments["HOSTS"]
                    #hosts = hostlist.expand_hostlist(arguments["HOSTS"])
                    hosts = arguments["HOSTS"]
                # call Heng's rest service for implementation
                # two steps:
                # 1. check and modify the 'system' in cobbler
                # 2. do the provision
                not_implemented()
                print "fake provision in progress..."
                for host in hosts:
                    print "now provision host: %s" % host

if __name__ == '__main__':
    rain_command(arguments)

    """
    timeparse('1.2 minutes')32m

2h32m
3d2h32m
1w3d2h32m
1w 3d 2h 32m
1 w 3 d 2 h 32 m
4:13
4:13:02
4:13:02.266
2:04:13:02.266
2 days,  4:13:02 (uptime format)
2 days,  4:13:02.266
5hr34m56s
5 hours, 34 minutes, 56 seconds
5 hrs, 34 mins, 56 secs
2 days, 5 hours, 34 minutes, 56 seconds
1.2 m
1.2 min
1.2 mins
1.2 minute
1.2 minutes
172 hours
172 hr
172 h
172 hrs
172 hour
1.24 days
5 d
5 day
5 days
5.6 wk
5.6 week
5.6 weeks
"""

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import logging

from scrapy import signals

from ..util import save_stats
from ..model import update_spider_stats

logger = logging.getLogger(__name__)

class ExtensionStats:
    def __init__(self, stats):
        self.stats = stats

    @classmethod
    def from_crawler(cls, crawler):
        ext = cls(crawler.stats)
        crawler.signals.connect(ext.spider_opened,
                                signal=signals.spider_opened)
        crawler.signals.connect(ext.spider_closed,
                                signal=signals.spider_closed)
        crawler.signals.connect(ext.item_scraped,
                                signal=signals.item_scraped)

        return ext

    def spider_opened(self, spider):
        self.stats.set_value(spider._id,
                             0)

    def spider_closed(self, spider):
        value = self.stats.get_value(spider._id)
        save_stats(spider.settings['SPIDER_STATS_URL'],
                   spider._id,
                   value)
        if spider.settings['BOT_NAME'] != 'TestSpider':
            logger.info('spider[%s] crawled %d articles',
                        spider.name,
                        value)
            if value == 0:
                update_spider_stats(spider,
                                    {'fail': 1})

    def item_scraped(self, item, spider):
        self.stats.inc_value(spider._id)
<commit_msg>log stats of corestats<commit_after># -*- coding: utf-8 -*-

import logging

from scrapy import signals

from ..util import save_stats
from ..model import update_spider_stats

logger = logging.getLogger(__name__)

class ExtensionStats:
    def __init__(self, stats):
        self.stats = stats

    @classmethod
    def from_crawler(cls, crawler):
        ext = cls(crawler.stats)
        crawler.signals.connect(ext.spider_opened,
                                signal=signals.spider_opened)
        crawler.signals.connect(ext.spider_closed,
                                signal=signals.spider_closed)
        crawler.signals.connect(ext.item_scraped,
                                signal=signals.item_scraped)

        return ext

    def spider_opened(self, spider):
        pass

    def item_scraped(self, item, spider):
        pass

    def spider_closed(self, spider):
        value = self.stats.get_value('item_scraped_count',
                                     0)
        save_stats(spider.settings['SPIDER_STATS_URL'],
                   spider._id,
                   value)
        if spider.settings['BOT_NAME'] != 'TestSpider':
            logger.info('spider[%s] crawled %d articles',
                        spider.name,
                        value)
            if value == 0:
                update_spider_stats(spider,
                                    {'fail': 1})
<|endoftext|>