<commit_before>
__all__ = ['ilr',
           'quantile',
           'regress_poly',
           'regress_components',
           'get_average_of_timeseries',
           'compcor',
           'bandpass_filter_matrix' ]

import numpy as np
from numpy.polynomial import Legendre
from scipy import linalg
from scipy.stats import pearsonr
import pandas as pd
from pandas import DataFrame
from sklearn import linear_model
import statsmodels.api as sm
import statsmodels.formula.api as smf

from .. import utils
from .. import core

def ilr( data_frame, voxmats, ilr_formula, verbose = False ):
    """
    Image-based linear regression.

    This function simplifies calculating p-values from linear models
    in which there is a similar formula that is applied many times
    with a change in image-based predictors.  Image-based variables
    are stored in the input matrix list. They should be named
    consistently in the input formula and in the image list.  If they
    are not, an error will be thrown.  All input matrices should have
    the same number of rows and columns.

    This function takes advantage of statsmodels R-style formulas.

    ANTsR function: `ilr`

    Arguments
    ---------

    data_frame: This data frame contains all relevant predictors except for
        the matrices associated with the image variables.  One should convert
        any categorical predictors ahead of time using `pd.get_dummies`.

    voxmats: The named list of matrices that contains the changing
        predictors.

    ilr_formula: This is a character string that defines a valid regression
        formula in the R-style.

    verbose:  will print a little bit of diagnostic information that allows
        a degree of model checking

    Returns
    -------

    A list of different matrices that contain names derived from the
    formula and the coefficients of the regression model.  The size of
    the output values ( p-values, t-values, parameter values ) will match
    the input matrix and, as such, can be converted to an image via `make_image`

    Example
    -------

    >>> nsub = 20
    >>> mu, sigma = 0, 1
    >>> outcome = np.random.normal( mu, sigma, nsub )
    >>> covar = np.random.normal( mu, sigma, nsub )
    >>> mat = np.random.normal( mu, sigma, (nsub, 500 ) )
    >>> mat2 = np.random.normal( mu, sigma, (nsub, 500 ) )
    >>> data = {'covar':covar,'outcome':outcome}
    >>> df = pd.DataFrame( data )
    >>> vlist = { "mat1": mat, "mat2": mat2 }
    >>> myform = " outcome ~ covar * mat1 "
    >>> result = ilr( df, vlist, myform)
    >>> myform = " mat2 ~ covar + mat1 "
    >>> result = ants.ilr( df, vlist, myform)

    """

    nvoxmats = len( voxmats )
    if nvoxmats < 1 :
        raise ValueError('Pass at least one matrix to voxmats list')
    keylist = list(voxmats.keys())
    firstmat = keylist[0]
    voxshape = voxmats[firstmat].shape
    nvox = voxshape[1]
    nmats = len( keylist )
    for k in keylist:
        if voxmats[firstmat].shape != voxmats[k].shape:
            raise ValueError('Matrices must have same number of rows (samples)')

    # test voxel
    vox = 0
    nrows = data_frame.shape[0]
    data_frame_vox = data_frame.copy()
    for k in range( nmats ):
        data = {keylist[k]: np.random.normal(0,1,nrows) }
        temp = pd.DataFrame( data )
        data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 )
    mod = smf.ols(formula=ilr_formula, data=data_frame_vox )
    res = mod.fit()
    modelNames = res.model.exog_names
    if verbose:
        print( data_frame_vox )
        print(res.summary())
    nOutcomes = len( modelNames )
    tValsOut = list()
    pValsOut = list()
    bValsOut = list()
    for k in range( len( modelNames ) ):
        bValsOut.append( np.zeros( nvox ) )
        pValsOut.append( np.zeros( nvox ) )
        tValsOut.append( np.zeros( nvox ) )

    for k in range( nvox ):
        # first get the correct data frame
        data_frame_vox = data_frame.copy()
        for v in range( nmats ):
            data = {keylist[v]: voxmats[keylist[v]][:,k] }
            temp = pd.DataFrame( data )
            data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 )
        # then get the local model results
        mod = smf.ols(formula=ilr_formula, data=data_frame_vox )
        res = mod.fit()
        tvals = res.tvalues
        pvals = res.pvalues
        bvals = res.params
        for v in range( len( modelNames ) ):
            bValsOut[v][k] = bvals[v]
            pValsOut[v][k] = pvals[v]
            tValsOut[v][k] = tvals[v]

    bValsOutDict = { }
    tValsOutDict = { }
    pValsOutDict = { }
    for v in range( len( modelNames ) ):
        bValsOutDict[ 'coef_' + modelNames[v] ] = bValsOut[v]
        tValsOutDict[ 'tval_' + modelNames[v] ] = tValsOut[v]
        pValsOutDict[ 'pval_' + modelNames[v] ] = pValsOut[v]

    return {
        'modelNames': modelNames,
        'coefficientValues': bValsOutDict,
        'pValues': pValsOutDict,
        'tValues': tValsOutDict }

def quantile(image, q, nonzero=True):
    """
    Get the quantile values from an ANTsImage
    """
    img_arr = image.numpy()
    if isinstance(q, (list,tuple)):
        q = [qq*100. if qq <= 1. else qq for qq in q]
        if nonzero:
            img_arr = img_arr[img_arr>0]
        vals = [np.percentile(img_arr, qq) for qq in q]
        return tuple(vals)
    elif isinstance(q, (float,int)):
        if q <= 1.:
            q = q*100.
        if nonzero:
            img_arr = img_arr[img_arr>0]
        return np.percentile(img_arr[img_arr>0], q)
    else:
        raise ValueError('q argument must be list/tuple or float/int')

def regress_poly(degree, data, remove_mean=True, axis=-1):
    """
    Returns data with degree polynomial regressed out.
    :param bool remove_mean: whether or not demean data (i.e. degree 0),
    :param int axis: numpy array axes along which regression is performed
    """
    timepoints = data.shape[0]
    # Generate design matrix
    X = np.ones((timepoints, 1))  # quick way to calc degree 0
    for i in range(degree):
        polynomial_func = Legendre.basis(i + 1)
        value_array = np.linspace(-1, 1, timepoints)
        X = np.hstack((X, polynomial_func(value_array)[:, np.newaxis]))
    non_constant_regressors = X[:, :-1] if X.shape[1] > 1 else np.array([])
    betas = np.linalg.pinv(X).dot(data)
    if remove_mean:
        datahat = X.dot(betas)
    else:  # disregard the first layer of X, which is degree 0
        datahat = X[:, 1:].dot(betas[1:, ...])
    regressed_data = data - datahat
    return regressed_data, non_constant_regressors

def regress_components( data, components, remove_mean=True ):
    """
    Returns data with components regressed out.
    :param bool remove_mean: whether or not demean data (i.e. degree 0),
    :param int axis: numpy array axes along which regression is performed
    """
    timepoints = data.shape[0]
    betas = np.linalg.pinv(components).dot(data)
    if remove_mean:
        datahat = components.dot(betas)
    else:  # disregard the first layer of X, which is degree 0
        datahat = components[:, 1:].dot(betas[1:, ...])
    regressed_data = data - datahat
    return regressed_data

def get_average_of_timeseries( image, idx=None ):
    """Average the timeseries into a dimension-1 image.
    image: input time series image
    idx:  indices over which to average
    """
    imagedim = image.dimension
    if idx is None:
        idx = range( image.shape[ imagedim - 1 ] )
    i0 = utils.slice_image( image, axis=image.dimension-1, idx=idx[0] ) * 0
    wt = 1.0 / len( idx )
    for k in idx:
        i0 = i0 + utils.slice_image( image, axis=image.dimension-1, idx=k ) * wt
    return( i0 )

def bandpass_filter_matrix( matrix,
    tr=1, lowf=0.01, highf=0.1, order = 3):
    """
    Bandpass filter the input time series image

    ANTsR function: `frequencyFilterfMRI`

    Arguments
    ---------

    image: input time series image

    tr:    sampling time interval (inverse of sampling rate)

    lowf:  low frequency cutoff

    highf: high frequency cutoff

    order: order of the butterworth filter run using `filtfilt`

    Returns
    -------
    filtered matrix

    Example
    -------

    >>> import numpy as np
    >>> import ants
    >>> import matplotlib.pyplot as plt
    >>> brainSignal = np.random.randn( 400, 1000 )
    >>> tr = 1
    >>> filtered = ants.bandpass_filter_matrix( brainSignal, tr = tr )
    >>> nsamples = brainSignal.shape[0]
    >>> t = np.linspace(0, tr*nsamples, nsamples, endpoint=False)
    >>> k = 20
    >>> plt.plot(t, brainSignal[:,k], label='Noisy signal')
    >>> plt.plot(t, filtered[:,k], label='Filtered signal')
    >>> plt.xlabel('time (seconds)')
    >>> plt.grid(True)
    >>> plt.axis('tight')
    >>> plt.legend(loc='upper left')
    >>> plt.show()
    """
    from scipy.signal import butter, filtfilt

    def butter_bandpass(lowcut, highcut, fs, order ):
        nyq = 0.5 * fs
        low = lowcut / nyq
        high = highcut / nyq
        b, a = butter(order, [low, high], btype='band')
        return b, a

    def butter_bandpass_filter(data, lowcut, highcut, fs, order ):
        b, a = butter_bandpass(lowcut, highcut, fs, order=order)
        y = filtfilt(b, a, data)
        return y

    fs = 1/tr   # sampling rate based on tr
    nsamples = matrix.shape[0]
    ncolumns = matrix.shape[1]
    matrixOut = matrix.copy()
    for k in range( ncolumns ):
        matrixOut[:,k] = butter_bandpass_filter(
            matrix[:,k], lowf, highf, fs, order=order )
    return matrixOut

def compcor( boldImage, ncompcor=4, quantile=0.975, mask=None, filter_type=False, degree=2 ):
    """
    Compute noise components from the input image

    ANTsR function: `compcor`

    this is adapted from nipy code https://github.com/nipy/nipype/blob/e29ac95fc0fc00fedbcaa0adaf29d5878408ca7c/nipype/algorithms/confounds.py

    Arguments
    ---------

    boldImage: input time series image

    ncompcor:  number of noise components to return

    quantile:  quantile defining high-variance

    mask:      mask defining brain or specific tissues

    filter_type: type off filter to apply to time series before computing
                 noise components.

        'polynomial' - Legendre polynomial basis
        False - None (mean-removal only)

    degree: order of polynomial used to remove trends from the timeseries

    Returns
    -------
    dictionary containing:

        components: a numpy array

        basis: a numpy array containing the (non-constant) filter regressors

    Example
    -------
    >>> cc = ants.compcor( ants.image_read(ants.get_ants_data("ch2")) )

    """

    def compute_tSTD(M, quantile, x=0, axis=0):
        stdM = np.std(M, axis=axis)
        # set bad values to x
        stdM[stdM == 0] = x
        stdM[np.isnan(stdM)] = x
        tt = round( quantile*100 )
        threshold_std = np.percentile( stdM, tt )
    #    threshold_std = quantile( stdM, quantile )
        return { 'tSTD': stdM, 'threshold_std': threshold_std}
    if mask is None:
        temp = utils.slice_image( boldImage, axis=boldImage.dimension-1, idx=0 )
        mask = utils.get_mask( temp )
    imagematrix = core.timeseries_to_matrix( boldImage, mask )
    temp = compute_tSTD( imagematrix, quantile, 0 )
    tsnrmask = core.make_image( mask,  temp['tSTD'] )
    tsnrmask = utils.threshold_image( tsnrmask, temp['threshold_std'], temp['tSTD'].max() )
    M = core.timeseries_to_matrix( boldImage, tsnrmask )
    components = None
    basis = np.array([])
    if filter_type in ('polynomial', False):
        M, basis = regress_poly(degree, M)
#    M = M / compute_tSTD(M, 1.)['tSTD']
    # "The covariance matrix C = MMT was constructed and decomposed into its
    # principal components using a singular value decomposition."
    u, _, _ = linalg.svd(M, full_matrices=False)
    if components is None:
        components = u[:, :ncompcor]
    else:
        components = np.hstack((components, u[:, :ncompcor]))
    if components is None and ncompcor > 0:
        raise ValueError('No components found')
    return { 'components': components, 'basis': basis }
<commit_msg>PERF: slightly faster<commit_after>
__all__ = ['ilr',
           'quantile',
           'regress_poly',
           'regress_components',
           'get_average_of_timeseries',
           'compcor',
           'bandpass_filter_matrix' ]

import numpy as np
from numpy.polynomial import Legendre
from scipy import linalg
from scipy.stats import pearsonr
import pandas as pd
from pandas import DataFrame
from sklearn import linear_model
import statsmodels.api as sm
import statsmodels.formula.api as smf

from .. import utils
from .. import core

def ilr( data_frame, voxmats, ilr_formula, verbose = False ):
    """
    Image-based linear regression.

    This function simplifies calculating p-values from linear models
    in which there is a similar formula that is applied many times
    with a change in image-based predictors.  Image-based variables
    are stored in the input matrix list. They should be named
    consistently in the input formula and in the image list.  If they
    are not, an error will be thrown.  All input matrices should have
    the same number of rows and columns.

    This function takes advantage of statsmodels R-style formulas.

    ANTsR function: `ilr`

    Arguments
    ---------

    data_frame: This data frame contains all relevant predictors except for
        the matrices associated with the image variables.  One should convert
        any categorical predictors ahead of time using `pd.get_dummies`.

    voxmats: The named list of matrices that contains the changing
        predictors.

    ilr_formula: This is a character string that defines a valid regression
        formula in the R-style.

    verbose:  will print a little bit of diagnostic information that allows
        a degree of model checking

    Returns
    -------

    A list of different matrices that contain names derived from the
    formula and the coefficients of the regression model.  The size of
    the output values ( p-values, t-values, parameter values ) will match
    the input matrix and, as such, can be converted to an image via `make_image`

    Example
    -------

    >>> nsub = 20
    >>> mu, sigma = 0, 1
    >>> outcome = np.random.normal( mu, sigma, nsub )
    >>> covar = np.random.normal( mu, sigma, nsub )
    >>> mat = np.random.normal( mu, sigma, (nsub, 500 ) )
    >>> mat2 = np.random.normal( mu, sigma, (nsub, 500 ) )
    >>> data = {'covar':covar,'outcome':outcome}
    >>> df = pd.DataFrame( data )
    >>> vlist = { "mat1": mat, "mat2": mat2 }
    >>> myform = " outcome ~ covar * mat1 "
    >>> result = ilr( df, vlist, myform)
    >>> myform = " mat2 ~ covar + mat1 "
    >>> result = ants.ilr( df, vlist, myform)

    """

    nvoxmats = len( voxmats )
    if nvoxmats < 1 :
        raise ValueError('Pass at least one matrix to voxmats list')
    keylist = list(voxmats.keys())
    firstmat = keylist[0]
    voxshape = voxmats[firstmat].shape
    nvox = voxshape[1]
    nmats = len( keylist )
    for k in keylist:
        if voxmats[firstmat].shape != voxmats[k].shape:
            raise ValueError('Matrices must have same number of rows (samples)')

    # test voxel
    vox = 0
    nrows = data_frame.shape[0]
    data_frame_vox = data_frame.copy()
    for k in range( nmats ):
        data = {keylist[k]: np.random.normal(0,1,nrows) }
        temp = pd.DataFrame( data )
        data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 )
    mod = smf.ols(formula=ilr_formula, data=data_frame_vox )
    res = mod.fit()
    modelNames = res.model.exog_names
    if verbose:
        print( data_frame_vox )
        print(res.summary())
    nOutcomes = len( modelNames )
    tValsOut = list()
    pValsOut = list()
    bValsOut = list()
    for k in range( len( modelNames ) ):
        bValsOut.append( np.zeros( nvox ) )
        pValsOut.append( np.zeros( nvox ) )
        tValsOut.append( np.zeros( nvox ) )

    data_frame_vox = data_frame.copy()
    for v in range( nmats ):
        data = {keylist[v]: voxmats[keylist[v]][:,k] }
        temp = pd.DataFrame( data )
        data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 )
    for k in range( nvox ):
        # first get the correct data frame
        for v in range( nmats ):
            data_frame_vox[ keylist[v] ] = voxmats[keylist[v]][:,k]
        # then get the local model results
        mod = smf.ols(formula=ilr_formula, data=data_frame_vox )
        res = mod.fit()
        tvals = res.tvalues
        pvals = res.pvalues
        bvals = res.params
        for v in range( len( modelNames ) ):
            bValsOut[v][k] = bvals[v]
            pValsOut[v][k] = pvals[v]
            tValsOut[v][k] = tvals[v]

    bValsOutDict = { }
    tValsOutDict = { }
    pValsOutDict = { }
    for v in range( len( modelNames ) ):
        bValsOutDict[ 'coef_' + modelNames[v] ] = bValsOut[v]
        tValsOutDict[ 'tval_' + modelNames[v] ] = tValsOut[v]
        pValsOutDict[ 'pval_' + modelNames[v] ] = pValsOut[v]

    return {
        'modelNames': modelNames,
        'coefficientValues': bValsOutDict,
        'pValues': pValsOutDict,
        'tValues': tValsOutDict }

def quantile(image, q, nonzero=True):
    """
    Get the quantile values from an ANTsImage
    """
    img_arr = image.numpy()
    if isinstance(q, (list,tuple)):
        q = [qq*100. if qq <= 1. else qq for qq in q]
        if nonzero:
            img_arr = img_arr[img_arr>0]
        vals = [np.percentile(img_arr, qq) for qq in q]
        return tuple(vals)
    elif isinstance(q, (float,int)):
        if q <= 1.:
            q = q*100.
        if nonzero:
            img_arr = img_arr[img_arr>0]
        return np.percentile(img_arr[img_arr>0], q)
    else:
        raise ValueError('q argument must be list/tuple or float/int')

def regress_poly(degree, data, remove_mean=True, axis=-1):
    """
    Returns data with degree polynomial regressed out.
    :param bool remove_mean: whether or not demean data (i.e. degree 0),
    :param int axis: numpy array axes along which regression is performed
    """
    timepoints = data.shape[0]
    # Generate design matrix
    X = np.ones((timepoints, 1))  # quick way to calc degree 0
    for i in range(degree):
        polynomial_func = Legendre.basis(i + 1)
        value_array = np.linspace(-1, 1, timepoints)
        X = np.hstack((X, polynomial_func(value_array)[:, np.newaxis]))
    non_constant_regressors = X[:, :-1] if X.shape[1] > 1 else np.array([])
    betas = np.linalg.pinv(X).dot(data)
    if remove_mean:
        datahat = X.dot(betas)
    else:  # disregard the first layer of X, which is degree 0
        datahat = X[:, 1:].dot(betas[1:, ...])
    regressed_data = data - datahat
    return regressed_data, non_constant_regressors

def regress_components( data, components, remove_mean=True ):
    """
    Returns data with components regressed out.
    :param bool remove_mean: whether or not demean data (i.e. degree 0),
    :param int axis: numpy array axes along which regression is performed
    """
    timepoints = data.shape[0]
    betas = np.linalg.pinv(components).dot(data)
    if remove_mean:
        datahat = components.dot(betas)
    else:  # disregard the first layer of X, which is degree 0
        datahat = components[:, 1:].dot(betas[1:, ...])
    regressed_data = data - datahat
    return regressed_data

def get_average_of_timeseries( image, idx=None ):
    """Average the timeseries into a dimension-1 image.
    image: input time series image
    idx:  indices over which to average
    """
    imagedim = image.dimension
    if idx is None:
        idx = range( image.shape[ imagedim - 1 ] )
    i0 = utils.slice_image( image, axis=image.dimension-1, idx=idx[0] ) * 0
    wt = 1.0 / len( idx )
    for k in idx:
        i0 = i0 + utils.slice_image( image, axis=image.dimension-1, idx=k ) * wt
    return( i0 )

def bandpass_filter_matrix( matrix,
    tr=1, lowf=0.01, highf=0.1, order = 3):
    """
    Bandpass filter the input time series image

    ANTsR function: `frequencyFilterfMRI`

    Arguments
    ---------

    image: input time series image

    tr:    sampling time interval (inverse of sampling rate)

    lowf:  low frequency cutoff

    highf: high frequency cutoff

    order: order of the butterworth filter run using `filtfilt`

    Returns
    -------
    filtered matrix

    Example
    -------

    >>> import numpy as np
    >>> import ants
    >>> import matplotlib.pyplot as plt
    >>> brainSignal = np.random.randn( 400, 1000 )
    >>> tr = 1
    >>> filtered = ants.bandpass_filter_matrix( brainSignal, tr = tr )
    >>> nsamples = brainSignal.shape[0]
    >>> t = np.linspace(0, tr*nsamples, nsamples, endpoint=False)
    >>> k = 20
    >>> plt.plot(t, brainSignal[:,k], label='Noisy signal')
    >>> plt.plot(t, filtered[:,k], label='Filtered signal')
    >>> plt.xlabel('time (seconds)')
    >>> plt.grid(True)
    >>> plt.axis('tight')
    >>> plt.legend(loc='upper left')
    >>> plt.show()
    """
    from scipy.signal import butter, filtfilt

    def butter_bandpass(lowcut, highcut, fs, order ):
        nyq = 0.5 * fs
        low = lowcut / nyq
        high = highcut / nyq
        b, a = butter(order, [low, high], btype='band')
        return b, a

    def butter_bandpass_filter(data, lowcut, highcut, fs, order ):
        b, a = butter_bandpass(lowcut, highcut, fs, order=order)
        y = filtfilt(b, a, data)
        return y

    fs = 1/tr   # sampling rate based on tr
    nsamples = matrix.shape[0]
    ncolumns = matrix.shape[1]
    matrixOut = matrix.copy()
    for k in range( ncolumns ):
        matrixOut[:,k] = butter_bandpass_filter(
            matrix[:,k], lowf, highf, fs, order=order )
    return matrixOut

def compcor( boldImage, ncompcor=4, quantile=0.975, mask=None, filter_type=False, degree=2 ):
    """
    Compute noise components from the input image

    ANTsR function: `compcor`

    this is adapted from nipy code https://github.com/nipy/nipype/blob/e29ac95fc0fc00fedbcaa0adaf29d5878408ca7c/nipype/algorithms/confounds.py

    Arguments
    ---------

    boldImage: input time series image

    ncompcor:  number of noise components to return

    quantile:  quantile defining high-variance

    mask:      mask defining brain or specific tissues

    filter_type: type off filter to apply to time series before computing
                 noise components.

        'polynomial' - Legendre polynomial basis
        False - None (mean-removal only)

    degree: order of polynomial used to remove trends from the timeseries

    Returns
    -------
    dictionary containing:

        components: a numpy array

        basis: a numpy array containing the (non-constant) filter regressors

    Example
    -------
    >>> cc = ants.compcor( ants.image_read(ants.get_ants_data("ch2")) )

    """

    def compute_tSTD(M, quantile, x=0, axis=0):
        stdM = np.std(M, axis=axis)
        # set bad values to x
        stdM[stdM == 0] = x
        stdM[np.isnan(stdM)] = x
        tt = round( quantile*100 )
        threshold_std = np.percentile( stdM, tt )
    #    threshold_std = quantile( stdM, quantile )
        return { 'tSTD': stdM, 'threshold_std': threshold_std}
    if mask is None:
        temp = utils.slice_image( boldImage, axis=boldImage.dimension-1, idx=0 )
        mask = utils.get_mask( temp )
    imagematrix = core.timeseries_to_matrix( boldImage, mask )
    temp = compute_tSTD( imagematrix, quantile, 0 )
    tsnrmask = core.make_image( mask,  temp['tSTD'] )
    tsnrmask = utils.threshold_image( tsnrmask, temp['threshold_std'], temp['tSTD'].max() )
    M = core.timeseries_to_matrix( boldImage, tsnrmask )
    components = None
    basis = np.array([])
    if filter_type in ('polynomial', False):
        M, basis = regress_poly(degree, M)
#    M = M / compute_tSTD(M, 1.)['tSTD']
    # "The covariance matrix C = MMT was constructed and decomposed into its
    # principal components using a singular value decomposition."
    u, _, _ = linalg.svd(M, full_matrices=False)
    if components is None:
        components = u[:, :ncompcor]
    else:
        components = np.hstack((components, u[:, :ncompcor]))
    if components is None and ncompcor > 0:
        raise ValueError('No components found')
    return { 'components': components, 'basis': basis }
<|endoftext|>


<commit_before><commit_msg>Reduce weather update frequency<commit_after><|endoftext|>


<commit_before>import sys
import os
import stat
from jinja2 import Environment, FileSystemLoader

def choice(question, options, default):
    "Ask the user to choose from a short list of named options"
    while True:
        sys.stdout.write("%s (%s) [%s]: " % (question, "/".join(options), default))
        answer = sys.stdin.readline().strip()
        if len(answer) == 0:
            return default
        for opt in options:
            if answer == opt:
                return answer

def yesno(question, default="y"):
    "Ask the user a yes/no question"
    while True:
        sys.stdout.write("%s (y/n) [%s]: " % (question, default))
        answer = sys.stdin.readline().strip().lower()
        if len(answer) == 0:
            answer = default
        if answer == 'y':
            return True
        elif answer == 'n':
            return False

def query(question, default=None):
    "Ask the user a question and return the response"
    while True:
        if default:
            sys.stdout.write("%s [%s]: " % (question, default))
        else:
            sys.stdout.write("%s: " % question)
        answer = sys.stdin.readline().strip().replace(' ', '_')
        if answer == "":
            if default:
                return default
        else:
            return answer

def optionlist(question, options, default=0):
    "Ask the user to choose from a list of options"
    for i, option in enumerate(options):
        print "%d: %s" % (i+1, option[0])
    while True:
        sys.stdout.write("%s (1-%d) [%d]: " % (question, len(options), default+1))
        answer = sys.stdin.readline().strip()
        if len(answer) == 0:
            return options[default][1]
        try:
            optno = int(answer)
            if optno > 0 and optno <= len(options):
                return options[optno-1][1]
        except:
            pass

class Workflow(object):
    def __init__(self, workflowdir, sharedir):
        self.jinja = Environment(loader=FileSystemLoader(sharedir), trim_blocks=True)
        self.name = os.path.basename(workflowdir)
        self.workflowdir = workflowdir
        self.sharedir = sharedir
        self.properties = {}
        self.home = os.environ["HOME"]
        self.user = os.environ["USER"]
        self.generate_tutorial = False
        self.tutorial_setup  = None
        self.compute_queue = "default"
        sysname, _, _, _, machine = os.uname()
        if sysname == 'Darwin':
            self.os = "MACOSX"
        else:
            # Probably Linux
            self.os = sysname.upper()
        self.arch = machine

    def copy_template(self, template, dest, mode=0644):
        "Copy template to dest in workflowdir with mode"
        path = os.path.join(self.workflowdir, dest)
        t = self.jinja.get_template(template)
        t.stream(**self.__dict__).dump(path)
        os.chmod(path, mode)

    def mkdir(self, path):
        "Make relative directory in workflowdir"
        path = os.path.join(self.workflowdir, path)
        if not os.path.exists(path):
            os.makedirs(path)

    def configure(self):
        # The tutorial is a special case
        if yesno("Do you want to generate a tutorial workflow?", "n"):
            self.config = "tutorial"
            self.daxgen = "tutorial"
            self.generate_tutorial = True

            # determine the environment to setup tutorial for
            self.tutorial_setup = optionlist("What environment is tutorial to be setup for?", [
                ("Local Machine", "submit-host"),
                ("USC HPCC Cluster", "usc-hpcc"),
                ("OSG from ISI submit node", "osg"),
                ("XSEDE, with Bosco", "xsede-bosco"),
                ("Bluewaters, with Glite", "bw-glite")
            ])

            # figure out what example options to provide
            examples = [
                ("Process", "process"),
                ("Pipeline", "pipeline"),
                ("Split", "split"),
                ("Merge", "merge")
            ]
            if self.tutorial_setup != "osg":
                examples.append( ("Diamond", "diamond"))

            if self.tutorial_setup == "bw-glite":
                examples.append( ("MPI Hello World", "y" ))

            self.tutorial = optionlist("What tutorial workflow do you want?", examples)

            self.setup_tutorial()
            return

        # Determine which DAX generator API to use
        self.daxgen = choice("What DAX generator API do you want to use?", ["python","perl","java"], "python")

        # Determine what kind of site catalog we need to generate
        self.config = optionlist("What does your computing infrastructure look like?", [
            ("Local Machine", "condorpool"),
            ("Remote Cluster using Globus GRAM", "globus"),
            ("Remote Cluster using CREAMCE", "creamce"),
            ("Local PBS Cluster with Glite", "glite"),
            ("Remote PBS Cluster with BOSCO and SSH", "bosco")
        ])

        # Find out some information about the site
        self.sitename = query("What do you want to call your compute site?", "compute")
        self.os = choice("What OS does your compute site have?", ["LINUX", "MACOSX"], self.os)
        self.arch = choice("What architecture does your compute site have?", ["x86_64", "x86"], self.arch)

    def setup_tutorial(self):
        """
        Set up tutorial for pre-defined computing environments
        :return:
        """
        if self.tutorial_setup is None:
            self.tutorial_setup = "submit-host"

        if self.tutorial_setup == "submit-host":
            self.sitename = "condorpool"
        elif self.tutorial_setup == "usc-hpcc":
            self.sitename = "usc-hpcc"
            self.config   = "glite"
            self.compute_queue = "default"
            # for running the whole workflow as mpi job
            self.properties["pegasus.job.aggregator"]="mpiexec"
        elif self.tutorial_setup == "osg":
            self.sitename = "osg"
            self.os = "linux"
            if not yesno("Do you want to use Condor file transfers", "y"):
                self.staging_site = "isi_workflow"
        elif self.tutorial_setup == "xsede-bosco":
            self.sitename = "condorpool"
        elif self.tutorial_setup == "bw-glite":
            self.sitename = "bluewaters"
            self.config   = "glite"
            self.compute_queue = "normal"
        return

    def generate(self):
        os.makedirs(self.workflowdir)
        self.mkdir("input")
        self.mkdir("output")

        if self.generate_tutorial:
            self.copy_template("%s/tc.txt" % self.tutorial, "tc.txt")
            self.copy_template("%s/daxgen.py" % self.tutorial, "daxgen.py")

            if self.tutorial == "diamond":

                # Executables used by the diamond workflow
                self.mkdir("bin")
                self.copy_template("diamond/transformation.py", "bin/preprocess", mode=0755)
                self.copy_template("diamond/transformation.py", "bin/findrange", mode=0755)
                self.copy_template("diamond/transformation.py", "bin/analyze", mode=0755)

                # Diamond input file
                self.copy_template("diamond/f.a", "input/f.a")
            elif self.tutorial == "split":
                # Split workflow input file
                self.copy_template("split/pegasus.html", "input/pegasus.html")
                self.copy_template("plan_cluster_dax.sh", "plan_cluster_dax.sh", mode=0755)
            elif self.tutorial == "mpi-hw":
                # copy the mpi wrapper, c code and mpi example
                # Executables used by the mpi-hw workflow
                self.mkdir("bin")
                self.copy_template( "%s/pegasus-mpi-hw.c" % self.tutorial , "pegasus-mpi-hw.c")
                self.copy_template( "%s/Makefile" % self.tutorial, "Makefile")
                self.copy_template( "%s/mpi-hello-world-wrapper" % self.tutorial, "bin/mpi-hello-world-wrapper", mode=0755)
                self.copy_template( "split/pegasus.html", "input/f.in")
        else:
            self.copy_template("tc.txt", "tc.txt")
            if self.daxgen == "python":
                self.copy_template("daxgen/daxgen.py", "daxgen.py")
            elif self.daxgen == "perl":
                self.copy_template("daxgen/daxgen.pl", "daxgen.pl")
            elif self.daxgen == "java":
                self.copy_template("daxgen/DAXGen.java", "DAXGen.java")
            else:
                assert False

        self.copy_template("sites.xml", "sites.xml")
        self.copy_template("plan_dax.sh", "plan_dax.sh", mode=0755)
        self.copy_template("generate_dax.sh", "generate_dax.sh", mode=0755)
        self.copy_template("README.md", "README.md")
        self.copy_template("rc.txt", "rc.txt")
        self.copy_template("pegasus.properties", "pegasus.properties")

        if self.generate_tutorial:
            sys.stdout.write("Pegasus Tutorial setup for example workflow - %s for execution on %s in directory %s\n"
                             % (self.tutorial, self.tutorial_setup, self.workflowdir))

def usage():
    print "Usage: %s WORKFLOW_DIR" % sys.argv[0]

def main(pegasus_share_dir):
    if len(sys.argv) != 2:
        usage()
        exit(1)

    if "-h" in sys.argv:
        usage()
        exit(1)

    workflowdir = sys.argv[1]
    if os.path.exists(workflowdir):
        print "ERROR: WORKFLOW_DIR '%s' already exists" % workflowdir
        exit(1)

    workflowdir = os.path.abspath(workflowdir)
    sharedir = os.path.join(pegasus_share_dir, "init")
    w = Workflow(workflowdir, sharedir)
    w.configure()
    w.generate()

<commit_msg>PM-1117 stupid typo<commit_after>import sys
import os
import stat
from jinja2 import Environment, FileSystemLoader

def choice(question, options, default):
    "Ask the user to choose from a short list of named options"
    while True:
        sys.stdout.write("%s (%s) [%s]: " % (question, "/".join(options), default))
        answer = sys.stdin.readline().strip()
        if len(answer) == 0:
            return default
        for opt in options:
            if answer == opt:
                return answer

def yesno(question, default="y"):
    "Ask the user a yes/no question"
    while True:
        sys.stdout.write("%s (y/n) [%s]: " % (question, default))
        answer = sys.stdin.readline().strip().lower()
        if len(answer) == 0:
            answer = default
        if answer == 'y':
            return True
        elif answer == 'n':
            return False

def query(question, default=None):
    "Ask the user a question and return the response"
    while True:
        if default:
            sys.stdout.write("%s [%s]: " % (question, default))
        else:
            sys.stdout.write("%s: " % question)
        answer = sys.stdin.readline().strip().replace(' ', '_')
        if answer == "":
            if default:
                return default
        else:
            return answer

def optionlist(question, options, default=0):
    "Ask the user to choose from a list of options"
    for i, option in enumerate(options):
        print "%d: %s" % (i+1, option[0])
    while True:
        sys.stdout.write("%s (1-%d) [%d]: " % (question, len(options), default+1))
        answer = sys.stdin.readline().strip()
        if len(answer) == 0:
            return options[default][1]
        try:
            optno = int(answer)
            if optno > 0 and optno <= len(options):
                return options[optno-1][1]
        except:
            pass

class Workflow(object):
    def __init__(self, workflowdir, sharedir):
        self.jinja = Environment(loader=FileSystemLoader(sharedir), trim_blocks=True)
        self.name = os.path.basename(workflowdir)
        self.workflowdir = workflowdir
        self.sharedir = sharedir
        self.properties = {}
        self.home = os.environ["HOME"]
        self.user = os.environ["USER"]
        self.generate_tutorial = False
        self.tutorial_setup  = None
        self.compute_queue = "default"
        sysname, _, _, _, machine = os.uname()
        if sysname == 'Darwin':
            self.os = "MACOSX"
        else:
            # Probably Linux
            self.os = sysname.upper()
        self.arch = machine

    def copy_template(self, template, dest, mode=0644):
        "Copy template to dest in workflowdir with mode"
        path = os.path.join(self.workflowdir, dest)
        t = self.jinja.get_template(template)
        t.stream(**self.__dict__).dump(path)
        os.chmod(path, mode)

    def mkdir(self, path):
        "Make relative directory in workflowdir"
        path = os.path.join(self.workflowdir, path)
        if not os.path.exists(path):
            os.makedirs(path)

    def configure(self):
        # The tutorial is a special case
        if yesno("Do you want to generate a tutorial workflow?", "n"):
            self.config = "tutorial"
            self.daxgen = "tutorial"
            self.generate_tutorial = True

            # determine the environment to setup tutorial for
            self.tutorial_setup = optionlist("What environment is tutorial to be setup for?", [
                ("Local Machine", "submit-host"),
                ("USC HPCC Cluster", "usc-hpcc"),
                ("OSG from ISI submit node", "osg"),
                ("XSEDE, with Bosco", "xsede-bosco"),
                ("Bluewaters, with Glite", "bw-glite")
            ])

            # figure out what example options to provide
            examples = [
                ("Process", "process"),
                ("Pipeline", "pipeline"),
                ("Split", "split"),
                ("Merge", "merge")
            ]
            if self.tutorial_setup != "osg":
                examples.append( ("Diamond", "diamond"))

            if self.tutorial_setup == "bw-glite":
                examples.append( ("MPI Hello World", "mpi-hw" ))

            self.tutorial = optionlist("What tutorial workflow do you want?", examples)

            self.setup_tutorial()
            return

        # Determine which DAX generator API to use
        self.daxgen = choice("What DAX generator API do you want to use?", ["python","perl","java"], "python")

        # Determine what kind of site catalog we need to generate
        self.config = optionlist("What does your computing infrastructure look like?", [
            ("Local Machine", "condorpool"),
            ("Remote Cluster using Globus GRAM", "globus"),
            ("Remote Cluster using CREAMCE", "creamce"),
            ("Local PBS Cluster with Glite", "glite"),
            ("Remote PBS Cluster with BOSCO and SSH", "bosco")
        ])

        # Find out some information about the site
        self.sitename = query("What do you want to call your compute site?", "compute")
        self.os = choice("What OS does your compute site have?", ["LINUX", "MACOSX"], self.os)
        self.arch = choice("What architecture does your compute site have?", ["x86_64", "x86"], self.arch)

    def setup_tutorial(self):
        """
        Set up tutorial for pre-defined computing environments
        :return:
        """
        if self.tutorial_setup is None:
            self.tutorial_setup = "submit-host"

        if self.tutorial_setup == "submit-host":
            self.sitename = "condorpool"
        elif self.tutorial_setup == "usc-hpcc":
            self.sitename = "usc-hpcc"
            self.config   = "glite"
            self.compute_queue = "default"
            # for running the whole workflow as mpi job
            self.properties["pegasus.job.aggregator"]="mpiexec"
        elif self.tutorial_setup == "osg":
            self.sitename = "osg"
            self.os = "linux"
            if not yesno("Do you want to use Condor file transfers", "y"):
                self.staging_site = "isi_workflow"
        elif self.tutorial_setup == "xsede-bosco":
            self.sitename = "condorpool"
        elif self.tutorial_setup == "bw-glite":
            self.sitename = "bluewaters"
            self.config   = "glite"
            self.compute_queue = "normal"
        return

    def generate(self):
        os.makedirs(self.workflowdir)
        self.mkdir("input")
        self.mkdir("output")

        if self.generate_tutorial:
            self.copy_template("%s/tc.txt" % self.tutorial, "tc.txt")
            self.copy_template("%s/daxgen.py" % self.tutorial, "daxgen.py")

            if self.tutorial == "diamond":

                # Executables used by the diamond workflow
                self.mkdir("bin")
                self.copy_template("diamond/transformation.py", "bin/preprocess", mode=0755)
                self.copy_template("diamond/transformation.py", "bin/findrange", mode=0755)
                self.copy_template("diamond/transformation.py", "bin/analyze", mode=0755)

                # Diamond input file
                self.copy_template("diamond/f.a", "input/f.a")
            elif self.tutorial == "split":
                # Split workflow input file
                self.copy_template("split/pegasus.html", "input/pegasus.html")
                self.copy_template("plan_cluster_dax.sh", "plan_cluster_dax.sh", mode=0755)
            elif self.tutorial == "mpi-hw":
                # copy the mpi wrapper, c code and mpi example
                # Executables used by the mpi-hw workflow
                self.mkdir("bin")
                self.copy_template( "%s/pegasus-mpi-hw.c" % self.tutorial , "pegasus-mpi-hw.c")
                self.copy_template( "%s/Makefile" % self.tutorial, "Makefile")
                self.copy_template( "%s/mpi-hello-world-wrapper" % self.tutorial, "bin/mpi-hello-world-wrapper", mode=0755)
                self.copy_template( "split/pegasus.html", "input/f.in")
        else:
            self.copy_template("tc.txt", "tc.txt")
            if self.daxgen == "python":
                self.copy_template("daxgen/daxgen.py", "daxgen.py")
            elif self.daxgen == "perl":
                self.copy_template("daxgen/daxgen.pl", "daxgen.pl")
            elif self.daxgen == "java":
                self.copy_template("daxgen/DAXGen.java", "DAXGen.java")
            else:
                assert False

        self.copy_template("sites.xml", "sites.xml")
        self.copy_template("plan_dax.sh", "plan_dax.sh", mode=0755)
        self.copy_template("generate_dax.sh", "generate_dax.sh", mode=0755)
        self.copy_template("README.md", "README.md")
        self.copy_template("rc.txt", "rc.txt")
        self.copy_template("pegasus.properties", "pegasus.properties")

        if self.generate_tutorial:
            sys.stdout.write("Pegasus Tutorial setup for example workflow - %s for execution on %s in directory %s\n"
                             % (self.tutorial, self.tutorial_setup, self.workflowdir))

def usage():
    print "Usage: %s WORKFLOW_DIR" % sys.argv[0]

def main(pegasus_share_dir):
    if len(sys.argv) != 2:
        usage()
        exit(1)

    if "-h" in sys.argv:
        usage()
        exit(1)

    workflowdir = sys.argv[1]
    if os.path.exists(workflowdir):
        print "ERROR: WORKFLOW_DIR '%s' already exists" % workflowdir
        exit(1)

    workflowdir = os.path.abspath(workflowdir)
    sharedir = os.path.join(pegasus_share_dir, "init")
    w = Workflow(workflowdir, sharedir)
    w.configure()
    w.generate()

<|endoftext|>


<commit_before>import yaml

#todo: refactor! experiment

def read_conf(file_path='../../conf/conf.yaml'):
    f = open(file_path, 'r+')
    file_conf = yaml.load(f.read())

    conf = file_conf['conf']
    specs = file_conf['specs']

    new_conf = []

    for pair in conf:
        reader = pair['reader']
        writer = pair['writer']
        new_reader = reader
        new_writer = writer
        if 'spec' in reader:
            spec = reader['spec']
            new_reader.pop('spec')
            new_reader.update(specs[spec])

        if 'spec' in writer:
            spec = writer['spec']
            new_writer.pop('spec')
            new_writer.update(specs[spec])
        new_pair = {'reader': new_reader, 'writer' : new_writer}
        new_conf.append(new_pair)
    return new_conf

if __name__ == '__main__':
    print read_conf()
<commit_msg>fail fast when spec not found in conf<commit_after>import yaml

#todo: refactor! experiment

def read_conf(file_path='../../conf/conf.yaml'):
    f = open(file_path, 'r+')
    file_conf = yaml.load(f.read())

    conf = file_conf['conf']
    specs = file_conf['specs']

    new_conf = []

    for pair in conf:
        reader = pair['reader']
        writer = pair['writer']
        new_reader = reader
        new_writer = writer
        if 'spec' in reader:
            spec = reader['spec']
            new_reader.pop('spec')
            if not spec in specs:
                print "spec '%s' not found in conf" % spec
                exit(-1)
            new_reader.update(specs[spec])

        if 'spec' in writer:
            spec = writer['spec']
            new_writer.pop('spec')
            if not spec in specs:
                print "spec '%s' not found in conf" % spec
                exit(-1)
            new_writer.update(specs[spec])
        new_pair = {'reader': new_reader, 'writer' : new_writer}
        new_conf.append(new_pair)
    return new_conf

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


<commit_before>"""Model definition for various aomi contexts"""

import re
import hvac
import aomi.exceptions
from aomi.helpers import hard_path, log, \
    warning, merge_dicts, cli_hash, random_word, \
    is_tagged
from aomi.vault import is_mounted
from aomi.error import output as error_output
from aomi.validation import sanitize_mount, check_obj

def wrap_vault(msg):
    def wrap_call(func):
        def func_wrapper(self, vault_client, opt):
            try:
                return func(self, vault_client, opt)
            except (hvac.exceptions.InvalidRequest,
                    hvac.exceptions.Forbidden) as vault_exception:
                if vault_exception.errors[0] == 'permission denied':
                    error_output("Permission denied %s from %s" %
                                 (msg, self.path), opt)
                else:
                    raise

        return func_wrapper
    return wrap_call

class Resource(object):
    required_fields = []
    config_key = None
    resource_key = None
    resource = 'vault resource'
    child = False

    def resources(self):
        return [self]

    def grok_state(self, obj):
        if 'state' in obj:
            my_state = obj['state'].lower()
            if my_state != 'absent' and my_state != 'present':
                raise aomi \
                    .exceptions \
                    .Validation('state must be either "absent" or "present"')

        self.present = obj.get('state', 'present').lower() == 'present'

    def validate(self, obj):
        if 'tags' in obj:
            if not isinstance(obj['tags'], list):
                raise aomi.exceptions.Validation('tags must be a list')

        if self.present:
            check_obj(self.required_fields, self.resource, obj)

    def __str__(self):
        return "%s %s" % (self.resource, self.path)

    def __init__(self, obj):
        self.grok_state(obj)
        self.validate(obj)
        self.path = None
        self.existing = None
        self.obj = {}
        self.tags = obj.get('tags', [])

    def fetch(self, vault_client, opt):
        result = self.read(vault_client, opt)
        if result:
            if 'data' in result:
                self.existing = result['data']
            else:
                self.existing = result
        else:
            self.existing = None

    def sync(self, vault_client, opt):
        if self.present and not self.existing:
            log("Writing new data to %s" % self, opt)
            self.write(vault_client, opt)
        elif self.present and self.existing:
            log("Updating data in %s" % self, opt)
            self.write(vault_client, opt)
        elif not self.present and not self.existing:
            log("No data to remove from %s" % self, opt)
        elif not self.present and self.existing:
            log("Removing data from %s" % self, opt)
            self.delete(vault_client, opt)

    def filtered(self, opt):
        if not is_tagged(opt.tags, self.tags):
            log("Skipping %s as it does not have requested tags" %
                self.path, opt)
            return False

        if not aomi.validation.specific_path_check(self.path, opt):
            log("Skipping %s as it does not match specified paths" %
                self.path, opt)
            return False

        return True

    @wrap_vault("reading")
    def read(self, client, opt):
        """Read from Vault while handling non surprising errors."""
        log("Reading from %s" % self, opt)
        return client.read(self.path)

    @wrap_vault("writing")
    def write(self, client, opt):
        """Write to Vault while handling non-surprising errors."""
        client.write(self.path, **self.obj)

    @wrap_vault("deleting")
    def delete(self, client, opt):
        """Delete from Vault while handling non-surprising errors."""
        log("Deleting %s" % self, opt)
        client.delete(self.path)

class Secret(Resource):
    config_key = 'secrets'

    def __init__(self, obj):
        super(Secret, self).__init__(obj)

def filtered_context(context, opt):
    ctx = Context()
    for resource in context.resources():
        if resource.child:
            continue
        if resource.filtered(opt):
            ctx.add(resource)

    return ctx

class Context(object):
    """The overall context of an aomi session"""

    def __init__(self):
        self._mounts = []
        self._resources = []
        self._auths = []
        self._logs = []

    def mounts(self):
        return self._mounts

    def logs(self):
        return self._logs

    def auths(self):
        return self._auths

    def resources(self):
        res = []
        for resource in self._resources:
            res = res + resource.resources()

        return res

    def find_backend(self, path, backends):
        for backend in backends:
            if backend.path == path:
                return backend

        return None

    def ensure_backend(self, resource, Backend, backends):
        existing_mount = self.find_backend(resource.mount, backends)
        if not existing_mount:
            new_mount = None
            if Backend == LogBackend:
                new_mount = Backend(resource)
            else:
                new_mount = Backend(resource.mount, resource.backend)

            backends.append(new_mount)
            return new_mount

        return existing_mount

    def add(self, resource):
        if isinstance(resource, Resource):
            if isinstance(resource, (Secret, Mount)):
                self.ensure_backend(resource, SecretBackend, self._mounts)
            elif isinstance(resource, (Auth)):
                self.ensure_backend(resource, AuthBackend, self._auths)
            elif isinstance(resource, (AuditLog)):
                self.ensure_backend(resource, LogBackend, self._logs)

            self._resources.append(resource)
        else:
            msg = "Unknown resource %s being " \
                  "added to context" % resource.__class__
            raise aomi.exceptions.AomiError(msg)

    def uses_mount(self, path):
        rsrcs = []
        for r in self._resources:
            if hasattr(r, 'mount') and r.mount == path:
                rsrcs.append(r)

        return rsrcs

    def remove(self, resource):
        if isinstance(resource, Resource):
            self._resources.remove(resource)

    def sync(self, vault_client, opt):
        for mount in self.mounts():
            if not mount.existing:
                mount.sync(vault_client, opt)
        for auth in self.auths():
            if not auth.existing:
                auth.sync(vault_client, opt)
        for blog in self.logs():
            if not blog.existing:
                blog.sync(vault_client, opt)
        for resource in self.resources():
            resource.sync(vault_client, opt)

    def fetch(self, vault_client, opt):
        for m in self.mounts():
            m.fetch(getattr(vault_client, SecretBackend.list_fun)())
        for a in self.auths():
            a.fetch(getattr(vault_client, AuthBackend.list_fun)())
        for l in self.logs():
            l.fetch(getattr(vault_client, LogBackend.list_fun)())

        for resource in self.resources():
            if issubclass(type(resource), aomi.model.Secret):
                if self.find_backend(resource.mount, self._mounts).existing:
                    resource.fetch(vault_client, opt)
            elif issubclass(type(resource), aomi.model.Auth):
                if self.find_backend(resource.mount, self._auths).existing:
                    resource.fetch(vault_client, opt)
            else:
                resource.fetch(vault_client, opt)

class Auth(Resource):
    def __init__(self, backend, obj):
        super(Auth, self).__init__(obj)
        self.backend = backend

    def backend_exists(self, client):
        """Checks to see if an auth backend exists yet"""
        backends = client.list_auth_backends().keys()
        backends = [x.rstrip('/') for x in backends]
        return self.backend in backends

    def ensure_auth(self, client):
        """Will ensure a particular auth endpoint is mounted"""
        if not self.backend_exists(client):
            client.enable_auth_backend(self.backend)

class Mount(Resource):
    resource = 'Generic Backend'
    required_fields = ['path']
    config_key = 'mounts'
    backend = 'generic'

    def __init__(self, obj, _opt):
        super(Mount, self).__init__(obj)
        self.mount = obj['path']

    def write(self, client, opt):
        pass

    def read(self, client, opt):
        pass

    def delete(self, client, opt):
        pass

class VaultBackend(object):
    backend = None
    list_fun = None
    unmount_fun = None
    mount_fun = None

    @staticmethod
    def list(vault_client, backend_class):
        return getattr(vault_client, backend_class.list_fun)()

    def __str__(self):
        if self.backend == self.path:
            return self.backend

        return "%s %s" % (self.backend, self.path)

    def __init__(self, path, backend):
        self.path = sanitize_mount(path)
        self.backend = backend
        self.existing = False

    def unmount(self, client):
        """Unmount a given mountpoint"""
        backends = getattr(client, self.list_fun)()
        if is_mounted(self.backend, self.path, backends):
            getattr(client, self.unmount_fun)(self.path)

    def sync(self, vault_client, opt):
        if not self.existing:
            self.actually_mount(vault_client)
            log("Mounting %s backend on %s" %
                (self.backend, self.path), opt)
        else:
            log("%s backend already mounted on %s" %
                (self.backend, self.path), opt)

    def fetch(self, backends):
        self.existing = is_mounted(self.backend, self.path, backends)

    def maybe_mount(self, client, opt):
        """Will ensure a mountpoint exists, or bail with a polite error"""
        backends = getattr(client, self.list_fun)()
        if not is_mounted(self.backend, self.path, backends):
            if self.backend == 'generic':
                log("Specifying a inline generic mountpoint is deprecated",
                    opt)

        self.actually_mount(client)

    def actually_mount(self, client):
        """Actually mount something in Vault"""
        try:
            getattr(client, self.mount_fun)(self.backend,
                                            mount_point=self.path)
        except hvac.exceptions.InvalidRequest as exception:
            match = re.match('existing mount at (?P<path>.+)', str(exception))
            if match:
                e_msg = "%s has a mountpoint conflict with %s" % \
                        (self.path, match.group('path'))
                raise aomi.exceptions.VaultConstraint(e_msg)
            else:
                raise

class SecretBackend(VaultBackend):
    list_fun = 'list_secret_backends'
    unmount_fun = 'disable_secret_backend'
    mount_fun = 'enable_secret_backend'

class AuthBackend(VaultBackend):
    list_fun = 'list_auth_backends'
    unmount_fun = 'disable_auth_backend'
    mount_fun = 'enable_auth_backend'

class LogBackend(VaultBackend):
    list_fun = 'list_audit_backends'
    unmount_fun = 'disable_audit_backend'
    mount_fun = 'enable_audit_backend'

    def __init__(self, resource):
        self.description = None
        super(LogBackend, self).__init__(resource.path, resource.backend)
        if resource.description:
            self.description = resource.description

        self.obj = resource.obj

    def actually_mount(self, client):
        if self.description:
            client.enable_audit_backend(self.backend,
                                        name=self.backend,
                                        description=self.description,
                                        options=self.obj)
        else:
            client.enable_audit_backend(self.backend,
                                        name=self.backend,
                                        options=self.obj)

class AuditLog(Resource):
    required_fields = ['type']
    resource = 'Audit Logs'
    config_key = 'audit_logs'

    def write(self, _vault_client, _opt):
        pass

    def __init__(self, log_obj, _opt):
        super(AuditLog, self).__init__(log_obj)
        self.backend = log_obj['type']
        self.mount = self.backend
        self.description = None
        self.path = log_obj.get('path', self.backend)
        obj = {
            'type': log_obj['type']
        }
        if self.backend == 'file':
            obj['file_path'] = log_obj['file_path']

        if self.backend == 'syslog':
            if 'tag' in log_obj:
                obj['tag'] = log_obj['tag']

            if 'facility' in log_obj:
                obj['facility'] = log_obj['facility']

        if 'description' in obj:
            self.description = obj['description']

        self.obj = obj
<commit_msg>aomi.model class docs<commit_after>"""Model definition for various aomi contexts"""

import re
import hvac
import aomi.exceptions
from aomi.helpers import log, is_tagged
from aomi.vault import is_mounted
from aomi.error import output as error_output
from aomi.validation import sanitize_mount, check_obj

def wrap_vault(msg):
    """Error catching Vault API wrapper
    This decorator wraps API interactions with Vault. It will
    catch and return appropriate error output on common
    problems"""
    # pylint: disable=missing-docstring
    def wrap_call(func):
        # pylint: disable=missing-docstring
        def func_wrapper(self, vault_client, opt):
            try:
                return func(self, vault_client, opt)
            except (hvac.exceptions.InvalidRequest,
                    hvac.exceptions.Forbidden) as vault_exception:
                if vault_exception.errors[0] == 'permission denied':
                    error_output("Permission denied %s from %s" %
                                 (msg, self.path), opt)
                else:
                    raise

        return func_wrapper
    return wrap_call

class Resource(object):
    """The base Vault Resource
    All aomi derived Vault resources should extend this
    class. It provides functionality for validation and
    API CRUD operations."""
    required_fields = []
    config_key = None
    resource_key = None
    resource = 'vault resource'
    child = False

    def resources(self):
        """List of included resources"""
        return [self]

    def grok_state(self, obj):
        """Determine the desired state of this
        resource based on data present"""
        if 'state' in obj:
            my_state = obj['state'].lower()
            if my_state != 'absent' and my_state != 'present':
                raise aomi \
                    .exceptions \
                    .Validation('state must be either "absent" or "present"')

        self.present = obj.get('state', 'present').lower() == 'present'

    def validate(self, obj):
        """Base validation method. Will inspect class attributes
        to termine just what should be present"""
        if 'tags' in obj and not isinstance(obj['tags'], list):
            raise aomi.exceptions.Validation('tags must be a list')

        if self.present:
            check_obj(self.required_fields, self.resource, obj)

    def __str__(self):
        return "%s %s" % (self.resource, self.path)

    def __init__(self, obj):
        self.grok_state(obj)
        self.validate(obj)
        self.path = None
        self.existing = None
        self.obj = {}
        self.tags = obj.get('tags', [])

    def fetch(self, vault_client, opt):
        """Populate internal representation of remote
        Vault resource contents"""
        result = self.read(vault_client, opt)
        if result:
            if 'data' in result:
                self.existing = result['data']
            else:
                self.existing = result
        else:
            self.existing = None

    def sync(self, vault_client, opt):
        """Update remove Vault resource contents if needed"""
        if self.present and not self.existing:
            log("Writing new data to %s" % self, opt)
            self.write(vault_client, opt)
        elif self.present and self.existing:
            log("Updating data in %s" % self, opt)
            self.write(vault_client, opt)
        elif not self.present and not self.existing:
            log("No data to remove from %s" % self, opt)
        elif not self.present and self.existing:
            log("Removing data from %s" % self, opt)
            self.delete(vault_client, opt)

    def filtered(self, opt):
        """Determines whether or not resource is filtered.
        Resources may be filtered if the tags do not match
        or the user has specified explict paths to include
        or exclude via command line options"""
        if not is_tagged(opt.tags, self.tags):
            log("Skipping %s as it does not have requested tags" %
                self.path, opt)
            return False

        if not aomi.validation.specific_path_check(self.path, opt):
            log("Skipping %s as it does not match specified paths" %
                self.path, opt)
            return False

        return True

    @wrap_vault("reading")
    def read(self, client, opt):
        """Read from Vault while handling non surprising errors."""
        log("Reading from %s" % self, opt)
        return client.read(self.path)

    @wrap_vault("writing")
    def write(self, client, _opt):
        """Write to Vault while handling non-surprising errors."""
        client.write(self.path, **self.obj)

    @wrap_vault("deleting")
    def delete(self, client, opt):
        """Delete from Vault while handling non-surprising errors."""
        log("Deleting %s" % self, opt)
        client.delete(self.path)

class Secret(Resource):
    """Vault Secrets
    These Vault resources will have some kind of secret backend
    underneath them. Seems to work with generic and AWS"""
    config_key = 'secrets'

def filtered_context(context, opt):
    """Filters a context
    This will return a new context with only the resources that
    are actually available for use. Uses tags and command line
    options to make determination."""

    ctx = Context()
    for resource in context.resources():
        if resource.child:
            continue
        if resource.filtered(opt):
            ctx.add(resource)

    return ctx

class Context(object):
    """The overall context of an aomi session"""

    def __init__(self):
        self._mounts = []
        self._resources = []
        self._auths = []
        self._logs = []

    def mounts(self):
        """Secret backends within context"""
        return self._mounts

    def logs(self):
        """Audit log backends within context"""
        return self._logs

    def auths(self):
        """Authentication backends within context"""
        return self._auths

    def resources(self):
        """Vault resources within context"""
        res = []
        for resource in self._resources:
            res = res + resource.resources()

        return res

    def find_backend(self, path, backends):
        for backend in backends:
            if backend.path == path:
                return backend

        return None

    def ensure_backend(self, resource, Backend, backends):
        existing_mount = self.find_backend(resource.mount, backends)
        if not existing_mount:
            new_mount = None
            if Backend == LogBackend:
                new_mount = Backend(resource)
            else:
                new_mount = Backend(resource.mount, resource.backend)

            backends.append(new_mount)
            return new_mount

        return existing_mount

    def add(self, resource):
        if isinstance(resource, Resource):
            if isinstance(resource, (Secret, Mount)):
                self.ensure_backend(resource, SecretBackend, self._mounts)
            elif isinstance(resource, (Auth)):
                self.ensure_backend(resource, AuthBackend, self._auths)
            elif isinstance(resource, (AuditLog)):
                self.ensure_backend(resource, LogBackend, self._logs)

            self._resources.append(resource)
        else:
            msg = "Unknown resource %s being " \
                  "added to context" % resource.__class__
            raise aomi.exceptions.AomiError(msg)

    def uses_mount(self, path):
        rsrcs = []
        for r in self._resources:
            if hasattr(r, 'mount') and r.mount == path:
                rsrcs.append(r)

        return rsrcs

    def remove(self, resource):
        if isinstance(resource, Resource):
            self._resources.remove(resource)

    def sync(self, vault_client, opt):
        for mount in self.mounts():
            if not mount.existing:
                mount.sync(vault_client, opt)
        for auth in self.auths():
            if not auth.existing:
                auth.sync(vault_client, opt)
        for blog in self.logs():
            if not blog.existing:
                blog.sync(vault_client, opt)
        for resource in self.resources():
            resource.sync(vault_client, opt)

    def fetch(self, vault_client, opt):
        for m in self.mounts():
            m.fetch(getattr(vault_client, SecretBackend.list_fun)())
        for a in self.auths():
            a.fetch(getattr(vault_client, AuthBackend.list_fun)())
        for l in self.logs():
            l.fetch(getattr(vault_client, LogBackend.list_fun)())

        for resource in self.resources():
            if issubclass(type(resource), aomi.model.Secret):
                if self.find_backend(resource.mount, self._mounts).existing:
                    resource.fetch(vault_client, opt)
            elif issubclass(type(resource), aomi.model.Auth):
                if self.find_backend(resource.mount, self._auths).existing:
                    resource.fetch(vault_client, opt)
            else:
                resource.fetch(vault_client, opt)

class Auth(Resource):
    def __init__(self, backend, obj):
        super(Auth, self).__init__(obj)
        self.backend = backend

    def backend_exists(self, client):
        """Checks to see if an auth backend exists yet"""
        backends = client.list_auth_backends().keys()
        backends = [x.rstrip('/') for x in backends]
        return self.backend in backends

    def ensure_auth(self, client):
        """Will ensure a particular auth endpoint is mounted"""
        if not self.backend_exists(client):
            client.enable_auth_backend(self.backend)

class Mount(Resource):
    resource = 'Generic Backend'
    required_fields = ['path']
    config_key = 'mounts'
    backend = 'generic'

    def __init__(self, obj, _opt):
        super(Mount, self).__init__(obj)
        self.mount = obj['path']

    def write(self, client, opt):
        pass

    def read(self, client, opt):
        pass

    def delete(self, client, opt):
        pass

class VaultBackend(object):
    backend = None
    list_fun = None
    unmount_fun = None
    mount_fun = None

    @staticmethod
    def list(vault_client, backend_class):
        return getattr(vault_client, backend_class.list_fun)()

    def __str__(self):
        if self.backend == self.path:
            return self.backend

        return "%s %s" % (self.backend, self.path)

    def __init__(self, path, backend):
        self.path = sanitize_mount(path)
        self.backend = backend
        self.existing = False

    def unmount(self, client):
        """Unmount a given mountpoint"""
        backends = getattr(client, self.list_fun)()
        if is_mounted(self.backend, self.path, backends):
            getattr(client, self.unmount_fun)(self.path)

    def sync(self, vault_client, opt):
        """Synchronizes the local and remote Vault resources. Has the net
        effect of adding backend if needed"""
        if not self.existing:
            self.actually_mount(vault_client)
            log("Mounting %s backend on %s" %
                (self.backend, self.path), opt)
        else:
            log("%s backend already mounted on %s" %
                (self.backend, self.path), opt)

    def fetch(self, backends):
        """Updates local resource with context on whether this
        backend is actually mounted and available"""
        self.existing = is_mounted(self.backend, self.path, backends)

    def maybe_mount(self, client, opt):
        """Will ensure a mountpoint exists, or bail with a polite error"""
        backends = getattr(client, self.list_fun)()
        if not is_mounted(self.backend, self.path, backends):
            if self.backend == 'generic':
                log("Specifying a inline generic mountpoint is deprecated",
                    opt)

        self.actually_mount(client)

    def actually_mount(self, client):
        """Actually mount something in Vault"""
        try:
            getattr(client, self.mount_fun)(self.backend,
                                            mount_point=self.path)
        except hvac.exceptions.InvalidRequest as exception:
            match = re.match('existing mount at (?P<path>.+)', str(exception))
            if match:
                e_msg = "%s has a mountpoint conflict with %s" % \
                        (self.path, match.group('path'))
                raise aomi.exceptions.VaultConstraint(e_msg)
            else:
                raise

class SecretBackend(VaultBackend):
    """Secret Backends for actual Vault resources"""
    list_fun = 'list_secret_backends'
    unmount_fun = 'disable_secret_backend'
    mount_fun = 'enable_secret_backend'

class AuthBackend(VaultBackend):
    """Authentication backends for Vault access"""
    list_fun = 'list_auth_backends'
    unmount_fun = 'disable_auth_backend'
    mount_fun = 'enable_auth_backend'

class LogBackend(VaultBackend):
    """Audit Log backends"""
    list_fun = 'list_audit_backends'
    unmount_fun = 'disable_audit_backend'
    mount_fun = 'enable_audit_backend'

    def __init__(self, resource):
        self.description = None
        super(LogBackend, self).__init__(resource.path, resource.backend)
        if resource.description:
            self.description = resource.description

        self.obj = resource.obj

    def actually_mount(self, client):
        if self.description:
            client.enable_audit_backend(self.backend,
                                        name=self.backend,
                                        description=self.description,
                                        options=self.obj)
        else:
            client.enable_audit_backend(self.backend,
                                        name=self.backend,
                                        options=self.obj)

class AuditLog(Resource):
    """Vault Resource for Audit Logging.
    Only supports syslog and file"""
    required_fields = ['type']
    resource = 'Audit Logs'
    config_key = 'audit_logs'

    def write(self, _vault_client, _opt):
        pass

    def __init__(self, log_obj, _opt):
        super(AuditLog, self).__init__(log_obj)
        self.backend = log_obj['type']
        self.mount = self.backend
        self.description = None
        self.path = log_obj.get('path', self.backend)
        obj = {
            'type': log_obj['type']
        }
        if self.backend == 'file':
            obj['file_path'] = log_obj['file_path']

        if self.backend == 'syslog':
            if 'tag' in log_obj:
                obj['tag'] = log_obj['tag']

            if 'facility' in log_obj:
                obj['facility'] = log_obj['facility']

        if 'description' in obj:
            self.description = obj['description']

        self.obj = obj
<|endoftext|>


<commit_before>#!/usr/bin/env python
# vim: sts=4 sw=4 et
#    This is a component of EMC
#    util.py Copyright 2010 Michael Haberler
#
#
#    This program 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.
#
#    This program 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 this program; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA'''
'''
    persistence support for gladevcp widgets
    Michael Haberler 11/2010
'''
import os
import sys
import time
import gtk
from configobj import ConfigObj, flatten_errors
from validate import Validator
from hashlib import sha1
from hal_widgets import _HalWidgetBase
from gladevcp.gladebuilder import widget_name

class UselessIniError(Exception):
    pass

class NoGetterSetterError(Exception):
    pass

class BadDescriptorDictError(Exception):
    pass

# map python type names to ConfigObj type names
co_map = {'str': 'string','int':'integer', 'str': 'string', 'bool': 'boolean'}
debug = 0
version_number = 1

def warn(*args):
    print >> sys.stderr,''.join(args)

def dbg(level,*args):
    global debug
    if debug < level: return
    print ''.join(args)

def set_debug(value):
    global debug
    debug = value

def select_widgets(widgets, hal_only=False,output_only = False):
    '''
    input: a list of widget instances
    output: a list of widget instances, pruned according to the flags:
        hal_only = True:  only add HAL widgets
        output_only = True: only add widgets which let users
        set a value (scale,button..)
    '''
    wlist = []
    for w in widgets:
        if hal_only and not isinstance(w, _HalWidgetBase):
            continue
        if output_only and not isinstance(w, (gtk.Range,
                                              gtk.SpinButton,
                                              gtk.ComboBox,
                                              gtk.CheckButton,
                                              gtk.ToggleButton,
                                              gtk.RadioButton)):
            continue
        wlist.append(w)
    return wlist

def accessors(w):
    '''
    retrieve getter/setter name of an 'interesting' widget
    '''
    if isinstance(w, (gtk.Range, gtk.SpinButton,gtk.ComboBox)):
        return ('get_value','set_value')
    if isinstance(w, (gtk.CheckButton, gtk.ToggleButton,gtk.RadioButton)):
        return ('get_active','set_active')

    raise NoGetterSetterError

def store_value(widget,value):
    try:
        (_,setter) = accessors(widget)
        method = getattr(widget,setter,None)
        method(value)
    except NoGetterSetterError:
        dbg(3,"store_value: no setter method for: ",str(widget))
        raise

def get_value(widget):
    try:
        (getter,_) = accessors(widget)
        method = getattr(widget,getter,None)
        return method()
    except NoGetterSetterError:
        dbg(3,"store_value: no getter method for: ",str(widget))
        raise

def widget_defaults(widgets):
    '''
    return a dict name/current value of wvalues which are to be restored
    '''
    wvalues = dict()
    for w in widgets:
        k = widget_name(w)
        try:
            v = get_value(w)
            wvalues[k] = v
        except NoGetterSetterError:
            continue
    return wvalues

class IniFile(object):

    # well known section and variable names in .ini files
    vars = 'vars'
    widgets = 'widgets'
    ini = 'ini'
    signature = 'signature'
    version = 'version'

    def _gen_spec(self,vdict):
        '''
        given a nested dict of sections,vars and default values,
        derive a ConfigObj typed spec
        NB: only single level nesting supported

        FIXME: add safeguards against broken vdict()
        '''
        spec = ''
        for section in sorted(vdict.keys()):
            spec += '[' + section + ']\n'
            for varname in sorted(vdict[section].keys()):
                typename = type(vdict[section][varname]).__name__
                if co_map.has_key(typename):
                    typename = co_map[typename]
                spec += '\t' + varname + ' = ' + typename  + '\n'
        return spec

    def restore_state(self,obj,builder):
        '''
        restore attributes from ini file 'vars' section as obj attributes,
        as well as any widget state in 'widgets' section
        '''
        if not self.defaults.has_key(IniFile.ini):
            raise BadDescriptorDictError("defaults dict lacks 'ini' section")

        if  self.defaults[IniFile.ini][IniFile.signature] != (
                self.config[IniFile.ini][IniFile.signature]):
            warn("signature mismatch in %s -  resetting to default" %
                 (self.filename))
            dbg(1, "expected: %s, got %s" %
                (self.defaults[IniFile.ini][IniFile.signature],
                 self.config[IniFile.ini][IniFile.signature]))
            self.create_default_ini()
        else:
            dbg(1,"signature verified OK for %s " % (self.filename))

        if self.config.has_key(IniFile.vars):
            for k,v in self.defaults[IniFile.vars].items():
                setattr(obj,k,self.config[IniFile.vars][k])

        if self.config.has_key(IniFile.widgets):
            for k,v in self.config[IniFile.widgets].items():
                store_value(builder.get_object(k),v)

    def save_state(self, obj,builder):
        '''
        save obj attributes as listed in ini file 'IniFile.vars' section and
        widget state to 'widgets' section
        '''
        if self.defaults.has_key(IniFile.vars):
            for k,v in self.defaults[IniFile.vars].items():
                self.config[IniFile.vars][k] = getattr(obj,k,None)

        if self.config.has_key(IniFile.widgets):
            for k in self.defaults[IniFile.widgets].keys():
                self.config[IniFile.widgets][k] = get_value(builder.get_object(k))

        self.config.final_comment = ['last update  by %s.save_state() on %s ' %
                                         (__name__,time.asctime())]
        self.write()

    def create_default_ini(self):
        '''
        create a default ini file with defaults derived from configspec
        '''
        self.config = ConfigObj(self.filename, interpolation=False,
                                configspec=self.spec)
        vdt = Validator()
        self.config.validate(vdt, copy=True)
        self.config.update(self.defaults)
        self.config.filename = self.filename
        self.config.initial_comment.append('generated by %s.create_default_ini() on %s' %
                                           (__name__, time.asctime()))
        self.config.write()

    def read_ini(self):
        '''
        make sure current file validates OK, this will also type-convert values
        recreate default ini file if bad things happen
        '''
        retries = 2
        while True:
            try:
                self.config = ConfigObj(self.filename,
                                        file_error=True,
                                        interpolation=False,
                                        configspec=self.spec,
                                        raise_errors=True)
                validator = Validator()
                results = self.config.validate(validator,preserve_errors=True)
                if results != True:
                    for (section_list, key, error
                         ) in flatten_errors(self.config, results):
                        if key is not None:
                            warn("key '%s' in section '%s' : %s" %
                                 (key, ', '.join(section_list),
                                  error if error else "missing"))
                        else:
                            warn("section '%s' missing" %
                                 ', '.join(section_list))
                    badfilename = self.filename + '.BAD'
                    os.rename(self.filename,badfilename)
                    retries -= 1
                    if retries:
                        raise UselessIniError("cant make sense of '%s', renamed to '%s'"
                                              % (self.filename, badfilename))
                    else:
                        raise Exception(errors)

            except (IOError, TypeError,UselessIniError),msg:
                warn("%s - creating default" % (msg))
                self.create_default_ini()
                continue

            except: # pass on ConfigObj's other troubles
                raise

            else: # all great
                break

    def write(self,filename=None):
        '''
        write out changed config
        '''
        if filename:
            self.filename = filename
        self.config.write()

    def __init__(self,filename,defaults):
        defaults[IniFile.ini] =   { IniFile.signature : 'astring',
                                    IniFile.version : version_number }
        self.defaults = defaults
        self.filename = filename
        spec = self._gen_spec(self.defaults)
        self.signature = sha1(spec).hexdigest()
        self.defaults[IniFile.ini][IniFile.signature] = self.signature

        dbg(2, "auto-generated spec:\n%s\nsignature = %s" %
            (spec, self.signature))
        self.spec = spec.splitlines()
        self.read_ini()
<commit_msg>gladevcp: persistence.py simplified<commit_after>#!/usr/bin/env python
# vim: sts=4 sw=4 et
#    This is a component of EMC
#    util.py Copyright 2010 Michael Haberler
#
#
#    This program 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.
#
#    This program 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 this program; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA'''
'''
    persistence support for gladevcp widgets
    Michael Haberler 11/2010
'''
import os
import sys
import time
import gtk
from configobj import ConfigObj, flatten_errors
from validate import Validator
from hashlib import sha1
from hal_widgets import _HalWidgetBase
from gladevcp.gladebuilder import widget_name

class UselessIniError(Exception):
    pass

class NoGetterSetterError(Exception):
    pass

class BadDescriptorDictError(Exception):
    pass

# map python type names to ConfigObj type names
co_map = {'str': 'string','int':'integer', 'str': 'string', 'bool': 'boolean'}
debug = 0
version_number = 1

def warn(*args):
    print >> sys.stderr,''.join(args)

def dbg(level,*args):
    global debug
    if debug < level: return
    print ''.join(args)

def set_debug(value):
    global debug
    debug = value

def select_widgets(widgets, hal_only=False,output_only = False):
    '''
    input: a list of widget instances
    output: a list of widget instances, pruned according to the flags:
        hal_only = True:  only add HAL widgets
        output_only = True: only add widgets which let users
        set a value (scale,button..)
    '''
    wlist = []
    for w in widgets:
        if hal_only and not isinstance(w, _HalWidgetBase):
            continue
        if output_only and not isinstance(w, (gtk.Range,
                                              gtk.SpinButton,
                                              gtk.ComboBox,
                                              gtk.CheckButton,
                                              gtk.ToggleButton,
                                              gtk.RadioButton)):
            continue
        wlist.append(w)
    return wlist

def accessors(w):
    '''
    retrieve getter/setter name of an 'interesting' widget
    '''
    if isinstance(w, (gtk.Range, gtk.SpinButton,gtk.ComboBox)):
        return ('get_value','set_value')
    if isinstance(w, (gtk.CheckButton, gtk.ToggleButton,gtk.RadioButton)):
        return ('get_active','set_active')

    raise NoGetterSetterError

def store_value(widget,value):
    try:
        (_,setter) = accessors(widget)
        method = getattr(widget,setter,None)
        method(value)
    except NoGetterSetterError:
        dbg(3,"store_value: no setter method for: ",str(widget))
        raise

def get_value(widget):
    try:
        (getter,_) = accessors(widget)
        method = getattr(widget,getter,None)
        return method()
    except NoGetterSetterError:
        dbg(3,"store_value: no getter method for: ",str(widget))
        raise

def widget_defaults(widgets):
    '''
    return a dict name/current value of wvalues which are to be restored
    '''
    wvalues = dict()
    for w in widgets:
        k = widget_name(w)
        try:
            v = get_value(w)
            wvalues[k] = v
        except NoGetterSetterError:
            continue
    return wvalues

class IniFile(object):

    # well known section and variable names in .ini files
    vars = 'vars'
    widgets = 'widgets'
    ini = 'ini'
    signature = 'signature'
    version = 'version'

    def _gen_spec(self,vdict):
        '''
        given a nested dict of sections,vars and default values,
        derive a ConfigObj typed spec
        NB: only single level nesting supported

        FIXME: add safeguards against broken vdict()
        '''
        spec = ''
        for section in sorted(vdict.keys()):
            spec += '[' + section + ']\n'
            for varname in sorted(vdict[section].keys()):
                typename = type(vdict[section][varname]).__name__
                if co_map.has_key(typename):
                    typename = co_map[typename]
                spec += '\t' + varname + ' = ' + typename  + '\n'
        return spec

    def restore_state(self,obj):
        '''
        restore attributes from ini file 'vars' section as obj attributes,
        as well as any widget state in 'widgets' section
        '''
        dbg(1, "restore_state() from %s" % (self.filename))

        if not self.defaults.has_key(IniFile.ini):
            raise BadDescriptorDictError("defaults dict lacks 'ini' section")

        if  self.defaults[IniFile.ini][IniFile.signature] != (
                self.config[IniFile.ini][IniFile.signature]):
            warn("signature mismatch in %s -  resetting to default" %
                 (self.filename))
            dbg(1, "expected: %s, got %s" %
                (self.defaults[IniFile.ini][IniFile.signature],
                 self.config[IniFile.ini][IniFile.signature]))
            self.create_default_ini()
        else:
            dbg(1,"signature verified OK for %s " % (self.filename))

        if self.config.has_key(IniFile.vars):
            for k,v in self.defaults[IniFile.vars].items():
                setattr(obj,k,self.config[IniFile.vars][k])

        if self.config.has_key(IniFile.widgets):
            for k,v in self.config[IniFile.widgets].items():
                store_value(self.builder.get_object(k),v)

    def save_state(self, obj):
        '''
        save obj attributes as listed in ini file 'IniFile.vars' section and
        widget state to 'widgets' section
        '''
        if self.defaults.has_key(IniFile.vars):
            for k,v in self.defaults[IniFile.vars].items():
                self.config[IniFile.vars][k] = getattr(obj,k,None)

        if self.config.has_key(IniFile.widgets):
            for k in self.defaults[IniFile.widgets].keys():
                self.config[IniFile.widgets][k] = get_value(self.builder.get_object(k))

        self.config.final_comment = ['last update  by %s.save_state() on %s ' %
                                         (__name__,time.asctime())]
        self.write()
        dbg(1, "save_state() to %s" % (self.filename))

    def create_default_ini(self):
        '''
        create a default ini file with defaults derived from configspec
        '''
        self.config = ConfigObj(self.filename, interpolation=False,
                                configspec=self.spec)
        vdt = Validator()
        self.config.validate(vdt, copy=True)
        self.config.update(self.defaults)
        self.config.filename = self.filename
        self.config.initial_comment.append('generated by %s.create_default_ini() on %s' %
                                           (__name__, time.asctime()))
        self.config.write()

    def read_ini(self):
        '''
        make sure current file validates OK, this will also type-convert values
        recreate default ini file if bad things happen
        '''
        retries = 2
        while True:
            try:
                self.config = ConfigObj(self.filename,
                                        file_error=True,
                                        interpolation=False,
                                        configspec=self.spec,
                                        raise_errors=True)
                validator = Validator()
                results = self.config.validate(validator,preserve_errors=True)
                if results != True:
                    for (section_list, key, error
                         ) in flatten_errors(self.config, results):
                        if key is not None:
                            warn("key '%s' in section '%s' : %s" %
                                 (key, ', '.join(section_list),
                                  error if error else "missing"))
                        else:
                            warn("section '%s' missing" %
                                 ', '.join(section_list))
                    badfilename = self.filename + '.BAD'
                    os.rename(self.filename,badfilename)
                    retries -= 1
                    if retries:
                        raise UselessIniError("cant make sense of '%s', renamed to '%s'"
                                              % (self.filename, badfilename))
                    else:
                        raise Exception(error)

            except (IOError, TypeError,UselessIniError),msg:
                warn("%s - creating default" % (msg))
                self.create_default_ini()
                continue

            except: # pass on ConfigObj's other troubles
                raise

            else: # all great
                break

    def write(self,filename=None):
        '''
        write out changed config
        '''
        if filename:
            self.filename = filename
        self.config.write()

    def __init__(self,filename,defaults,builder):
        defaults[IniFile.ini] =   { IniFile.signature : 'astring',
                                    IniFile.version : version_number }
        self.defaults = defaults
        self.filename = filename
        self.builder = builder
        spec = self._gen_spec(self.defaults)
        self.signature = sha1(spec).hexdigest()
        self.defaults[IniFile.ini][IniFile.signature] = self.signature

        dbg(2, "auto-generated spec:\n%s\nsignature = %s" %
            (spec, self.signature))
        self.spec = spec.splitlines()
        self.read_ini()
<|endoftext|>