<commit_before>from sklearn import hmm
import utils
from copy import deepcopy
import numpy as np
import pandas as pd
from collections import OrderedDict

def init_HMM(pi_prior,a_prior,mean_prior,cov_prior):
    '''
    Initializes a trace object from a series and a metadata dictionary.
    Series must be sampled at a particular sample rate
    pi_prior is the starting probability of the HMM
    a_prior is the transition matrix of the HMM
    mean_prior is the initial mean value of each state
    cov_prior is the initial covariance of each state

    For an n-state HMM:
    *pi_prior is a 1-D numpy array of size n
    *a_prior is a 2-D numpy array of size n x n
    *mean_prior is an numpy array of size n
    *cov_prior is a 3-D numpy array that has been tiled into two rows,
    one column, and n third dimensional states.
    ex) np.tile(1,(2,1,n)) for uniform covariance to start with.
    '''
    model = hmm.GaussianHMM(pi_prior.size,'full',pi_prior,a_prior)
    model.means_ = mean_prior
    model.covars_ = cov_prior
    return model

def fit_trace_to_HMM(model,trace):
    '''
    Fits the given trace to the model. NaNs are turned into zeroes.
    '''
    trace_values = utils.trace_series_to_numpy_array(trace.series)
    model.fit([trace_values])
    startprob, means, covars, transmat = sort_learnt_parameters(model.startprob_,
            model.means_, model.covars_ , model.transmat_)
    model=hmm.GaussianHMM(startprob.size, 'full', startprob, transmat)
    model.means_ = means
    model.covars_ = covars
    return model

def fit_instance_to_HMM(model,instance):
    '''
    Fits the given instance to the model. NaNs are turned into zeroes.
    '''
    for trace in instance.traces:
        model=fit_trace_to_HMM(model,trace)
    return model

def generate_HMMs_from_type(type,pi_prior,a_prior,
        mean_prior,cov_prior,key_for_model_name=None):
    '''
    Generates a dictionary of HMMs using each instance of given type.
    The key to the dictionary is defined by the parameter 'key_for_model_name'
    which looks at each instances metadata and uses the value from that key
    in order to name the model. If no key is given, the model is named based on
    its index.
    '''
    instance_models=OrderedDict()
    for i,instance in enumerate(type.instances):
        if(key_for_model_name):
            instance_name=instance.traces[0].metadata[key_for_model_name]
        else:
            instance_name=i
        instance_models[instance_name]=init_HMM(pi_prior,a_prior,mean_prior,cov_prior)
        instance_models[instance_name]=fit_instance_to_HMM(instance_models[instance_name],
                instance)
    return instance_models

def generate_FHMM_from_HMMs(type_models):
    '''
    Takes a dictionary of models, where the keys are the device type name, and
    generates an FHMM of these models.
    '''
    list_pi=[]
    list_A=[]
    list_means=[]
    for device_type_name in type_models:
        list_pi.append(type_models[device_type_name].startprob_)
        list_A.append(type_models[device_type_name].transmat_)
        list_means.append(type_models[device_type_name].means_.flatten().
            tolist())
    pi_combined=compute_pi_fhmm(list_pi)
    A_combined=compute_A_fhmm(list_A)
    [mean_combined, cov_combined]=compute_means_fhmm(list_means)
    model_fhmm=create_combined_hmm(len(pi_combined),pi_combined,
            A_combined, mean_combined, cov_combined)

def predict_with_FHMM(model_fhmm,test_data,plot=False):
    '''
    Predicts the decoded states and power for the given test data with the
    given FHMM. test_data is a dictionary that should contain a key called
    'aggregated', as well as keys for each device that is in the FHMM.
    '''
    learnt_states=model_fhmm.predict(test_data['total'])
    decode_hmm(len(learnt_states), mean_prior,
            [appliance for appliance in type_models], learnt_states)
    plot_FHMM_and_predictions(test_data,decoded_power)

def plot_FHMM_and_predictions(test_data,decoded_power):
    '''
    This plots the actual and predicted power based on the FHMM.
    '''
    for i,device_type in enumerate(test_data):
        if(device_type is not 'use'):
            plt.figure()
            plt.plot(test_data[device_type],'g')
            plt.title('Ground Truth Power for %s' %device_type)
            plt.plot(decoded_power[device_type],'b')
            plt.title('Predicted Power for %s' %device_type)
            plt.ylabel('Power (W)')
            plt.xlabel('Time')
            plt.ylim((np.min(test_data[device_type])-10, np.max(test_data[device_type])+10))
            plt.tight_layout()

def get_best_instance_model(instance_models,device_type,key_for_model_name):
    dfs_model = {}
    best_model_score = 0
    for model_name in instance_models:
        instances_of_model = []
        for instance in device_type.instances:
            test_trace = instance.traces[0]
            instance_name = test_trace.metadata[key_for_model_name]
            test = utils.trace_series_to_numpy_array(test_trace.series)
            model_score = instance_models[model_name].score(test)
            instances_of_model.append([model_name,instance_name,model_score])
            if(model_score > best_model_score):
                best_model = instance_models[model_name]
        dfs_model[model_name] = pd.DataFrame(data=instances_of_model,
                columns=['Model_Instance','Test_Instance','Value'])
    model_averages = []
    for key in dfs_model:
        sum=0
        for row in dfs_model[key].iterrows():
            sum = sum+row[1]['Value']
        model_averages.append([key,sum/len(dfs_model[key].index)])
    print
    avg_model_df = pd.DataFrame(data=model_averages,
            columns=['Model_Instance','Avg Probability'])
    print avg_model_df.sort('Avg Probability',ascending=False)
    bestModel = avg_model_df.sort('Avg Probability',
            ascending=False).sort('Avg Probability',
                    ascending=False).head(1)['Model_Instance'].values[0]
    print str(bestModel) + ' is best.'
    return bestModel

def sort_startprob(mapping, startprob):
    '''
    Sort the startprob of the HMM according to power means; as returned by mapping
    '''

    num_elements = len(startprob)
    new_startprob = np.zeros(num_elements)
    for i in xrange(len(startprob)):
        new_startprob[i] = startprob[mapping[i]]
    return new_startprob

def sort_covars(mapping, covars):
    num_elements = len(covars)
    new_covars = np.zeros_like(covars)
    for i in xrange(len(covars)):
        new_covars[i] = covars[mapping[i]]
    return new_covars

def sort_transition_matrix(mapping, A):
    '''
    Sorts the transition matrix of the HMM according to power means; as returned by mapping
    '''
    num_elements = len(A)
    A_new = np.zeros((num_elements, num_elements))
    for i in range(num_elements):
        for j in range(num_elements):
            A_new[i,j] = A[mapping[i], mapping[j]]
    return A_new

def return_sorting_mapping(means):
    means_copy = deepcopy(means)
    # Sorting
    means_copy = np.sort(means_copy, axis = 0)
    # Finding mapping
    mapping = {}
    mapping_set=set()
    x=0
    for i, val in enumerate(means_copy):
        x= np.where(val==means)[0]
        for val in x:
            if val not in mapping_set:
                mapping_set.add(val)
                mapping[i]=val
                break
    return mapping

def sort_learnt_parameters(startprob, means, covars, transmat):
    '''
    Sorts the learnt parameters for the HMM
    '''
    mapping = return_sorting_mapping(means)
    means_new = np.sort(means, axis = 0)

    startprob_new = sort_startprob(mapping, startprob)
    covars_new = sort_covars(mapping, covars)
    transmat_new = sort_transition_matrix(mapping, transmat)
    assert np.shape(means_new) == np.shape(means)
    assert np.shape(startprob_new) == np.shape(startprob)
    assert np.shape(transmat_new) == np.shape(transmat)
    return [startprob_new, means_new, covars_new, transmat_new]

def compute_pi_fhmm(list_pi):
    '''
    Input: list_pi: List of PI's of individual learnt HMMs
    Output: Combined Pi for the FHMM
    '''
    result=list_pi[0]
    for i in range(len(list_pi)-1):
        result=np.kron(result,list_pi[i+1])
    return result

def compute_A_fhmm(list_A):
    '''
    Input: list_pi: List of PI's of individual learnt HMMs
    Output: Combined Pi for the FHMM
    '''
    result=list_A[0]
    for i in range(len(list_A)-1):
        result=np.kron(result,list_A[i+1])
    return result

def compute_means_fhmm(list_means):
    '''
    Returns [mu, sigma]
    '''

    #list_of_appliances_centroids=[ [appliance[i][0] for i in range(len(appliance))] for appliance in list_B]
    states_combination=list(itertools.product(*list_means))
    print states_combination
    num_combinations=len(states_combination)
    print num_combinations
    means_stacked=np.array([sum(x) for x in states_combination])
    means=np.reshape(means_stacked,(num_combinations,1))
    cov=np.tile(5*np.identity(1), (num_combinations, 1, 1))
    return [means, cov]

def create_combined_hmm(n, pi, A, mean, cov):
    combined_model=hmm.GaussianHMM(n_components=n,covariance_type='full', startprob=pi, transmat=A)
    combined_model.covars_=cov
    combined_model.means_=mean
    return combined_model

def decode_hmm(length_sequence, centroids, appliance_list, states):
    '''
    Decodes the HMM state sequence
    '''
    power_states_dict={}
    hmm_states={}
    hmm_power={}
    total_num_combinations=1
    for appliance in appliance_list:
        total_num_combinations*=len(centroids[appliance])

    for appliance in appliance_list:
        hmm_states[appliance]=np.zeros(length_sequence,dtype=np.int)
        hmm_power[appliance]=np.zeros(length_sequence)

    for i in range(length_sequence):
        factor=total_num_combinations
        for appliance in appliance_list:
            #assuming integer division (will cause errors in Python 3x)
            factor=factor//len(centroids[appliance])

            temp=int(states[i])/factor
            hmm_states[appliance][i]=temp%len(centroids[appliance])
            hmm_power[appliance][i]=centroids[appliance][hmm_states[appliance][i]]

    return [hmm_states,hmm_power]
<commit_msg>finished fhmm.py<commit_after>from sklearn import hmm
import utils
from copy import deepcopy
import numpy as np
import pandas as pd
from collections import OrderedDict
import itertools
import matplotlib.pyplot as plt

def init_HMM(pi_prior,a_prior,mean_prior,cov_prior):
    '''
    Initializes a trace object from a series and a metadata dictionary.
    Series must be sampled at a particular sample rate
    pi_prior is the starting probability of the HMM
    a_prior is the transition matrix of the HMM
    mean_prior is the initial mean value of each state
    cov_prior is the initial covariance of each state

    For an n-state HMM:
    *pi_prior is a 1-D numpy array of size n
    *a_prior is a 2-D numpy array of size n x n
    *mean_prior is an numpy array of size n
    *cov_prior is a 3-D numpy array that has been tiled into two rows,
    one column, and n third dimensional states.
    ex) np.tile(1,(2,1,n)) for uniform covariance to start with.
    '''
    model = hmm.GaussianHMM(pi_prior.size,'full',pi_prior,a_prior)
    model.means_ = mean_prior
    model.covars_ = cov_prior
    return model

def fit_trace_to_HMM(model,trace):
    '''
    Fits the given trace to the model. NaNs are turned into zeroes.
    '''
    trace_values = utils.trace_series_to_numpy_array(trace.series)
    model.fit([trace_values])
    startprob, means, covars, transmat = sort_learnt_parameters(model.startprob_,
            model.means_, model.covars_ , model.transmat_)
    model=hmm.GaussianHMM(startprob.size, 'full', startprob, transmat)
    model.means_ = means
    model.covars_ = covars
    return model

def fit_instance_to_HMM(model,instance):
    '''
    Fits the given instance to the model. NaNs are turned into zeroes.
    '''
    for trace in instance.traces:
        model=fit_trace_to_HMM(model,trace)
    return model

def generate_HMMs_from_type(type,pi_prior,a_prior,
        mean_prior,cov_prior,key_for_model_name=None):
    '''
    Generates a dictionary of HMMs using each instance of given type.
    The key to the dictionary is defined by the parameter 'key_for_model_name'
    which looks at each instances metadata and uses the value from that key
    in order to name the model. If no key is given, the model is named based on
    its index.
    '''
    instance_models=OrderedDict()
    for i,instance in enumerate(type.instances):
        if(key_for_model_name):
            instance_name=instance.traces[0].metadata[key_for_model_name]
        else:
            instance_name=i
        instance_models[instance_name]=init_HMM(pi_prior,a_prior,mean_prior,cov_prior)
        instance_models[instance_name]=fit_instance_to_HMM(instance_models[instance_name],
                instance)
    return instance_models

def generate_FHMM_from_HMMs(type_models):
    '''
    Takes a dictionary of models, where the keys are the device type name, and
    generates an FHMM of these models. It returns the fhmm model as well as
    a dictionary with the key being device type and each value being a list
    containing the means for each state of that device type.
    '''
    list_pi=[]
    list_A=[]
    list_means=[]
    means={}
    for device_type_name in type_models:
        list_pi.append(type_models[device_type_name].startprob_)
        list_A.append(type_models[device_type_name].transmat_)
        list_means.append(type_models[device_type_name].means_.flatten().
            tolist())
        means[device_type_name]=type_models[device_type_name].means_
    pi_combined=compute_pi_fhmm(list_pi)
    A_combined=compute_A_fhmm(list_A)
    [mean_combined, cov_combined]=compute_means_fhmm(list_means)
    model_fhmm=create_combined_hmm(len(pi_combined),pi_combined,
            A_combined, mean_combined, cov_combined)
    return model_fhmm,means

def predict_with_FHMM(model_fhmm,means,test_data,power_total,plot=False):
    '''
    Predicts the decoded states and power for the given test data with the
    given FHMM. test_data is a dictionary containing keys for each device
    that is in the FHMM.
    '''
    learnt_states=model_fhmm.predict(power_total)
    [decoded_states,decoded_power]=decode_hmm(len(learnt_states), means,
            [appliance for appliance in test_data], learnt_states)
    if(plot):
        plot_FHMM_and_predictions(test_data,decoded_power)
    return decoded_states,decoded_power

def plot_FHMM_and_predictions(test_data,decoded_power):
    '''
    This plots the actual and predicted power based on the FHMM.
    '''
    for i,device_type in enumerate(test_data):
        if(device_type is not 'use'):
            plt.figure()
            plt.plot(test_data[device_type],'g')
            plt.plot(decoded_power[device_type],'b')
            plt.title('Ground Truth (Green) and Predicted (Blue) for %s' %device_type)
            plt.ylabel('Power (W)')
            plt.xlabel('Time')
            plt.ylim((np.min(test_data[device_type])-10, np.max(test_data[device_type])+10))
            plt.tight_layout()

def get_best_instance_model(instance_models,device_type,key_for_model_name):
    dfs_model = {}
    best_model_score = 0
    for model_name in instance_models:
        instances_of_model = []
        for instance in device_type.instances:
            test_trace = instance.traces[0]
            instance_name = test_trace.metadata[key_for_model_name]
            test = utils.trace_series_to_numpy_array(test_trace.series)
            model_score = instance_models[model_name].score(test)
            instances_of_model.append([model_name,instance_name,model_score])
            if(model_score > best_model_score):
                best_model = instance_models[model_name]
        dfs_model[model_name] = pd.DataFrame(data=instances_of_model,
                columns=['Model_Instance','Test_Instance','Value'])
    model_averages = []
    for key in dfs_model:
        sum=0
        for row in dfs_model[key].iterrows():
            sum = sum+row[1]['Value']
        model_averages.append([key,sum/len(dfs_model[key].index)])
    print
    avg_model_df = pd.DataFrame(data=model_averages,
            columns=['Model_Instance','Avg Probability'])
    print avg_model_df.sort('Avg Probability',ascending=False)
    bestModel = avg_model_df.sort('Avg Probability',
            ascending=False).sort('Avg Probability',
                    ascending=False).head(1)['Model_Instance'].values[0]
    print str(bestModel) + ' is best.'
    return bestModel

def sort_startprob(mapping, startprob):
    '''
    Sort the startprob of the HMM according to power means; as returned by mapping
    '''

    num_elements = len(startprob)
    new_startprob = np.zeros(num_elements)
    for i in xrange(len(startprob)):
        new_startprob[i] = startprob[mapping[i]]
    return new_startprob

def sort_covars(mapping, covars):
    num_elements = len(covars)
    new_covars = np.zeros_like(covars)
    for i in xrange(len(covars)):
        new_covars[i] = covars[mapping[i]]
    return new_covars

def sort_transition_matrix(mapping, A):
    '''
    Sorts the transition matrix of the HMM according to power means; as returned by mapping
    '''
    num_elements = len(A)
    A_new = np.zeros((num_elements, num_elements))
    for i in range(num_elements):
        for j in range(num_elements):
            A_new[i,j] = A[mapping[i], mapping[j]]
    return A_new

def return_sorting_mapping(means):
    means_copy = deepcopy(means)
    # Sorting
    means_copy = np.sort(means_copy, axis = 0)
    # Finding mapping
    mapping = {}
    mapping_set=set()
    x=0
    for i, val in enumerate(means_copy):
        x= np.where(val==means)[0]
        for val in x:
            if val not in mapping_set:
                mapping_set.add(val)
                mapping[i]=val
                break
    return mapping

def sort_learnt_parameters(startprob, means, covars, transmat):
    '''
    Sorts the learnt parameters for the HMM
    '''
    mapping = return_sorting_mapping(means)
    means_new = np.sort(means, axis = 0)

    startprob_new = sort_startprob(mapping, startprob)
    covars_new = sort_covars(mapping, covars)
    transmat_new = sort_transition_matrix(mapping, transmat)
    assert np.shape(means_new) == np.shape(means)
    assert np.shape(startprob_new) == np.shape(startprob)
    assert np.shape(transmat_new) == np.shape(transmat)
    return [startprob_new, means_new, covars_new, transmat_new]

def compute_pi_fhmm(list_pi):
    '''
    Input: list_pi: List of PI's of individual learnt HMMs
    Output: Combined Pi for the FHMM
    '''
    result=list_pi[0]
    for i in range(len(list_pi)-1):
        result=np.kron(result,list_pi[i+1])
    return result

def compute_A_fhmm(list_A):
    '''
    Input: list_pi: List of PI's of individual learnt HMMs
    Output: Combined Pi for the FHMM
    '''
    result=list_A[0]
    for i in range(len(list_A)-1):
        result=np.kron(result,list_A[i+1])
    return result

def compute_means_fhmm(list_means):
    '''
    Returns [mu, sigma]
    '''

    #list_of_appliances_centroids=[ [appliance[i][0] for i in range(len(appliance))] for appliance in list_B]
    states_combination=list(itertools.product(*list_means))
    print states_combination
    num_combinations=len(states_combination)
    print num_combinations
    means_stacked=np.array([sum(x) for x in states_combination])
    means=np.reshape(means_stacked,(num_combinations,1))
    cov=np.tile(5*np.identity(1), (num_combinations, 1, 1))
    return [means, cov]

def create_combined_hmm(n, pi, A, mean, cov):
    combined_model=hmm.GaussianHMM(n_components=n,covariance_type='full', startprob=pi, transmat=A)
    combined_model.covars_=cov
    combined_model.means_=mean
    return combined_model

def decode_hmm(length_sequence, centroids, appliance_list, states):
    '''
    Decodes the HMM state sequence
    '''
    power_states_dict={}
    hmm_states={}
    hmm_power={}
    total_num_combinations=1
    for appliance in appliance_list:
        total_num_combinations*=len(centroids[appliance])

    for appliance in appliance_list:
        hmm_states[appliance]=np.zeros(length_sequence,dtype=np.int)
        hmm_power[appliance]=np.zeros(length_sequence)

    for i in range(length_sequence):
        factor=total_num_combinations
        for appliance in appliance_list:
            #assuming integer division (will cause errors in Python 3x)
            factor=factor//len(centroids[appliance])

            temp=int(states[i])/factor
            hmm_states[appliance][i]=temp%len(centroids[appliance])
            hmm_power[appliance][i]=centroids[appliance][hmm_states[appliance][i]]

    return [hmm_states,hmm_power]
<|endoftext|>


<commit_before>from __future__ import division
from builtins import object

import numpy as np

from sporco import interp

class TestSet01(object):

    def test_01(self):
        x = np.arange(0, 11).astype(np.float32)
        m0 = 2.0
        c0 = 1.0
        y0 = m0 * x + c0
        y = y0.copy()
        y[4] += 2.0
        A = np.vstack([x, np.ones(x.shape[0:])]).T
        m1, c1 = interp.lstabsdev(A, y)
        assert np.abs(m0 - m1) < 1e-5
        assert np.abs(c0 - c1) < 1e-5

    def test_02(self):
        x = np.arange(0, 11).astype(np.float32)
        m0 = 2.0
        c0 = 1.0
        y0 = m0 * x + c0
        y = y0.copy()
        y[0::2] += 2.0
        y[1::2] -= 2.0
        A = np.vstack([x, np.ones(x.shape[0:])]).T
        m1, c1 = interp.lstmaxdev(A, y)
        assert np.abs(m0 - m1) < 1e-5
        assert np.abs(c0 - c1) < 1e-5

    def test_03(self):
        x = np.random.randn(9, 8)
        y = interp.bilinear_demosaic(x)
        assert np.array_equal(x[1::2, 1::2], y[1::2, 1::2, 0])
        assert np.array_equal(x[0::2, 1::2], y[0::2, 1::2, 1])
        assert np.array_equal(x[1::2, 0::2], y[1::2, 0::2, 1])
        assert np.array_equal(x[0::2, 0::2], y[0::2, 0::2, 2])
<commit_msg>Add tests for new interpolation functions<commit_after>from __future__ import division
from builtins import object

import numpy as np

from sporco import interp

class TestSet01(object):

    def test_01(self):
        x = np.arange(0, 11).astype(np.float32)
        m0 = 2.0
        c0 = 1.0
        y0 = m0 * x + c0
        y = y0.copy()
        y[4] += 2.0
        A = np.vstack([x, np.ones(x.shape[0:])]).T
        m1, c1 = interp.lstabsdev(A, y)
        assert np.abs(m0 - m1) < 1e-5
        assert np.abs(c0 - c1) < 1e-5

    def test_02(self):
        x = np.arange(0, 11).astype(np.float32)
        m0 = 2.0
        c0 = 1.0
        y0 = m0 * x + c0
        y = y0.copy()
        y[0::2] += 2.0
        y[1::2] -= 2.0
        A = np.vstack([x, np.ones(x.shape[0:])]).T
        m1, c1 = interp.lstmaxdev(A, y)
        assert np.abs(m0 - m1) < 1e-5
        assert np.abs(c0 - c1) < 1e-5

    def test_03(self):
        y = interp.lanczos_kernel(np.array([-0.5, 0.0, 0.5, 1.0]), a=2)
        assert np.abs(y[0] - y[2]) < 1e-16
        assert np.abs(y[1] - 1.0) < 1e-16
        assert np.abs(y[3]) < 1e-16

    def test_04(self):
        y = interp.lanczos_filters(4, a=2)
        assert y.shape == (5, 4)

    def test_05(self):
        y = interp.lanczos_filters((3, 4))
        assert y.shape == (7, 7, 12)
        y = interp.lanczos_filters((3, 4), collapse_axes=False)
        assert y.shape == (7, 7, 3, 4)
        y = interp.lanczos_filters((2, (0.1, 0.2, 0.3)))
        assert y.shape == (7, 7, 6)

    def test_06(self):
        x = np.random.randn(9, 8)
        y = interp.bilinear_demosaic(x)
        assert np.array_equal(x[1::2, 1::2], y[1::2, 1::2, 0])
        assert np.array_equal(x[0::2, 1::2], y[0::2, 1::2, 1])
        assert np.array_equal(x[1::2, 0::2], y[1::2, 0::2, 1])
        assert np.array_equal(x[0::2, 0::2], y[0::2, 0::2, 2])
<|endoftext|>


<commit_before>from ztag.transform import ZGrab2Transform
from ztag import protocols

class PostgresTransform(ZGrab2Transform):

    name = "postgres/banner"
    port = None
    protocol = protocols.POSTGRES
    subprotocol = protocols.POSTGRES.BANNER

    def __init__(self, *args, **kwargs):
        super(PostgresTransform, self).__init__(*args, **kwargs)

    def clean_error(self, err):
        return {k: self.clean_banner(v) for k, v in err.items()}

    def _transform_object(self, obj):
        zout = super(PostgresTransform, self)._transform_object(obj)
        results = self.get_scan_results(obj)
        if not results:
            return zout

        to_copy = ["supported_versions", "protocol_error", "startup_error",
                   "backend_key_data"]

        for f in to_copy:
            if results.get(f) is not None:
                zout.transformed[f] = results[f]

        to_clean = ["supported_versions", "protocol_error", "startup_error"]
        for f in to_clean:
            if f in zout.transformed:
                zout.transformed[f] = self.clean_banner(zout.transformed[f])

        if "is_ssl" in results:
            zout.transformed["is_ssl"] = bool(results.get("is_ssl", False))

        if results.get("authentication_mode") is not None:
            zout.transformed["authentication_mode"] = [
                x["mode"] for x in results["authentication_mode"]
            ]

        if results.get("backend_key_data") is not None:
            zout.transformed["backend_key_data"] = results["backend_key_data"]

        return zout
<commit_msg>actually use clean_error on error objects; use clean_banner on supported_versions [a string]<commit_after>from ztag.transform import ZGrab2Transform
from ztag import protocols

class PostgresTransform(ZGrab2Transform):

    name = "postgres/banner"
    port = None
    protocol = protocols.POSTGRES
    subprotocol = protocols.POSTGRES.BANNER

    def __init__(self, *args, **kwargs):
        super(PostgresTransform, self).__init__(*args, **kwargs)

    def clean_error(self, err):
        return {k: self.clean_banner(v) for k, v in err.items()}

    def _transform_object(self, obj):
        zout = super(PostgresTransform, self)._transform_object(obj)
        results = self.get_scan_results(obj)
        if not results:
            return zout

        to_copy = ["supported_versions", "protocol_error", "startup_error",
                   "backend_key_data"]

        for f in to_copy:
            if results.get(f) is not None:
                zout.transformed[f] = results[f]

        to_clean_error = ["protocol_error", "startup_error"]
        for f in to_clean_error:
            if f in zout.transformed:
                zout.transformed[f] = self.clean_error(zout.transformed[f])

        to_clean_banner = ["supported_versions"]
        for f in to_clean_banner:
            if f in zout.transformed:
                zout.transformed[f] = self.clean_banner(zout.transformed[f])

        if "is_ssl" in results:
            zout.transformed["is_ssl"] = bool(results.get("is_ssl", False))

        if results.get("authentication_mode") is not None:
            zout.transformed["authentication_mode"] = [
                x["mode"] for x in results["authentication_mode"]
            ]

        if results.get("backend_key_data") is not None:
            zout.transformed["backend_key_data"] = results["backend_key_data"]

        return zout
<|endoftext|>


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

"""
The MIT License (MIT)

Copyright (c) 2015-present Rapptz

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
"""

class _RawReprMixin:
    def __repr__(self):
        value = ' '.join('%s=%r' % (attr, getattr(self, attr)) for attr in self.__slots__)
        return '<%s %s>' % (self.__class__.__name__, value)

class RawMessageDeleteEvent(_RawReprMixin):
    """Represents the event payload for a :func:`on_raw_message_delete` event.

    Attributes
    ------------
    channel_id: :class:`int`
        The channel ID where the deletion took place.
    guild_id: Optional[:class:`int`]
        The guild ID where the deletion took place, if applicable.
    message_id: :class:`int`
        The message ID that got deleted.
    cached_message: Optional[:class:`Message`]
        The cached message, if found in the internal message cache.
    """

    __slots__ = ('message_id', 'channel_id', 'guild_id', 'cached_message')

    def __init__(self, data):
        self.message_id = int(data['id'])
        self.channel_id = int(data['channel_id'])
        self.cached_message = None
        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawBulkMessageDeleteEvent(_RawReprMixin):
    """Represents the event payload for a :func:`on_raw_bulk_message_delete` event.

    Attributes
    -----------
    message_ids: Set[:class:`int`]
        A :class:`set` of the message IDs that were deleted.
    channel_id: :class:`int`
        The channel ID where the message got deleted.
    guild_id: Optional[:class:`int`]
        The guild ID where the message got deleted, if applicable.
    cached_messages: List[:class:`Message`]
        The cached messages, if found in the internal message cache.
    """

    __slots__ = ('message_ids', 'channel_id', 'guild_id', 'cached_messages')

    def __init__(self, data):
        self.message_ids = {int(x) for x in data.get('ids', [])}
        self.channel_id = int(data['channel_id'])
        self.cached_messages = []

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawMessageUpdateEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_message_edit` event.

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got updated.
    channel_id: :class:`int`
        The channel ID where the update took place.

        .. versionadded:: 1.3

    data: :class:`dict`
        The raw data given by the `gateway <https://discord.com/developers/docs/topics/gateway#message-update>`_
    cached_message: Optional[:class:`Message`]
        The cached message, if found in the internal message cache. Represents the message before
        it is modified by the data in :attr:`RawMessageUpdateEvent.data`.
    """

    __slots__ = ('message_id', 'channel_id', 'data', 'cached_message')

    def __init__(self, data):
        self.message_id = int(data['id'])
        self.channel_id = int(data['channel_id'])
        self.data = data
        self.cached_message = None

class RawReactionActionEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_reaction_add` or
    :func:`on_raw_reaction_remove` event.

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got or lost a reaction.
    user_id: :class:`int`
        The user ID who added the reaction or whose reaction was removed.
    channel_id: :class:`int`
        The channel ID where the reaction got added or removed.
    guild_id: Optional[:class:`int`]
        The guild ID where the reaction got added or removed, if applicable.
    emoji: :class:`PartialEmoji`
        The custom or unicode emoji being used.
    member: Optional[:class:`Member`]
        The member who added the reaction. Only available if `event_type` is `REACTION_ADD` and the reaction is inside a guild.

        .. versionadded:: 1.3

    event_type: :class:`str`
        The event type that triggered this action. Can be
        ``REACTION_ADD`` for reaction addition or
        ``REACTION_REMOVE`` for reaction removal.

        .. versionadded:: 1.3
    """

    __slots__ = ('message_id', 'user_id', 'channel_id', 'guild_id', 'emoji',
                 'event_type', 'member')

    def __init__(self, data, emoji, event_type):
        self.message_id = int(data['message_id'])
        self.channel_id = int(data['channel_id'])
        self.user_id = int(data['user_id'])
        self.emoji = emoji
        self.event_type = event_type
        self.member = None

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawReactionClearEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_reaction_clear` event.

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got its reactions cleared.
    channel_id: :class:`int`
        The channel ID where the reactions got cleared.
    guild_id: Optional[:class:`int`]
        The guild ID where the reactions got cleared.
    """

    __slots__ = ('message_id', 'channel_id', 'guild_id')

    def __init__(self, data):
        self.message_id = int(data['message_id'])
        self.channel_id = int(data['channel_id'])

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawReactionClearEmojiEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_reaction_clear_emoji` event.

    .. versionadded:: 1.3

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got its reactions cleared.
    channel_id: :class:`int`
        The channel ID where the reactions got cleared.
    guild_id: Optional[:class:`int`]
        The guild ID where the reactions got cleared.
    emoji: :class:`PartialEmoji`
        The custom or unicode emoji being removed.
    """

    __slots__ = ('message_id', 'channel_id', 'guild_id', 'emoji')

    def __init__(self, data, emoji):
        self.emoji = emoji
        self.message_id = int(data['message_id'])
        self.channel_id = int(data['channel_id'])

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None
<commit_msg>Add guild_id attribute to RawMessageUpdateEvent <commit_after># -*- coding: utf-8 -*-

"""
The MIT License (MIT)

Copyright (c) 2015-present Rapptz

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
"""

class _RawReprMixin:
    def __repr__(self):
        value = ' '.join('%s=%r' % (attr, getattr(self, attr)) for attr in self.__slots__)
        return '<%s %s>' % (self.__class__.__name__, value)

class RawMessageDeleteEvent(_RawReprMixin):
    """Represents the event payload for a :func:`on_raw_message_delete` event.

    Attributes
    ------------
    channel_id: :class:`int`
        The channel ID where the deletion took place.
    guild_id: Optional[:class:`int`]
        The guild ID where the deletion took place, if applicable.
    message_id: :class:`int`
        The message ID that got deleted.
    cached_message: Optional[:class:`Message`]
        The cached message, if found in the internal message cache.
    """

    __slots__ = ('message_id', 'channel_id', 'guild_id', 'cached_message')

    def __init__(self, data):
        self.message_id = int(data['id'])
        self.channel_id = int(data['channel_id'])
        self.cached_message = None
        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawBulkMessageDeleteEvent(_RawReprMixin):
    """Represents the event payload for a :func:`on_raw_bulk_message_delete` event.

    Attributes
    -----------
    message_ids: Set[:class:`int`]
        A :class:`set` of the message IDs that were deleted.
    channel_id: :class:`int`
        The channel ID where the message got deleted.
    guild_id: Optional[:class:`int`]
        The guild ID where the message got deleted, if applicable.
    cached_messages: List[:class:`Message`]
        The cached messages, if found in the internal message cache.
    """

    __slots__ = ('message_ids', 'channel_id', 'guild_id', 'cached_messages')

    def __init__(self, data):
        self.message_ids = {int(x) for x in data.get('ids', [])}
        self.channel_id = int(data['channel_id'])
        self.cached_messages = []

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawMessageUpdateEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_message_edit` event.

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got updated.
    channel_id: :class:`int`
        The channel ID where the update took place.

        .. versionadded:: 1.3
    guild_id: Optional[:class:`int`]
        The guild ID where the message got updated, if applicable.

        .. versionadded:: 1.7

    data: :class:`dict`
        The raw data given by the `gateway <https://discord.com/developers/docs/topics/gateway#message-update>`_
    cached_message: Optional[:class:`Message`]
        The cached message, if found in the internal message cache. Represents the message before
        it is modified by the data in :attr:`RawMessageUpdateEvent.data`.
    """

    __slots__ = ('message_id', 'channel_id', 'guild_id', 'data', 'cached_message')

    def __init__(self, data):
        self.message_id = int(data['id'])
        self.channel_id = int(data['channel_id'])
        self.data = data
        self.cached_message = None

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawReactionActionEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_reaction_add` or
    :func:`on_raw_reaction_remove` event.

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got or lost a reaction.
    user_id: :class:`int`
        The user ID who added the reaction or whose reaction was removed.
    channel_id: :class:`int`
        The channel ID where the reaction got added or removed.
    guild_id: Optional[:class:`int`]
        The guild ID where the reaction got added or removed, if applicable.
    emoji: :class:`PartialEmoji`
        The custom or unicode emoji being used.
    member: Optional[:class:`Member`]
        The member who added the reaction. Only available if `event_type` is `REACTION_ADD` and the reaction is inside a guild.

        .. versionadded:: 1.3

    event_type: :class:`str`
        The event type that triggered this action. Can be
        ``REACTION_ADD`` for reaction addition or
        ``REACTION_REMOVE`` for reaction removal.

        .. versionadded:: 1.3
    """

    __slots__ = ('message_id', 'user_id', 'channel_id', 'guild_id', 'emoji',
                 'event_type', 'member')

    def __init__(self, data, emoji, event_type):
        self.message_id = int(data['message_id'])
        self.channel_id = int(data['channel_id'])
        self.user_id = int(data['user_id'])
        self.emoji = emoji
        self.event_type = event_type
        self.member = None

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawReactionClearEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_reaction_clear` event.

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got its reactions cleared.
    channel_id: :class:`int`
        The channel ID where the reactions got cleared.
    guild_id: Optional[:class:`int`]
        The guild ID where the reactions got cleared.
    """

    __slots__ = ('message_id', 'channel_id', 'guild_id')

    def __init__(self, data):
        self.message_id = int(data['message_id'])
        self.channel_id = int(data['channel_id'])

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None

class RawReactionClearEmojiEvent(_RawReprMixin):
    """Represents the payload for a :func:`on_raw_reaction_clear_emoji` event.

    .. versionadded:: 1.3

    Attributes
    -----------
    message_id: :class:`int`
        The message ID that got its reactions cleared.
    channel_id: :class:`int`
        The channel ID where the reactions got cleared.
    guild_id: Optional[:class:`int`]
        The guild ID where the reactions got cleared.
    emoji: :class:`PartialEmoji`
        The custom or unicode emoji being removed.
    """

    __slots__ = ('message_id', 'channel_id', 'guild_id', 'emoji')

    def __init__(self, data, emoji):
        self.emoji = emoji
        self.message_id = int(data['message_id'])
        self.channel_id = int(data['channel_id'])

        try:
            self.guild_id = int(data['guild_id'])
        except KeyError:
            self.guild_id = None
<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
import pytest
from mock import MagicMock as mock
import numpy as np
from numpy.testing import assert_allclose

class TestLayer:
    @pytest.fixture
    def layer(self):
        from yadll.layers import Layer
        return Layer(mock())

    @pytest.fixture
    def named_layer(self):
        from yadll.layers import Layer
        return Layer(mock(), name='layer_name')

    def test_input_shape(self, layer):
        assert layer.input_shape == layer.input_layer.output_shape

    def test_get_params(self, layer):
        assert layer.get_params() == []

    def test_get_reguls(self, layer):
        assert layer.get_reguls() == 0

    def test_named_layer(self, named_layer):
        assert named_layer.name == 'layer_name'

    def test_get_output(self, layer):
        with pytest.raises(NotImplementedError):
            layer.get_output()

    @pytest.fixture
    def layer_from_shape(self):
        from yadll.layers import Layer
        return Layer((None, 20))

    def test_layer_from_shape(self, layer_from_shape):
        layer = layer_from_shape
        assert layer.input_layer is None
        assert layer.input_shape == (None, 20)

class Testinput_layer:
    @pytest.fixture
    def input_layer(self):
        from yadll.layers import InputLayer
        return InputLayer((3, 2))

    def test_input_layer(self, input_layer):
        assert input_layer.input_layer is None

    def test_shape(self, input_layer):
        assert input_layer.input_shape == (3, 2)

    def test_get_params(self, input_layer):
        assert input_layer.get_params() == []

    def test_get_reguls(self, input_layer):
        assert input_layer.get_reguls() == 0

class TestReshapeLayer:
    @pytest.fixture
    def reshape_layer(self):
        from yadll.layers import ReshapeLayer
        return ReshapeLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.ones((16, 3, 5, 7, 10)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (16, 3, None, None, 10)
        return InputLayer(shape, input_var=input_data)

    def test_reshape(self, reshape_layer, input_layer):
        layer = reshape_layer(input_layer, (16, 3, 5, 7, 2, 5))
        assert layer.output_shape == (16, 3, 5, 7, 2, 5)
        result = layer.get_output().eval()
        assert result.shape == (16, 3, 5, 7, 2, 5)
        layer = reshape_layer(input_layer, (None, 3, 5, 7, 2, 5))
        assert layer.output_shape == (None, 3, 5, 7, 2, 5)
        result = layer.get_output().eval()
        assert result.shape == (16, 3, 5, 7, 2, 5)

class TestFlattenLayer:
    @pytest.fixture
    def flatten_layer(self):
        from yadll.layers import FlattenLayer
        return FlattenLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((2, 3, 4, 5)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (2, 3, 4, 5,)
        return InputLayer(shape, input_var=input_data)

    def test_output_shape(self, flatten_layer, input_layer):
        layer = flatten_layer(input_layer)
        assert layer.output_shape == (2, 3 * 4 * 5)

    def test_get_output(self, flatten_layer, input_layer, input_data):
        layer = flatten_layer(input_layer)
        result = layer.get_output().eval()
        input = np.asarray(input_data.eval())
        assert (result == input.reshape(input.shape[0], -1)).all()

class TestDenseLayer:
    @pytest.fixture
    def dense_layer(self):
        from yadll.layers import DenseLayer
        return DenseLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((10, 20)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (10, 20)
        return InputLayer(shape, input_var=input_data)

    @pytest.fixture
    def layer(self, dense_layer, input_layer):
        return dense_layer(input_layer, nb_units=2, l1=1, l2=2)

    @pytest.fixture
    def layer_from_layer(self, dense_layer, input_layer, layer):
        return dense_layer(input_layer, W=layer.W, b=layer.b, nb_units=2, l1=1, l2=2)

    def test_get_params(self, layer):
        assert layer.get_params() == [layer.W, layer.b]

    def test_output_shape(self, layer):
        assert layer.output_shape == (10, 2)

    def test_get_output(self, layer, input_data):
        X = input_data.eval()
        W = layer.W.eval()
        b = layer.b.eval()
        assert_allclose(layer.get_output().eval(), np.tanh(np.dot(X, W) + b), rtol=1e-4)

    def test_reguls(self, layer):
        W = layer.W.eval()
        assert_allclose(layer.reguls.eval(), np.mean(np.abs(W)) + 2 * np.mean(np.power(W, 2)), rtol=1e-4)

    def test_layer_from_layer(self, layer, layer_from_layer):
        assert layer.W == layer_from_layer.W
        assert layer.b == layer_from_layer.b

class Testunsupervised_layer:
    @pytest.fixture
    def unsupervised_layer(self):
        from yadll.layers import UnsupervisedLayer
        return UnsupervisedLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((10, 20)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (10, 20)
        return InputLayer(shape, input_var=input_data)

    @pytest.fixture
    def hp(self):
        from yadll.hyperparameters import Hyperparameters
        hp = Hyperparameters()
        hp('batch_size', 10)
        hp('n_epochs', 10)
        hp('learning_rate', 0.1)
        hp('patience', 1000)
        return hp

    @pytest.fixture
    def layer(self, unsupervised_layer, input_layer, hp):
        return unsupervised_layer(input_layer, nb_units=2, hyperparameters=hp)

    def test_get_params(self, layer):
        assert layer.get_params() == [layer.W, layer.b]

    def test_output_shape(self, layer):
        assert layer.output_shape == (10, 2)

    def test_get_output(self, layer, input_data):
        X = input_data.eval()
        W = layer.W.eval()
        b = layer.b.eval()
        assert_allclose(layer.get_output().eval(), np.tanh(np.dot(X, W) + b), rtol=1e-3)

    def test_get_encoded_input(self, layer):
        with pytest.raises(NotImplementedError):
            layer.get_encoded_input()

    def test_get_unsupervised_cost(self, layer):
        with pytest.raises(NotImplementedError):
            layer.get_unsupervised_cost()

class TestLogisticRegression:
    @pytest.fixture
    def logistic_regression(self):
        from yadll.layers import LogisticRegression
        return LogisticRegression

class TestDropout:
    @pytest.fixture
    def dropout(self):
        from yadll.layers import Dropout
        return Dropout

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((10, 20)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (10, 20)
        return InputLayer(shape, input_var=input_data)

    @pytest.fixture
    def layer(self, dropout, input_layer):
        return dropout(input_layer, corruption_level=0.5)

    @pytest.fixture
    def layer_c0(self, dropout, input_layer):
        return dropout(input_layer, corruption_level=0)

    @pytest.fixture
    def layer_c1(self, dropout, input_layer):
        return dropout(input_layer, corruption_level=1)

    def test_get_output(self, input_layer, layer, layer_c0, layer_c1):
        np.testing.assert_array_equal(input_layer.get_output().eval(), layer_c0.get_output().eval())
        assert np.all(layer_c1.get_output().eval() == 0)

class TestDropConnect:
    @pytest.fixture
    def dropconnect(self):
        from yadll.layers import Dropconnect
        return Dropconnect

class TestDropPoolLayer:
    @pytest.fixture
    def pool_layer(self):
        from yadll.layers import PoolLayer
        return PoolLayer

class TestConvLayer:
    @pytest.fixture
    def conv_layer(self):
        from yadll.layers import ConvLayer
        return ConvLayer

class TestConvPoolLayer:
    @pytest.fixture
    def conv_pool_layer(self):
        from yadll.layers import ConvPoolLayer
        return ConvPoolLayer

class TestAutoEncoder:
    @pytest.fixture
    def auto_encoder(self):
        from yadll.layers import AutoEncoder
        return AutoEncoder

class TestRBM:
    @pytest.fixture
    def rbm(self):
        from yadll.layers import RBM
        return RBM

class TestRNN:
    @pytest.fixture
    def rnn(self):
        from yadll.layers import RNN
        return RNN

class TestLSTM:
    @pytest.fixture
    def lstm(self):
        from yadll.layers import LSTM
        return LSTM

<commit_msg>first yadll commit<commit_after># -*- coding: UTF-8 -*-
import pytest
from mock import MagicMock as mock
import numpy as np
from numpy.testing import assert_allclose

class TestLayer:
    @pytest.fixture
    def layer(self):
        from yadll.layers import Layer
        return Layer(mock())

    @pytest.fixture
    def named_layer(self):
        from yadll.layers import Layer
        return Layer(mock(), name='layer_name')

    def test_input_shape(self, layer):
        assert layer.input_shape == layer.input_layer.output_shape

    def test_get_params(self, layer):
        assert layer.get_params() == []

    def test_get_reguls(self, layer):
        assert layer.get_reguls() == 0

    def test_named_layer(self, named_layer):
        assert named_layer.name == 'layer_name'

    def test_get_output(self, layer):
        with pytest.raises(NotImplementedError):
            layer.get_output()

    @pytest.fixture
    def layer_from_shape(self):
        from yadll.layers import Layer
        return Layer((None, 20))

    def test_layer_from_shape(self, layer_from_shape):
        layer = layer_from_shape
        assert layer.input_layer is None
        assert layer.input_shape == (None, 20)

class Testinput_layer:
    @pytest.fixture
    def input_layer(self):
        from yadll.layers import InputLayer
        return InputLayer((3, 2))

    def test_input_layer(self, input_layer):
        assert input_layer.input_layer is None

    def test_shape(self, input_layer):
        assert input_layer.input_shape == (3, 2)

    def test_get_params(self, input_layer):
        assert input_layer.get_params() == []

    def test_get_reguls(self, input_layer):
        assert input_layer.get_reguls() == 0

class TestReshapeLayer:
    @pytest.fixture
    def reshape_layer(self):
        from yadll.layers import ReshapeLayer
        return ReshapeLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.ones((16, 3, 5, 7, 10)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (16, 3, None, None, 10)
        return InputLayer(shape, input_var=input_data)

    def test_reshape(self, reshape_layer, input_layer):
        layer = reshape_layer(input_layer, (16, 3, 5, 7, 2, 5))
        assert layer.output_shape == (16, 3, 5, 7, 2, 5)
        result = layer.get_output().eval()
        assert result.shape == (16, 3, 5, 7, 2, 5)
        layer = reshape_layer(input_layer, (None, 3, 5, 7, 2, 5))
        assert layer.output_shape == (None, 3, 5, 7, 2, 5)
        result = layer.get_output().eval()
        assert result.shape == (16, 3, 5, 7, 2, 5)

class TestFlattenLayer:
    @pytest.fixture
    def flatten_layer(self):
        from yadll.layers import FlattenLayer
        return FlattenLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((2, 3, 4, 5)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (2, 3, 4, 5,)
        return InputLayer(shape, input_var=input_data)

    def test_output_shape(self, flatten_layer, input_layer):
        layer = flatten_layer(input_layer)
        assert layer.output_shape == (2, 3 * 4 * 5)

    def test_get_output(self, flatten_layer, input_layer, input_data):
        layer = flatten_layer(input_layer)
        result = layer.get_output().eval()
        input = np.asarray(input_data.eval())
        assert (result == input.reshape(input.shape[0], -1)).all()

class TestDenseLayer:
    @pytest.fixture
    def dense_layer(self):
        from yadll.layers import DenseLayer
        return DenseLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((10, 20)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (10, 20)
        return InputLayer(shape, input_var=input_data)

    @pytest.fixture
    def layer(self, dense_layer, input_layer):
        return dense_layer(input_layer, nb_units=2, l1=1, l2=2)

    @pytest.fixture
    def layer_from_layer(self, dense_layer, input_layer, layer):
        return dense_layer(input_layer, W=layer.W, b=layer.b, nb_units=2, l1=1, l2=2)

    def test_get_params(self, layer):
        assert layer.get_params() == [layer.W, layer.b]

    def test_output_shape(self, layer):
        assert layer.output_shape == (10, 2)

    def test_get_output(self, layer, input_data):
        X = input_data.eval()
        W = layer.W.eval()
        b = layer.b.eval()
        assert_allclose(layer.get_output().eval(), np.tanh(np.dot(X, W) + b), rtol=1e-4)

    def test_reguls(self, layer):
        W = layer.W.eval()
        assert_allclose(layer.reguls.eval(), np.mean(np.abs(W)) + 2 * np.mean(np.power(W, 2)), rtol=1e-4)

    def test_layer_from_layer(self, layer, layer_from_layer):
        assert layer.W == layer_from_layer.W
        assert layer.b == layer_from_layer.b

class Testunsupervised_layer:
    @pytest.fixture
    def unsupervised_layer(self):
        from yadll.layers import UnsupervisedLayer
        return UnsupervisedLayer

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((10, 20)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (10, 20)
        return InputLayer(shape, input_var=input_data)

    @pytest.fixture
    def hp(self):
        from yadll.hyperparameters import Hyperparameters
        hp = Hyperparameters()
        hp('batch_size', 10)
        hp('n_epochs', 10)
        hp('learning_rate', 0.1)
        hp('patience', 1000)
        return hp

    @pytest.fixture
    def layer(self, unsupervised_layer, input_layer, hp):
        return unsupervised_layer(input_layer, nb_units=2, hyperparameters=hp)

    def test_get_params(self, layer):
        assert layer.get_params() == [layer.W, layer.b]

    def test_output_shape(self, layer):
        assert layer.output_shape == (10, 2)

    def test_get_output(self, layer, input_data):
        X = input_data.eval()
        W = layer.W.eval()
        b = layer.b.eval()
        assert_allclose(layer.get_output().eval(), np.tanh(np.dot(X, W) + b), rtol=1e-3)

    def test_get_encoded_input(self, layer):
        with pytest.raises(NotImplementedError):
            layer.get_encoded_input()

    def test_get_unsupervised_cost(self, layer):
        with pytest.raises(NotImplementedError):
            layer.get_unsupervised_cost()

class TestLogisticRegression:
    @pytest.fixture
    def logistic_regression(self):
        from yadll.layers import LogisticRegression
        return LogisticRegression

class TestDropout:
    @pytest.fixture
    def dropout(self):
        from yadll.layers import Dropout
        return Dropout

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((10, 20)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (10, 20)
        return InputLayer(shape, input_var=input_data)

    @pytest.fixture
    def layer(self, dropout, input_layer):
        return dropout(input_layer, corruption_level=0.5)

    @pytest.fixture
    def layer_c0(self, dropout, input_layer):
        return dropout(input_layer, corruption_level=0)

    @pytest.fixture
    def layer_c1(self, dropout, input_layer):
        return dropout(input_layer, corruption_level=1)

    def test_get_output(self, input_layer, layer, layer_c0, layer_c1):
        np.testing.assert_array_equal(input_layer.get_output().eval(), layer_c0.get_output().eval())
        assert np.all(layer_c1.get_output().eval() == 0)

class TestDropConnect:
    @pytest.fixture
    def dropconnect(self):
        from yadll.layers import Dropconnect
        return Dropconnect

    @pytest.fixture
    def input_data(self):
        from yadll.utils import shared_variable
        return shared_variable(np.random.random((10, 20)))

    @pytest.fixture
    def input_layer(self, input_data):
        from yadll.layers import InputLayer
        shape = (10, 20)
        return InputLayer(shape, input_var=input_data)

    @pytest.fixture
    def layer(self, dropconnect, input_layer):
        return dropconnect(input_layer, nb_units=10, corruption_level=0.5)

    @pytest.fixture
    def layer_c0(self, dropconnect, input_layer):
        return dropconnect(input_layer, nb_units=10, corruption_level=0)

    @pytest.fixture
    def layer_c1(self, dropconnect, input_layer):
        return dropconnect(input_layer, nb_units=10, corruption_level=1)

    def test_get_output(self, input_layer, layer, layer_c0, layer_c1):
        assert np.all(layer_c1.get_output().eval() == 0)

class TestDropPoolLayer:
    @pytest.fixture
    def pool_layer(self):
        from yadll.layers import PoolLayer
        return PoolLayer

class TestConvLayer:
    @pytest.fixture
    def conv_layer(self):
        from yadll.layers import ConvLayer
        return ConvLayer

class TestConvPoolLayer:
    @pytest.fixture
    def conv_pool_layer(self):
        from yadll.layers import ConvPoolLayer
        return ConvPoolLayer

class TestAutoEncoder:
    @pytest.fixture
    def auto_encoder(self):
        from yadll.layers import AutoEncoder
        return AutoEncoder

class TestRBM:
    @pytest.fixture
    def rbm(self):
        from yadll.layers import RBM
        return RBM

class TestRNN:
    @pytest.fixture
    def rnn(self):
        from yadll.layers import RNN
        return RNN

class TestLSTM:
    @pytest.fixture
    def lstm(self):
        from yadll.layers import LSTM
        return LSTM

<|endoftext|>


<commit_before>from ..constructs import Instruction

class Trie(object):
    BUCKET_LEN = 1
    BUCKET_MASK = (2**BUCKET_LEN)-1
    def __init__(self):
        self.children = [None for _ in range(2**Trie.BUCKET_LEN)]
        self.value = None

    def __setitem__(self, key, value):
        assert type(value) == Instruction

        node = self
        for bucket in [(key >> i) & Trie.BUCKET_MASK for \
                       i in range(64, -1, -Trie.BUCKET_LEN)]:
            if not node.children[bucket]:
                node.children[bucket] = Trie()
            node = node.children[bucket]

        node.value = value

    def __getitem__(self, item):
        if type(item) in (int, long):
            node = self
            for bucket in [(item >> i) & Trie.BUCKET_MASK for \
                           i in range(64, -1, -Trie.BUCKET_LEN)]:
                if not node.children[bucket]:
                    raise KeyError()
                node = node.children[bucket]

            return node.value

        elif type(item) == slice:
            uncommon_bits = (item.stop - item.start).bit_length() + 1

            node = self
            for bucket in [(item.start >> i) & Trie.BUCKET_MASK for \
                           i in range(64, uncommon_bits, -Trie.BUCKET_LEN)]:
                if not node.children[bucket]:
                    raise KeyError()
                node = node.children[bucket]

            return [v for v in iter(node) if item.start <= v.address < item.stop][::item.step]

    def __iter__(self):
        if self.value:
            yield self.value
        for child in filter(None, self.children):
            for v in child:
                yield v

    def __contains__(self, item):
        node = self
        for bucket in [(item >> i) & Trie.BUCKET_MASK for \
                       i in range(64, -1, -Trie.BUCKET_LEN)]:
            if not node.children[bucket]:
                return False
            node = node.children[bucket]
        return True

    def __delitem__(self, key):
        node = self
        for bucket in [(key >> i) & Trie.BUCKET_MASK for \
                       i in range(64, -1, -Trie.BUCKET_LEN)]:
            if not node.children[bucket]:
                raise KeyError()
            node = node.children[bucket]

        if node.value:
            del node.value
<commit_msg>Fixed bug in Trie slicing<commit_after>from ..constructs import Instruction

class Trie(object):
    BUCKET_LEN = 1
    BUCKET_MASK = (2**BUCKET_LEN)-1
    def __init__(self):
        self.children = [None for _ in range(2**Trie.BUCKET_LEN)]
        self.value = None

    def __setitem__(self, key, value):
        assert type(value) == Instruction

        node = self
        for bucket in [(key >> i) & Trie.BUCKET_MASK for \
                       i in range(64, -1, -Trie.BUCKET_LEN)]:
            if not node.children[bucket]:
                node.children[bucket] = Trie()
            node = node.children[bucket]

        node.value = value

    def __getitem__(self, item):
        if type(item) in (int, long):
            node = self
            for bucket in [(item >> i) & Trie.BUCKET_MASK for \
                           i in range(64, -1, -Trie.BUCKET_LEN)]:
                if not node.children[bucket]:
                    raise KeyError()
                node = node.children[bucket]

            return node.value

        elif type(item) == slice:
            uncommon_bits = (item.stop ^ item.start).bit_length()

            node = self
            for bucket in [(item.start >> i) & Trie.BUCKET_MASK for \
                           i in range(64, uncommon_bits, -Trie.BUCKET_LEN)]:
                if not node.children[bucket]:
                    raise KeyError()
                node = node.children[bucket]

            return [v for v in iter(node) if item.start <= v.address < item.stop][::item.step]

    def __iter__(self):
        if self.value:
            yield self.value
        for child in filter(None, self.children):
            for v in child:
                yield v

    def __contains__(self, item):
        node = self
        for bucket in [(item >> i) & Trie.BUCKET_MASK for \
                       i in range(64, -1, -Trie.BUCKET_LEN)]:
            if not node.children[bucket]:
                return False
            node = node.children[bucket]
        return True

    def __delitem__(self, key):
        node = self
        for bucket in [(key >> i) & Trie.BUCKET_MASK for \
                       i in range(64, -1, -Trie.BUCKET_LEN)]:
            if not node.children[bucket]:
                raise KeyError()
            node = node.children[bucket]

        if node.value:
            del node.value
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import pymongo
import dibble.fields
import dibble.model
import dibble.mapper
from nose import with_setup
from nose.tools import eq_, assert_dict_contains_subset, raises, assert_dict_equal

DBNAME = 'dibbletest'

class UserModel(dibble.model.Model):
    name = dibble.fields.Field()

class AdvancedUserModel(dibble.model.Model):
    logincount = dibble.fields.Field()
    username = dibble.fields.Field()
    usernames = dibble.fields.Field()

class ReloadTestModel(dibble.model.Model):
    counter = dibble.fields.Field()
    foo = dibble.fields.Field()
    bar = dibble.fields.Field()

def get_db():
    con = pymongo.Connection()

    return con[DBNAME]

def get_mapper():
    db = get_db()

    return dibble.mapper.ModelMapper(UserModel, db.user)

def setup_db():
    db = get_db()
    [db.drop_collection(x) for x in db.collection_names() if not x.startswith('system.')]
    #db.connection.drop_database(DBNAME)

@with_setup(setup_db)
def test_modelmapper_new():
    users = get_mapper()

    username = 'testuser'
    user = users(name=username)

    eq_(user['name'], username)
    eq_(user.name.value, username)

@with_setup(setup_db)
def test_modelmapper_count():
    mapper = get_mapper()
    eq_(mapper.count(), 0)

@with_setup(setup_db)
def test_modelmapper_save_and_find_one():
    users = get_mapper()

    dummy_user = {'name': 'test'}
    users.save(dummy_user)

    eq_(users.count(), 1)

@with_setup(setup_db)
def test_find_generator():
    users = get_mapper()

    dummy_user = {'name': 'test'}
    uid = users.save(dummy_user, safe=True)

    for x in range(10):
        users.save({'name': 'test_' + str(x)})

    cursor = users.find({'name': dummy_user['name']})

    eq_(cursor.count(), 1)

    db_user = list(cursor)[0]

    eq_(db_user['_id'], uid)
    eq_(db_user.name.value, dummy_user['name'])
    eq_(db_user['name'], dummy_user['name'])

@with_setup(setup_db)
def test_find_getitem():
    users = get_mapper()

    dummy_user = {'name': 'test'}
    uid = users.save(dummy_user, safe=True)

    for x in range(10):
        users.save({'name': 'test_' + str(x)})

    cursor = users.find({'name': dummy_user['name']})

    eq_(cursor.count(), 1)

    db_user = cursor[0]

    eq_(db_user['_id'], uid)
    eq_(db_user.name.value, dummy_user['name'])
    eq_(db_user['name'], dummy_user['name'])

@with_setup(setup_db)
def test_modelmapper_model_save():
    db = get_db()
    users = dibble.mapper.ModelMapper(AdvancedUserModel, db.user)

    user = users()
    user.logincount.inc(1)
    user.username.set('Foo Bar')
    user.usernames.push('Foo Bar')
    user.save()

    u = dict(users.collection.find_one())
    expected = {'logincount': 1, 'username': 'Foo Bar', 'usernames': ['Foo Bar']}

    assert_dict_contains_subset(expected, u)

    users.collection.update({}, {'$set': {'username': 'Fumm Fumm'}})

    user.logincount.inc(41)
    user.save()

    u = dict(users.collection.find_one())
    expected = {'logincount': 42, 'username': 'Fumm Fumm', 'usernames': ['Foo Bar']}

    assert_dict_contains_subset(expected, u)

@with_setup(setup_db)
def test_modelmapper_model_reload():
    db = get_db()
    users = dibble.mapper.ModelMapper(AdvancedUserModel, db.user)

    user = users()
    user.username.set('Foo Bar')
    user.save()

    users.collection.update({}, {'$set': {'username': 'Fumm Fumm'}})

    user.reload()

    expected = {'username': 'Fumm Fumm'}

    assert_dict_contains_subset(expected, dict(user))

@with_setup(setup_db)
@raises(dibble.model.UnboundModelError)
def test_modelmapper_model_reload_unbound():
    user = AdvancedUserModel()
    user.reload()

@with_setup(setup_db)
@raises(dibble.model.UnsavedModelError)
def test_modelmapper_model_reload_unsaved():
    db = get_db()
    users = dibble.mapper.ModelMapper(AdvancedUserModel, db.user)
    user = users()
    user.reload()

@with_setup(setup_db)
def test_modelmapper_custom_id():
    db = get_db()
    users = dibble.mapper.ModelMapper(UserModel, db.user)
    user = users()
    user.name.set('Foo Bar')

    user.save(_id='foobar')

    res = users.collection.find_one()
    expected = {'_id': 'foobar', 'name': 'Foo Bar'}

    assert_dict_equal(res, expected)

@with_setup(setup_db)
def test_modelpapper_reload_inc():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    # this is intentional and triggers a model reload
    m.counter.value

    m.counter.inc(1)
    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_set():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.set('foo')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, 'foo')

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_inc():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.counter.inc(1)

    eq_(m.foo.value, 'bar')
    eq_(m.counter.value, 1)

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_unset():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set('foo')
    m.save()

    m.bar.unset()

    eq_(m.foo.value, 'bar')
    eq_(m.bar.defined, False)

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_push():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.push('foo')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_push_all():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.push_all(['foo', 'bar', 'baz'])

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo', 'bar', 'baz'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_add_to_set():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.add_to_set('foo')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_pop():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set(['foo', 'bar', 'baz'])
    m.save()

    m.bar.pop()

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo', 'bar'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_pull():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set(['foo', 'bar', 'baz'])
    m.save()

    m.bar.pull('bar')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo', 'baz'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_pull_all():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set(['foo', 'bar', 'baz'])
    m.save()

    m.bar.pull_all(['bar', 'baz'])

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo'])

    m.save()

@with_setup(setup_db)
def test_modelmapper_unsafe_update():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.counter.set(1)
    m.save()

    m._update.inc('counter', 41)
    m.save()

    expected = {'foo': 'bar', 'counter': 42}
    assert_dict_contains_subset(expected, dict(m))
<commit_msg>added regression test for subfield error<commit_after># -*- coding: utf-8 -*-
import pymongo
import dibble.fields
import dibble.model
import dibble.mapper
from nose import with_setup
from nose.tools import eq_, assert_dict_contains_subset, raises, assert_dict_equal

DBNAME = 'dibbletest'

class UserModel(dibble.model.Model):
    name = dibble.fields.Field()

class AdvancedUserModel(dibble.model.Model):
    logincount = dibble.fields.Field()
    username = dibble.fields.Field()
    usernames = dibble.fields.Field()

class ReloadTestModel(dibble.model.Model):
    counter = dibble.fields.Field()
    foo = dibble.fields.Field()
    bar = dibble.fields.Field()

def get_db():
    con = pymongo.Connection()

    return con[DBNAME]

def get_mapper():
    db = get_db()

    return dibble.mapper.ModelMapper(UserModel, db.user)

def setup_db():
    db = get_db()
    [db.drop_collection(x) for x in db.collection_names() if not x.startswith('system.')]
    #db.connection.drop_database(DBNAME)

@with_setup(setup_db)
def test_modelmapper_new():
    users = get_mapper()

    username = 'testuser'
    user = users(name=username)

    eq_(user['name'], username)
    eq_(user.name.value, username)

@with_setup(setup_db)
def test_modelmapper_count():
    mapper = get_mapper()
    eq_(mapper.count(), 0)

@with_setup(setup_db)
def test_modelmapper_save_and_find_one():
    users = get_mapper()

    dummy_user = {'name': 'test'}
    users.save(dummy_user)

    eq_(users.count(), 1)

@with_setup(setup_db)
def test_find_generator():
    users = get_mapper()

    dummy_user = {'name': 'test'}
    uid = users.save(dummy_user, safe=True)

    for x in range(10):
        users.save({'name': 'test_' + str(x)})

    cursor = users.find({'name': dummy_user['name']})

    eq_(cursor.count(), 1)

    db_user = list(cursor)[0]

    eq_(db_user['_id'], uid)
    eq_(db_user.name.value, dummy_user['name'])
    eq_(db_user['name'], dummy_user['name'])

@with_setup(setup_db)
def test_find_getitem():
    users = get_mapper()

    dummy_user = {'name': 'test'}
    uid = users.save(dummy_user, safe=True)

    for x in range(10):
        users.save({'name': 'test_' + str(x)})

    cursor = users.find({'name': dummy_user['name']})

    eq_(cursor.count(), 1)

    db_user = cursor[0]

    eq_(db_user['_id'], uid)
    eq_(db_user.name.value, dummy_user['name'])
    eq_(db_user['name'], dummy_user['name'])

@with_setup(setup_db)
def test_modelmapper_model_save():
    db = get_db()
    users = dibble.mapper.ModelMapper(AdvancedUserModel, db.user)

    user = users()
    user.logincount.inc(1)
    user.username.set('Foo Bar')
    user.usernames.push('Foo Bar')
    user.save()

    u = dict(users.collection.find_one())
    expected = {'logincount': 1, 'username': 'Foo Bar', 'usernames': ['Foo Bar']}

    assert_dict_contains_subset(expected, u)

    users.collection.update({}, {'$set': {'username': 'Fumm Fumm'}})

    user.logincount.inc(41)
    user.save()

    u = dict(users.collection.find_one())
    expected = {'logincount': 42, 'username': 'Fumm Fumm', 'usernames': ['Foo Bar']}

    assert_dict_contains_subset(expected, u)

@with_setup(setup_db)
def test_modelmapper_model_reload():
    db = get_db()
    users = dibble.mapper.ModelMapper(AdvancedUserModel, db.user)

    user = users()
    user.username.set('Foo Bar')
    user.save()

    users.collection.update({}, {'$set': {'username': 'Fumm Fumm'}})

    user.reload()

    expected = {'username': 'Fumm Fumm'}

    assert_dict_contains_subset(expected, dict(user))

@with_setup(setup_db)
@raises(dibble.model.UnboundModelError)
def test_modelmapper_model_reload_unbound():
    user = AdvancedUserModel()
    user.reload()

@with_setup(setup_db)
@raises(dibble.model.UnsavedModelError)
def test_modelmapper_model_reload_unsaved():
    db = get_db()
    users = dibble.mapper.ModelMapper(AdvancedUserModel, db.user)
    user = users()
    user.reload()

@with_setup(setup_db)
def test_modelmapper_custom_id():
    db = get_db()
    users = dibble.mapper.ModelMapper(UserModel, db.user)
    user = users()
    user.name.set('Foo Bar')

    user.save(_id='foobar')

    res = users.collection.find_one()
    expected = {'_id': 'foobar', 'name': 'Foo Bar'}

    assert_dict_equal(res, expected)

@with_setup(setup_db)
def test_modelpapper_reload_inc():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    # this is intentional and triggers a model reload
    m.counter.value

    m.counter.inc(1)
    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_set():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.set('foo')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, 'foo')

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_inc():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.counter.inc(1)

    eq_(m.foo.value, 'bar')
    eq_(m.counter.value, 1)

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_unset():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set('foo')
    m.save()

    m.bar.unset()

    eq_(m.foo.value, 'bar')
    eq_(m.bar.defined, False)

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_push():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.push('foo')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_push_all():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.push_all(['foo', 'bar', 'baz'])

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo', 'bar', 'baz'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_add_to_set():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.save()

    m.bar.add_to_set('foo')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_pop():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set(['foo', 'bar', 'baz'])
    m.save()

    m.bar.pop()

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo', 'bar'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_pull():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set(['foo', 'bar', 'baz'])
    m.save()

    m.bar.pull('bar')

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo', 'baz'])

    m.save()

@with_setup(setup_db)
def test_modelpapper_reload_multifield_pull_all():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.bar.set(['foo', 'bar', 'baz'])
    m.save()

    m.bar.pull_all(['bar', 'baz'])

    eq_(m.foo.value, 'bar')
    eq_(m.bar.value, ['foo'])

    m.save()

@with_setup(setup_db)
def test_modelmapper_unsafe_update():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(ReloadTestModel, db.reloadtest)

    m = mapper()
    m.foo.set('bar')
    m.counter.set(1)
    m.save()

    m._update.inc('counter', 41)
    m.save()

    expected = {'foo': 'bar', 'counter': 42}
    assert_dict_contains_subset(expected, dict(m))

@with_setup(setup_db)
def test_subfield_invalidate():
    db = get_db()
    mapper = dibble.mapper.ModelMapper(UserModel, db.subfieldtest)

    m = mapper({'bar': 'baz'})
    m.save()
    m.name.set({'firstname': 'foo'})
    sf = m.name['firstname']

    m.reload()
    m.name.set({'firstname': 'bar'})
    m.save()
<|endoftext|>