<commit_before>import unittest
import os
from subprocess import Popen
import subprocess
import tarfile
import shutil
import json
import numpy as np
import math
from coremltools.converters import caffe as caffe_converter
nets_path = os.getenv('CAFFE_MODELS_PATH', '')
nets_path = nets_path + '/' if nets_path else ''
import coremltools
import pytest

def extract_tarfile(input_filename, dest_dir):
    with tarfile.open(input_filename, "r:gz") as tar:
        tar.extractall(dest_dir)

def traverse_caffe_nets(layer_type):
    for root, dirs, files in os.walk(
            '{}nets/{}/caffemodels/'.format(nets_path, layer_type)):
        return files

def traverse_data_files(layer_type):
    for root, dirs, files in os.walk(
            '{}nets/{}/data/'.format(nets_path, layer_type)):
        return set(files)

"""def convert(model, image_input_names=[], is_bgr=False,
            red_bias=0.0, blue_bias=0.0, green_bias=0.0, gray_bias=0.0,
            image_scale=1.0, class_labels=None, predicted_feature_name=None):
            """
def conversion_to_mlmodel(net_name, proto_name, layer_type, input_layer):
    filename= '{}nets/{}/mlkitmodels/{}.mlmodel'.format(
                              nets_path, layer_type, net_name)
    caffe_model_path = '{}nets/{}/caffemodels/{}.caffemodel'.format(
                              nets_path, layer_type, net_name),
    proto_path = '{}nets/{}/prototxt/{}.prototxt'.format(
                   nets_path, layer_type, proto_name)
    model_path = caffe_model_path[0]
    if isinstance(input_layer, str):
        input_layer = [input_layer]
    try:
        model = caffe_converter.convert(
            (model_path, proto_path),
        )
        model.save(filename)
    except RuntimeError as e:
        print(e)
        return False
    return True

def load_mlmodel(net_name, layer_type):
    load_args = [' /usr/local/bin/coremltest',
                 'load',
                 '-modelPath',
                 '{}nets/{}/mlkitmodels/{}.mlmodel'.format(
                     nets_path, layer_type, net_name),
                 ]
    print('Loading {}'.format(net_name))
    process = Popen((" ").join(load_args),
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    shell=True)
    stdout, err = process.communicate()

    if err == '':
        return True
    else:
        print(" The error is {}".format(err))
        return False

def evaluation_data(net_name, layer_type, data_files):
    if 'input' in net_name:
        for file in data_files:
            if 'output' in file:
                output_data_file = file
            elif 'input' in file:
                input_data_file = file
            else:
                return False

        with open('{}nets/{}/data/{}'.format(
                nets_path, layer_type, input_data_file)
        ) as data_file:
            input_net_data = json.load(data_file)
        with open('{}nets/{}/data/{}'.format(
                nets_path, layer_type, output_data_file)
        ) as data_file:
            output_net_data = json.load(data_file)
        if isinstance(input_net_data, list):
            input_data = []
            for data in input_net_data:
                input_data.append(np.array(data["input_data"]))
        else:
            input_data = np.array(input_net_data["input_data"], dtype = 'f')
        if isinstance(output_net_data, list):
            output_data = []
            for data in output_net_data:
                output_data.append(np.array(data["output_data"]))
        else:
            output_data = np.array(output_net_data["output_data"])

        return input_data, output_data

@unittest.skipIf(not nets_path, 'Caffe nets path environment variable not '
                                'found. Skipping all caffe nose tests')
class CaffeLayers(unittest.TestCase):
    """
    Unit test case for caffe layers
    """
    @classmethod
    def setUpClass(self):
        """
        Set up the unit test by loading common utilities.
        """

    def run_case(self, layer_type, input_layer, output_layer, delta=1e-2):
        self.maxDiff = None
        extract_tarfile('{}nets/{}.gz'.format(nets_path, layer_type),
                        '{}nets/'.format(nets_path))
        nets = traverse_caffe_nets(layer_type)
        data_files = traverse_data_files(layer_type)
        failed_tests_load = []
        failed_tests_conversion = []
        failed_tests_evaluation = []
        counter = 0
        for net_name_proto in nets:

            counter += 1
            net_data_files = []
            proto_name = \
                net_name_proto.split("_")[0] + \
                "_" + \
                net_name_proto.split("_")[1]
            for file in data_files:
                if proto_name + '_' in file:
                    net_data_files.append(file)
            net_name = net_name_proto.split(".")[0]
            conversion_result = conversion_to_mlmodel(
                net_name,
                proto_name,
                layer_type,
                input_layer
            )
            if conversion_result is False:
                failed_tests_conversion.append(net_name)
                continue
            load_result = load_mlmodel(net_name, layer_type)
            if load_result is False:
                failed_tests_load.append(net_name)
            if 'input' in net_name:
                evaluation_result, failed_tests_evaluation = \
                    self.evaluate_model(
                        net_name,
                        layer_type,
                        input_layer,
                        output_layer,
                        net_data_files,
                        failed_tests_evaluation,
                        counter,
                        delta)
        with open('./failed_tests_{}.json'.format(layer_type), mode='w') \
                as file:
            json.dump({'conversion': failed_tests_conversion,
                       'load': failed_tests_load,
                       'evaluation': failed_tests_evaluation},
                      file)

        self.assertEqual(failed_tests_conversion,
                         [])
        self.assertEqual(failed_tests_load,
                         [])
        self.assertEqual(failed_tests_evaluation,
                         [])
        shutil.rmtree('{}nets/{}'.format(nets_path, layer_type))

    def evaluate_model(self,
                       net_name,
                       layer_type,
                       input_layer,
                       output_layer,
                       data_files,
                       failed_tests,
                       counter,
                       delta=1e-2):
        input_data, output_data = evaluation_data(
            net_name,
            layer_type,
            data_files)
        model_path = '{}nets/{}/mlkitmodels/{}.mlmodel'.format(
            nets_path,
            layer_type,
            net_name
        )
        coremlmodel = coremltools.models.MLModel(model_path)
        mlmodel_input = {}
        if isinstance(input_layer, list):
            i = 0
            for layer in input_layer:
                if len(input_data[i].shape) == 4:
                    if input_data[i].shape[0] > 1:
                        mlmodel_input[layer] = np.expand_dims(input_data[i], axis=0)
                    else:
                        mlmodel_input[layer] = input_data[i][0,:,:,:]
                else:
                    mlmodel_input[layer] = input_data[i]
                i += 1
        else:
            if len(input_data.shape) == 4:
                if input_data.shape[0] > 1:
                    if str(coremlmodel.output_description).split('(')[1][:-1] == 'LayerEmbed':
                        mlmodel_input[input_layer] = np.expand_dims(input_data, axis=1)
                    else:
                        mlmodel_input[input_layer] = np.expand_dims(input_data, axis=0)
                else :
                    mlmodel_input[input_layer] = input_data[0,:,:,:]
            else:
                mlmodel_input[input_layer] = input_data[0]
        if isinstance(output_layer, list):
            output_preds = coremlmodel.predict(mlmodel_input)
            coreml_preds = []
            caffe_preds = []
            i = 0
            for key in sorted(output_preds):
                coreml_preds.extend(output_preds[key].flatten().tolist())
                caffe_preds.extend(output_data[i].flatten().tolist())
                i += 1
        else:
            output_layer_name = str(coremlmodel.output_description).split('(')[1].split(')')[0]
            coreml_preds = coremlmodel.predict(mlmodel_input)[output_layer_name].flatten()
            caffe_preds = output_data.flatten()

        if len(coreml_preds) != len(caffe_preds):
            failed_tests.append(net_name)
            return relative_error, failed_tests

        for i in range(len(caffe_preds)):
            max_den = max(1.0, np.abs(caffe_preds[i]), np.abs(coreml_preds[i]))
            relative_error = np.abs(caffe_preds[i]/max_den - coreml_preds[i]/max_den)
            if relative_error > delta and np.abs(caffe_preds[i]) < 1e10:
                failed_tests.append(net_name)
                break

        return relative_error, failed_tests

    def test_caffe_inner_product_layer(self):
        self.run_case(
            layer_type='inner_product',
            input_layer='data',
            output_layer='LayerInnerProduct'
        )

    def test_caffe_inner_product_activation_layer(self):
        self.run_case(
            layer_type='inner_product_activation',
            input_layer='data',
            output_layer='LayerActivation'
        )

    #@unittest.skip("Add Test cases where group is not 1: Radar: 32739970")
    @pytest.mark.slow
    def test_convolutional_layer(self):
        self.run_case(
            layer_type='convolutional',
            input_layer='data',
            output_layer='LayerConvolution'
        )

    #@unittest.skip("Add Test cases where group and dilation are not 1 Radar: 32739970")
    @pytest.mark.slow
    def test_deconvolution_layer(self):
        self.run_case(
            layer_type='deconvolution',
            input_layer='data',
            output_layer='LayerDeconvolution'
        )

    def test_reduction_layer(self):
        self.run_case(
            layer_type='reduction',
            input_layer='data',
            output_layer='LayerReduction'
            )

    def test_scale_layer(self):
        self.run_case(
            layer_type='scale',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerScale'
        )

    def test_slice_layer(self):
        self.run_case(
            layer_type='slice',
            input_layer='data',
            output_layer=['LayerSlice', 'LayerSlice1', 'LayerSlice2']
        )

    def test_bias_layer(self):
        self.run_case(
            layer_type='bias',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerBias'
        )

    def test_crop_layer(self):
        self.run_case(
            layer_type='crop',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerCrop'
        )

    @unittest.skip(" Radar: 32877551")
    def test_concat_layer(self):
        self.run_case(
            layer_type='concat',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerConcat'
        )

    @pytest.mark.slow
    def test_pooling_layer(self):
        self.run_case(
            layer_type='pooling',
            input_layer='data',
            output_layer='LayerPooling'
        )

    def test_lrn(self):
        self.run_case(
            layer_type='lrn',
            input_layer='data',
            output_layer='LayerLRN',
        )

    @unittest.skip(" Radar: 33056676")
    def test_mvn(self):
        self.run_case(
            layer_type='mvn',
            input_layer='data',
            output_layer='LayerMVN',
        )

    def test_reshape(self):
        self.run_case(
            layer_type='reshape',
            input_layer='data',
            output_layer='LayerReshape',
        )

    def test_embed(self):
        self.run_case(
            layer_type='embed',
            input_layer='data',
            output_layer='LayerEmbed',
        )

    def test_batchnorm(self):
        self.run_case(
            layer_type='batchnorm',
            input_layer='data',
            output_layer='LayerBatchNorm',
        )

    def test_flatten(self):
        self.run_case(
            layer_type='flatten',
            input_layer='data',
            output_layer='LayerFlatten',
        )

    def test_eltwise(self):
        self.run_case(
            layer_type='eltwise',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerEltwise',
        )

    @unittest.skip("Radar: 32739970")
    def test_parameter(self):
        self.run_case(
            layer_type='parameter',
            input_layer='data',
            output_layer='LayerParameter',
        )

    def test_split(self):
        self.run_case(
            layer_type='split',
            input_layer='data',
            output_layer=['LayerOutput_1', 'LayerOutput_2'],
        )
<commit_msg>Mark slow tests as such<commit_after>import unittest
import os
from subprocess import Popen
import subprocess
import tarfile
import shutil
import json
import numpy as np
import math
from coremltools.converters import caffe as caffe_converter
nets_path = os.getenv('CAFFE_MODELS_PATH', '')
nets_path = nets_path + '/' if nets_path else ''
import coremltools
import pytest

def extract_tarfile(input_filename, dest_dir):
    with tarfile.open(input_filename, "r:gz") as tar:
        tar.extractall(dest_dir)

def traverse_caffe_nets(layer_type):
    for root, dirs, files in os.walk(
            '{}nets/{}/caffemodels/'.format(nets_path, layer_type)):
        return files

def traverse_data_files(layer_type):
    for root, dirs, files in os.walk(
            '{}nets/{}/data/'.format(nets_path, layer_type)):
        return set(files)

"""def convert(model, image_input_names=[], is_bgr=False,
            red_bias=0.0, blue_bias=0.0, green_bias=0.0, gray_bias=0.0,
            image_scale=1.0, class_labels=None, predicted_feature_name=None):
            """
def conversion_to_mlmodel(net_name, proto_name, layer_type, input_layer):
    filename= '{}nets/{}/mlkitmodels/{}.mlmodel'.format(
                              nets_path, layer_type, net_name)
    caffe_model_path = '{}nets/{}/caffemodels/{}.caffemodel'.format(
                              nets_path, layer_type, net_name),
    proto_path = '{}nets/{}/prototxt/{}.prototxt'.format(
                   nets_path, layer_type, proto_name)
    model_path = caffe_model_path[0]
    if isinstance(input_layer, str):
        input_layer = [input_layer]
    try:
        model = caffe_converter.convert(
            (model_path, proto_path),
        )
        model.save(filename)
    except RuntimeError as e:
        print(e)
        return False
    return True

def load_mlmodel(net_name, layer_type):
    load_args = [' /usr/local/bin/coremltest',
                 'load',
                 '-modelPath',
                 '{}nets/{}/mlkitmodels/{}.mlmodel'.format(
                     nets_path, layer_type, net_name),
                 ]
    print('Loading {}'.format(net_name))
    process = Popen((" ").join(load_args),
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    shell=True)
    stdout, err = process.communicate()

    if err == '':
        return True
    else:
        print(" The error is {}".format(err))
        return False

def evaluation_data(net_name, layer_type, data_files):
    if 'input' in net_name:
        for file in data_files:
            if 'output' in file:
                output_data_file = file
            elif 'input' in file:
                input_data_file = file
            else:
                return False

        with open('{}nets/{}/data/{}'.format(
                nets_path, layer_type, input_data_file)
        ) as data_file:
            input_net_data = json.load(data_file)
        with open('{}nets/{}/data/{}'.format(
                nets_path, layer_type, output_data_file)
        ) as data_file:
            output_net_data = json.load(data_file)
        if isinstance(input_net_data, list):
            input_data = []
            for data in input_net_data:
                input_data.append(np.array(data["input_data"]))
        else:
            input_data = np.array(input_net_data["input_data"], dtype = 'f')
        if isinstance(output_net_data, list):
            output_data = []
            for data in output_net_data:
                output_data.append(np.array(data["output_data"]))
        else:
            output_data = np.array(output_net_data["output_data"])

        return input_data, output_data

@unittest.skipIf(not nets_path, 'Caffe nets path environment variable not '
                                'found. Skipping all caffe nose tests')
class CaffeLayers(unittest.TestCase):
    """
    Unit test case for caffe layers
    """
    @classmethod
    def setUpClass(self):
        """
        Set up the unit test by loading common utilities.
        """

    def run_case(self, layer_type, input_layer, output_layer, delta=1e-2):
        self.maxDiff = None
        extract_tarfile('{}nets/{}.gz'.format(nets_path, layer_type),
                        '{}nets/'.format(nets_path))
        nets = traverse_caffe_nets(layer_type)
        data_files = traverse_data_files(layer_type)
        failed_tests_load = []
        failed_tests_conversion = []
        failed_tests_evaluation = []
        counter = 0
        for net_name_proto in nets:

            counter += 1
            net_data_files = []
            proto_name = \
                net_name_proto.split("_")[0] + \
                "_" + \
                net_name_proto.split("_")[1]
            for file in data_files:
                if proto_name + '_' in file:
                    net_data_files.append(file)
            net_name = net_name_proto.split(".")[0]
            conversion_result = conversion_to_mlmodel(
                net_name,
                proto_name,
                layer_type,
                input_layer
            )
            if conversion_result is False:
                failed_tests_conversion.append(net_name)
                continue
            load_result = load_mlmodel(net_name, layer_type)
            if load_result is False:
                failed_tests_load.append(net_name)
            if 'input' in net_name:
                evaluation_result, failed_tests_evaluation = \
                    self.evaluate_model(
                        net_name,
                        layer_type,
                        input_layer,
                        output_layer,
                        net_data_files,
                        failed_tests_evaluation,
                        counter,
                        delta)
        with open('./failed_tests_{}.json'.format(layer_type), mode='w') \
                as file:
            json.dump({'conversion': failed_tests_conversion,
                       'load': failed_tests_load,
                       'evaluation': failed_tests_evaluation},
                      file)

        self.assertEqual(failed_tests_conversion,
                         [])
        self.assertEqual(failed_tests_load,
                         [])
        self.assertEqual(failed_tests_evaluation,
                         [])
        shutil.rmtree('{}nets/{}'.format(nets_path, layer_type))

    def evaluate_model(self,
                       net_name,
                       layer_type,
                       input_layer,
                       output_layer,
                       data_files,
                       failed_tests,
                       counter,
                       delta=1e-2):
        input_data, output_data = evaluation_data(
            net_name,
            layer_type,
            data_files)
        model_path = '{}nets/{}/mlkitmodels/{}.mlmodel'.format(
            nets_path,
            layer_type,
            net_name
        )
        coremlmodel = coremltools.models.MLModel(model_path)
        mlmodel_input = {}
        if isinstance(input_layer, list):
            i = 0
            for layer in input_layer:
                if len(input_data[i].shape) == 4:
                    if input_data[i].shape[0] > 1:
                        mlmodel_input[layer] = np.expand_dims(input_data[i], axis=0)
                    else:
                        mlmodel_input[layer] = input_data[i][0,:,:,:]
                else:
                    mlmodel_input[layer] = input_data[i]
                i += 1
        else:
            if len(input_data.shape) == 4:
                if input_data.shape[0] > 1:
                    if str(coremlmodel.output_description).split('(')[1][:-1] == 'LayerEmbed':
                        mlmodel_input[input_layer] = np.expand_dims(input_data, axis=1)
                    else:
                        mlmodel_input[input_layer] = np.expand_dims(input_data, axis=0)
                else :
                    mlmodel_input[input_layer] = input_data[0,:,:,:]
            else:
                mlmodel_input[input_layer] = input_data[0]
        if isinstance(output_layer, list):
            output_preds = coremlmodel.predict(mlmodel_input)
            coreml_preds = []
            caffe_preds = []
            i = 0
            for key in sorted(output_preds):
                coreml_preds.extend(output_preds[key].flatten().tolist())
                caffe_preds.extend(output_data[i].flatten().tolist())
                i += 1
        else:
            output_layer_name = str(coremlmodel.output_description).split('(')[1].split(')')[0]
            coreml_preds = coremlmodel.predict(mlmodel_input)[output_layer_name].flatten()
            caffe_preds = output_data.flatten()

        if len(coreml_preds) != len(caffe_preds):
            failed_tests.append(net_name)
            return relative_error, failed_tests

        for i in range(len(caffe_preds)):
            max_den = max(1.0, np.abs(caffe_preds[i]), np.abs(coreml_preds[i]))
            relative_error = np.abs(caffe_preds[i]/max_den - coreml_preds[i]/max_den)
            if relative_error > delta and np.abs(caffe_preds[i]) < 1e10:
                failed_tests.append(net_name)
                break

        return relative_error, failed_tests

    @pytest.mark.slow
    def test_caffe_inner_product_layer(self):
        self.run_case(
            layer_type='inner_product',
            input_layer='data',
            output_layer='LayerInnerProduct'
        )

    @pytest.mark.slow
    def test_caffe_inner_product_activation_layer(self):
        self.run_case(
            layer_type='inner_product_activation',
            input_layer='data',
            output_layer='LayerActivation'
        )

    #@unittest.skip("Add Test cases where group is not 1: Radar: 32739970")
    @pytest.mark.slow
    def test_convolutional_layer(self):
        self.run_case(
            layer_type='convolutional',
            input_layer='data',
            output_layer='LayerConvolution'
        )

    #@unittest.skip("Add Test cases where group and dilation are not 1 Radar: 32739970")
    @pytest.mark.slow
    def test_deconvolution_layer(self):
        self.run_case(
            layer_type='deconvolution',
            input_layer='data',
            output_layer='LayerDeconvolution'
        )

    def test_reduction_layer(self):
        self.run_case(
            layer_type='reduction',
            input_layer='data',
            output_layer='LayerReduction'
            )

    def test_scale_layer(self):
        self.run_case(
            layer_type='scale',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerScale'
        )

    def test_slice_layer(self):
        self.run_case(
            layer_type='slice',
            input_layer='data',
            output_layer=['LayerSlice', 'LayerSlice1', 'LayerSlice2']
        )

    def test_bias_layer(self):
        self.run_case(
            layer_type='bias',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerBias'
        )

    @pytest.mark.slow
    def test_crop_layer(self):
        self.run_case(
            layer_type='crop',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerCrop'
        )

    @unittest.skip(" Radar: 32877551")
    def test_concat_layer(self):
        self.run_case(
            layer_type='concat',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerConcat'
        )

    @pytest.mark.slow
    def test_pooling_layer(self):
        self.run_case(
            layer_type='pooling',
            input_layer='data',
            output_layer='LayerPooling'
        )

    @pytest.mark.slow
    def test_lrn(self):
        self.run_case(
            layer_type='lrn',
            input_layer='data',
            output_layer='LayerLRN',
        )

    @unittest.skip(" Radar: 33056676")
    def test_mvn(self):
        self.run_case(
            layer_type='mvn',
            input_layer='data',
            output_layer='LayerMVN',
        )

    def test_reshape(self):
        self.run_case(
            layer_type='reshape',
            input_layer='data',
            output_layer='LayerReshape',
        )

    def test_embed(self):
        self.run_case(
            layer_type='embed',
            input_layer='data',
            output_layer='LayerEmbed',
        )

    def test_batchnorm(self):
        self.run_case(
            layer_type='batchnorm',
            input_layer='data',
            output_layer='LayerBatchNorm',
        )

    def test_flatten(self):
        self.run_case(
            layer_type='flatten',
            input_layer='data',
            output_layer='LayerFlatten',
        )

    def test_eltwise(self):
        self.run_case(
            layer_type='eltwise',
            input_layer=['LayerInput1', 'LayerInput2'],
            output_layer='LayerEltwise',
        )

    @unittest.skip("Radar: 32739970")
    def test_parameter(self):
        self.run_case(
            layer_type='parameter',
            input_layer='data',
            output_layer='LayerParameter',
        )

    def test_split(self):
        self.run_case(
            layer_type='split',
            input_layer='data',
            output_layer=['LayerOutput_1', 'LayerOutput_2'],
        )
<|endoftext|>


<commit_before>import asyncio
import time
import unittest

from aioes.transport import Endpoint, Transport

class TestTransport(unittest.TestCase):
    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(None)

    def tearDown(self):
        self.loop.close()

    def make_transport(self, endpoints=[{'host': 'localhost'}],
                       sniffer_interval=None):
        tr = Transport(endpoints, loop=self.loop,
                       sniffer_interval=sniffer_interval)
        self.addCleanup(tr.close)
        return tr

    def test_ctor(self):
        tr = self.make_transport()
        self.assertEqual(3, tr.max_retries)
        self.assertGreaterEqual(time.monotonic(), tr.last_sniff)
        self.assertIsNone(tr.sniffer_interval)
        self.assertAlmostEqual(0.1, tr.sniffer_timeout)
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_simple(self):
        tr = self.make_transport()

        @asyncio.coroutine
        def go():
            status, data = yield from tr.perform_request(
                'GET', '/_nodes/_all/clear')
            self.assertEqual(200, status)
            self.assertIn('nodes', data)
            # self.assertEqual(
            #     {'nodes':
            #      {'kagIbHGHS3a0dcyPmp0Jkw':
            #       {'version': '1.3.1',
            #        'ip': '127.0.1.1',
            #        'build': '2de6dc5',
            #        'name': 'Mandrill',
            #        'transport_address':
            #        'inet[/192.168.0.183:9300]',
            #        'http_address': 'inet[/192.168.0.183:9200]',
            #        'host': 'andrew-levelup'}},
            #      'cluster_name': 'elasticsearch'}, data)
        self.loop.run_until_complete(go())

    def test_set_endpoints(self):
        tr = self.make_transport([])
        self.assertEqual([], tr.endpoints)
        tr.endpoints = [{'host': 'localhost'}]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_endpoints_Endpoint(self):
        tr = self.make_transport([])
        self.assertEqual([], tr.endpoints)
        tr.endpoints = [Endpoint('localhost', 9200)]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_dont_recreate_existing_connections(self):
        tr = self.make_transport()
        connections = tr._pool.connections
        tr.endpoints = [{'host': 'localhost'}]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(connections, tr._pool.connections)

    def test_set_malformed_endpoints(self):
        tr = self.make_transport()
        with self.assertRaises(RuntimeError):
            tr.endpoints = [123]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_host_only_string(self):
        tr = self.make_transport()
        tr.endpoints = ['host']
        self.assertEqual([Endpoint('host', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_host_port_string(self):
        tr = self.make_transport()
        tr.endpoints = ['host:123']
        self.assertEqual([Endpoint('host', 123)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_sniff(self):
        tr = self.make_transport(sniffer_interval=0.001)

        @asyncio.coroutine
        def go():
            t0 = time.monotonic()
            yield from asyncio.sleep(0.001, loop=self.loop)
            yield from tr.get_connection()
            self.assertGreater(tr.last_sniff, t0)

        self.loop.run_until_complete(go())

    def test_get_connection_without_sniffing(self):
        tr = self.make_transport(sniffer_interval=1000)

        @asyncio.coroutine
        def go():
            t0 = tr.last_sniff
            yield from tr.get_connection()
            self.assertEqual(t0, tr.last_sniff)

        self.loop.run_until_complete(go())
<commit_msg>Add more tests<commit_after>import asyncio
import time
import unittest

from aioes.transport import Endpoint, Transport

class TestTransport(unittest.TestCase):
    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(None)

    def tearDown(self):
        self.loop.close()

    def make_transport(self, endpoints=[{'host': 'localhost'}],
                       sniffer_interval=None):
        tr = Transport(endpoints, loop=self.loop,
                       sniffer_interval=sniffer_interval)
        self.addCleanup(tr.close)
        return tr

    def test_ctor(self):
        tr = self.make_transport()
        self.assertEqual(3, tr.max_retries)
        self.assertGreaterEqual(time.monotonic(), tr.last_sniff)
        self.assertIsNone(tr.sniffer_interval)
        self.assertAlmostEqual(0.1, tr.sniffer_timeout)
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_simple(self):
        tr = self.make_transport()

        @asyncio.coroutine
        def go():
            status, data = yield from tr.perform_request(
                'GET', '/_nodes/_all/clear')
            self.assertEqual(200, status)
            self.assertIn('nodes', data)
            # self.assertEqual(
            #     {'nodes':
            #      {'kagIbHGHS3a0dcyPmp0Jkw':
            #       {'version': '1.3.1',
            #        'ip': '127.0.1.1',
            #        'build': '2de6dc5',
            #        'name': 'Mandrill',
            #        'transport_address':
            #        'inet[/192.168.0.183:9300]',
            #        'http_address': 'inet[/192.168.0.183:9200]',
            #        'host': 'andrew-levelup'}},
            #      'cluster_name': 'elasticsearch'}, data)
        self.loop.run_until_complete(go())

    def test_set_endpoints(self):
        tr = self.make_transport([])
        self.assertEqual([], tr.endpoints)
        tr.endpoints = [{'host': 'localhost'}]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_endpoints_Endpoint(self):
        tr = self.make_transport([])
        self.assertEqual([], tr.endpoints)
        tr.endpoints = [Endpoint('localhost', 9200)]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_dont_recreate_existing_connections(self):
        tr = self.make_transport()
        connections = tr._pool.connections
        tr.endpoints = [{'host': 'localhost'}]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(connections, tr._pool.connections)

    def test_set_malformed_endpoints(self):
        tr = self.make_transport()
        with self.assertRaises(RuntimeError):
            tr.endpoints = [123]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_host_only_string(self):
        tr = self.make_transport()
        tr.endpoints = ['host']
        self.assertEqual([Endpoint('host', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_host_port_string(self):
        tr = self.make_transport()
        tr.endpoints = ['host:123']
        self.assertEqual([Endpoint('host', 123)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_host_port_string_invalid(self):
        tr = self.make_transport()
        with self.assertRaises(RuntimeError):
            tr.endpoints = ['host:123:abc']
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_set_host_dict_invalid(self):
        tr = self.make_transport()
        with self.assertRaises(RuntimeError):
            tr.endpoints = [{'a': 'b'}]
        self.assertEqual([Endpoint('localhost', 9200)], tr.endpoints)
        self.assertEqual(1, len(tr._pool.connections))

    def test_sniff(self):
        tr = self.make_transport(sniffer_interval=0.001)

        @asyncio.coroutine
        def go():
            t0 = time.monotonic()
            yield from asyncio.sleep(0.001, loop=self.loop)
            yield from tr.get_connection()
            self.assertGreater(tr.last_sniff, t0)

        self.loop.run_until_complete(go())

    def test_get_connection_without_sniffing(self):
        tr = self.make_transport(sniffer_interval=1000)

        @asyncio.coroutine
        def go():
            t0 = tr.last_sniff
            yield from tr.get_connection()
            self.assertEqual(t0, tr.last_sniff)

        self.loop.run_until_complete(go())
<|endoftext|>


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

import sys

import web, urllib, requests
import markdown
import alignment_extension
import graph_com_ann_extension
import highlight_extension
from markdown.extensions.toc import TocExtension
import re

# read client_id and client_secret from CLI, otherwise set to 0 (causes credentials err.)
if len(sys.argv) == 4:
    client_id = sys.argv[2]
    client_secret = sys.argv[3]
else:
    client_id = 0
    client_secret = 0

    print  """
	Warning: Bad arguments.
	Usage: ./server.py <PORT> <OAUTH_CLIENT_ID> <OAUTH_CLIENT_SECRET>
	Application will run without OAuth (no ability to login).
	"""

# OAuth scopes (permissions)
scopes = ['repo', 'user']

# URL handling
urls = (
    '/', 'Index',
    '/markdown', 'Markdown',
    '/auth', 'Auth',
    '/test', 'Test'
)

# Application setup
app = web.application(urls, globals())
templates = web.template.render('templates')
# web.config.debug = False # disable debug mode because of sessions support

# Session setup
session = web.session.Session(
    app,
    web.session.DiskStore('sessions'),
    initializer={
        'access_token': 0
    }
)

class Index:
    def GET(self):
        login_link = "https://github.com/login/oauth/authorize?" + urllib.urlencode({
            "client_id": client_id,
            "scope": ','.join(scopes)
        })
        data = [
            [
                ["share", " Share", ""],
                ["export", " Export", "onclick='exportDocument()'"],
                ["print", " Print", "onclick='printDocument()'"],
                ["login", " Login", 'onclick="location.href=\'' + login_link + '\'"']
            ], [
                ["Heading 1", "H1", "onclick=\"putChar('# ', 2)\""],
                ["Heading 2", "H2", "onclick=\"putChar('## ', 3)\""],
                ["Heading 3", "H3", "onclick=\"putChar('### ', 4)\""],
                ["Heading 4", "H4", "onclick=\"putChar('#### ', 5)\""],
                ["Heading 5", "H5", "onclick=\"putChar('##### ', 6)\""],
                ["Heading 6", "H6", "onclick=\"putChar('###### ', 7)\""],
            ], [
                ["Bold", "", "onclick=\"putChar('++  ++', 3)\""],
                ["Italic", "", "onclick=\"putChar('~~  ~~', 3)\""],
                ["Underline", "", "onclick=\"putChar('__  __', 3)\""],
                ["StrikeThrough", "", "onclick=\"putChar('--  --', 3)\""],
                ["typewriting", "T", "onclick=\"putChar('```  ```', 4)\""],
            ], [
                ["align-left", "", "onclick=\"putChar('{{\\n', 3)\""],
                ["align-center", "", "onclick=\"putChar('}{\\n', 3)\""],
                ["align-block", "", "onclick=\"putChar('{}\\n', 3)\""],
                ["align-right", "", "onclick=\"putChar('}}\\n', 3)\""],
            ], [
                ["cislovany seznam", "", "onclick=\"putChar('1. ', 3)\""],
                ["odrazkovy seznam", "", "onclick=\"putChar('- ', 2)\""],
            ], [
                ["table", "", "id=\"tableButton\" onclick=\"createTable(6, 3)\""],
                ["include", "", ""],
                ["graph", "", "onclick=\"putChar('```graph\\n\\n```', 10)\""],
                ["code", "", "onclick=\"putChar('```\\n\\t\\n```', 5)\""]
            ], [
                ["preview", "Preview", "id=\"previewOpen\""],
                ["render", 'Render Mermaid', "onclick='switchMermaid();' id='mermaidBtn'"]
            ]
        ]
        return templates.index(data)

class Markdown:
    def POST(self):
        data = web.input()
        graph_com_ann_ext = graph_com_ann_extension.Extensions(data['final'], data['annotations'].split(',,,'))
        highlight_ext = highlight_extension.HighlightExtension()
        alignment_ext = alignment_extension.Extensions()
        md = markdown.Markdown(safe_mode='escape', extensions=
        [graph_com_ann_ext,  # graph, comment, annotation
         highlight_ext,  # strong, italic, underline, cross
         alignment_ext,  # alignment
         'markdown.extensions.tables',  # tables
         'markdown.extensions.sane_lists',  # using lists like in normal mardkown
         TocExtension(slugify=self.code, separator='-'),  # table of contents
         'markdown_include.include'  # option to include other files
         ])

        data = '<?xml version="1.0" encoding="utf-8" ?><reply><preview>
' + md.convert(data[
                                                                                'data']) + '
</preview><toc>' + md.toc + '</toc><comments>' + graph_com_ann_ext.comment_list + '</comments><annotations>'+graph_com_ann_ext.annotation_strings+'</annotations></reply>'
        return data

    def code(self, value, separator):
        value = re.sub(r"[^\w\s]", '', value)
        value = re.sub(r"\s+", '-', value)
        return 'header' + separator + value

class Auth:
    def GET(self):
        input = web.input()
        if 'code' in input:
            response = requests.post(
                'https://github.com/login/oauth/access_token',
                data={
                    'client_id': client_id,
                    'client_secret': client_secret,
                    'code': input.code
                },
                headers={
                    'Accept': 'application/json'
                }
            )
            json = response.json()
            if 'access_token' in json:
                session.access_token = json['access_token']
                raise web.seeother('/')
            else:
                return "Failed to get the access token"
        else:
            return "Failed to get the access code"

if __name__ == "__main__":
    app.run()
<commit_msg>use auth module; create request mappings<commit_after>#!/usr/bin/env python2

import sys
import re

import web

import markdown
import alignment_extension
import graph_com_ann_extension
import highlight_extension
from markdown.extensions.toc import TocExtension

import auth

# read client_id and client_secret from CLI, otherwise set to 0 (causes credentials err.)
if len(sys.argv) == 4:
    client_id = sys.argv[2]
    client_secret = sys.argv[3]
else:
    client_id = 0
    client_secret = 0

    print """
Warning: Bad arguments.
Usage: ./server.py <PORT> <OAUTH_CLIENT_ID> <OAUTH_CLIENT_SECRET>
Application will run without OAuth (no ability to log in).
"""

# OAuth scopes (permissions)
scopes = ['repo', 'user']

# URL handling
urls = (
    '/', 'Index',
    '/markdown', 'Markdown',
    '/auth', 'Auth',
    '/test', 'Test',
    '/list-repos', 'List_repos',
    '/list-repo-tree', 'List_repo_tree',
    '/commit-file', 'Commit_file',
    '/get-file', 'Get_file'
)

# Application setup
app = web.application(urls, globals())
templates = web.template.render('templates')
web.config.debug = False  # Must be disabled because conflicts with sessions (disable only temporarily)

# Session setup
session = web.session.Session(
    app,
    web.session.DiskStore('sessions'),
    initializer={
        'token': None
    }
)

class Index:
    def GET(self):
        login_link = auth.generate_auth_link(client_id, scopes)
        data = [
            [
                ["share", " Share", ""],
                ["export", " Export", "onclick='exportDocument()'"],
                ["print", " Print", "onclick='printDocument()'"],
                ["login", " Login", 'onclick="location.href=\'' + login_link + '\'"']
            ], [
                ["Heading 1", "H1", "onclick=\"putChar('# ', 2)\""],
                ["Heading 2", "H2", "onclick=\"putChar('## ', 3)\""],
                ["Heading 3", "H3", "onclick=\"putChar('### ', 4)\""],
                ["Heading 4", "H4", "onclick=\"putChar('#### ', 5)\""],
                ["Heading 5", "H5", "onclick=\"putChar('##### ', 6)\""],
                ["Heading 6", "H6", "onclick=\"putChar('###### ', 7)\""],
            ], [
                ["Bold", "", "onclick=\"putChar('++  ++', 3)\""],
                ["Italic", "", "onclick=\"putChar('~~  ~~', 3)\""],
                ["Underline", "", "onclick=\"putChar('__  __', 3)\""],
                ["StrikeThrough", "", "onclick=\"putChar('--  --', 3)\""],
                ["typewriting", "T", "onclick=\"putChar('```  ```', 4)\""],
            ], [
                ["align-left", "", "onclick=\"putChar('{{\\n', 3)\""],
                ["align-center", "", "onclick=\"putChar('}{\\n', 3)\""],
                ["align-block", "", "onclick=\"putChar('{}\\n', 3)\""],
                ["align-right", "", "onclick=\"putChar('}}\\n', 3)\""],
            ], [
                ["cislovany seznam", "", "onclick=\"putChar('1. ', 3)\""],
                ["odrazkovy seznam", "", "onclick=\"putChar('- ', 2)\""],
            ], [
                ["table", "", "id=\"tableButton\" onclick=\"createTable(6, 3)\""],
                ["include", "", ""],
                ["graph", "", "onclick=\"putChar('```graph\\n\\n```', 10)\""],
                ["code", "", "onclick=\"putChar('```\\n\\t\\n```', 5)\""]
            ], [
                ["preview", "Preview", "id=\"previewOpen\""],
                ["render", 'Render Mermaid', "onclick='switchMermaid();' id='mermaidBtn'"]
            ]
        ]
        return templates.index(data)

class Markdown:
    def POST(self):
        data = web.input()
        graph_com_ann_ext = graph_com_ann_extension.Extensions(data['final'], data['annotations'].split(',,,'))
        highlight_ext = highlight_extension.HighlightExtension()
        alignment_ext = alignment_extension.Extensions()
        md = markdown.Markdown(safe_mode='escape', extensions=
        [graph_com_ann_ext,  # graph, comment, annotation
         highlight_ext,  # strong, italic, underline, cross
         alignment_ext,  # alignment
         'markdown.extensions.tables',  # tables
         'markdown.extensions.sane_lists',  # using lists like in normal mardkown
         TocExtension(slugify=self.code, separator='-'),  # table of contents
         'markdown_include.include'  # option to include other files
         ])

        data = '<?xml version="1.0" encoding="utf-8" ?><reply><preview>
' + md.convert(data[
                                                                                'data']) + '
</preview><toc>' + md.toc + '</toc><comments>' + graph_com_ann_ext.comment_list + '</comments><annotations>'+graph_com_ann_ext.annotation_strings+'</annotations></reply>'
        return data

    def code(self, value, separator):
        value = re.sub(r"[^\w\s]", '', value)
        value = re.sub(r"\s+", '-', value)
        return 'header' + separator + value

class Auth:
    def GET(self):
        query = web.input()
        if 'code' in query:
            token = auth.get_auth_token(client_id, client_secret, query['code'])

            if token is None:
                return 'Login failed - no access token received.'

            session['token'] = token
            print token
            raise web.seeother('/')  # redirect users back to the editor
        else:
            return 'Login failed - no auth. code received.'

# TODO
class List_repos:
    def POST(self):
        return ''

# TODO
class List_repo_tree:
    def POST(self):
        return ''

# TODO
class Commit_file:
    def POST(self):
        return ''

# TODO
class Get_file:
    def POST(self):
        return ''

if __name__ == "__main__":
    app.run()
<|endoftext|>


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

'''
This module performs list of transformations on list of entries, and provides a
consistent input vs output interface for transformations on db entries.
'''

from __future__ import division

__author__ = "Shyue Ping Ong, Will Richards"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Mar 2, 2012"

import os
import re
import json
import datetime

from pymatgen.core.structure import Structure
from pymatgen.transformations.standard_transformations import transformation_from_dict
from pymatgen.io.cifio import CifParser
from pymatgen.io.vaspio import Poscar
from copy import deepcopy

class TransformedStructure(object):
    """
    Container object for new structures that include history of transformations.

    Each transformed structure is made up of a sequence of structures with associated
    transformation history.
    """

    def __init__(self, structure, transformations, history=None,
                 other_parameters=None):
        """
        Standard constructor for a TransformedStructure

        Args:
            structure:
                input structure
            transformations:
                Sequence of transformations to be applied to the input
                structure.
            history:
                optional history for the input structure, which provides a way
                to track structures having undergone multiple series of
                transformations.
            other_parameters:
                optional parameters to store along with the
                TransformedStructure. This can include tags (a list) or author
                which will be parsed.
        """
        history = [] if history == None else history
        self._source = {}
        self._structures = []
        self._transformations = []
        self._transformation_parameters = []
        self._redo_trans = []
        self._other_parameters = {} if other_parameters == None else deepcopy(other_parameters)
        if len(history) > 0:
            self._source = history[0]
            for i in xrange(1, len(history)):
                self._structures.append(Structure.from_dict(history[i]['input_structure']))
                self._transformations.append(transformation_from_dict(history[i]))
                self._transformation_parameters.append(history[i].get('output_parameters', {}))

        self._structures.append(structure)
        for t in transformations:
            self.append_transformation(t)

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.

        Raises:
            IndexError if already at the oldest change.
        """
        if len(self._transformations) == 0:
            raise IndexError("Can't undo. Already at oldest change.")
        self._structures.pop()
        self._transformation_parameters.pop()
        self._redo_trans.append(self._transformations.pop())

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.

        Raises:
            IndexError if already at the latest change.
        """
        if len(self._redo_trans) == 0:
            raise IndexError("Can't undo. Already at latest change.")
        t = self._redo_trans.pop()
        self.append_transformation(t, clear_redo=False)

    def __getitem__(self, index):
        return (self._structures[index], self._transformations[0:index])

    def __getattr__(self, name):
        return getattr(self._structures[-1], name)

    def __len__(self):
        return len(self._structures)

    @staticmethod
    def _alternative_transformed_structures(unmodified_transformed_structure, transformation, structure_dicts):
        for x in structure_dicts:
            new_structure = deepcopy(unmodified_transformed_structure)
            new_structure._structures.append(x.pop('structure'))
            new_structure._transformation_parameters.append(x)
            new_structure._transformations.append(transformation)
            yield new_structure

    def append_transformation(self, transformation, return_alternatives=False, clear_redo=True):
        """
        Appends a transformation to the TransformedStructure.

        Args:
            transformation:
                Transformation to append
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the redo
                list should not be cleared to allow multiple redos.
        """

        if return_alternatives:
            structures_dict_list = transformation.apply_transformation(self._structures[-1], return_ranked_list=True)
            alternative_structures = self._alternative_transformed_structures(deepcopy(self), transformation, structures_dict_list[1:])

            new_s = structures_dict_list[0]
            self._structures.append(new_s.pop('structure'))
            self._transformations.append(transformation)
            self._transformation_parameters.append(new_s)
            return alternative_structures

        else:
            new_s = transformation.apply_transformation(self._structures[-1])
            self._structures.append(new_s)
            self._transformation_parameters.append({})
            self._transformations.append(transformation)

        if clear_redo:
            self._redo_trans = []

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.

        Args:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def get_vasp_input(self, vasp_input_set, generate_potcar=True):
        """
        Returns VASP input as a dict of vaspio objects.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a POTCAR,
                which contains the POTCAR labels but not the actual POTCAR. Defaults
                to True.
        """
        d = vasp_input_set.get_all_vasp_input(self._structures[-1], generate_potcar)
        d['transformations.json'] = json.dumps(self.to_dict)
        return d

    def write_vasp_input(self, vasp_input_set, output_dir, create_directory=True):
        """
        Writes VASP input to an output_dir.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
        """
        vasp_input_set.write_input(self._structures[-1], output_dir, make_dir_if_not_present=create_directory)
        with open(os.path.join(output_dir, 'transformations.json'), 'w') as fp:
            json.dump(self.to_dict, fp)

    def __str__(self):
        output = ["Current structure"]
        output.append("------------")
        output.append(str(self._structures[-1]))
        output.append("\nSource")
        output.append("------------")
        output.append(str(self._source))
        output.append("\nTransformation history")
        output.append("------------")
        for i, t in enumerate(self._transformations):
            output.append(str(t.to_dict) + ' ' + str(self._transformation_parameters[i]))
        output.append("\nOther parameters")
        output.append("------------")
        output.append(str(self._other_parameters))
        return "\n".join(output)

    def set_parameter(self, key, value):
        self._other_parameters[key] = value

    @property
    def other_parameters(self):
        return self._other_parameters

    @property
    def was_modified(self):
        """
        boolean describing whether the last transformation on the structure made any alterations to it
        one example of when this would return false is in the case of performing a substitution transformation
        on the structure when the specie to replace isn't in the structure.
        """
        return not self._structures[-1] == self._structures[-2]

    @property
    def structures(self):
        """
        Returns a copy of all structures in the TransformedStructure. A structure
        is stored after every single transformation.
        """
        return [s for s in self._structures]

    @property
    def transformations(self):
        """
        Returns a copy of all transformations in the TransformedStructure.
        """
        return [t for t in self._transformations]

    @property
    def final_structure(self):
        """
        Returns the final structure in the TransformedStructure.
        """
        return self._structures[-1]

    @staticmethod
    def from_dict(d):
        """
        Creates a TransformedStructure from a dict.
        """
        s = Structure.from_dict(d)
        return TransformedStructure(s, [], d['history'], d.get('other_parameters', None))

    @property
    def history(self):
        history = [self._source]
        for i, t in enumerate(self._transformations):
            tdict = t.to_dict
            tdict['input_structure'] = self._structures[i].to_dict
            tdict['output_parameters'] = self._transformation_parameters[i]
            history.append(tdict)
        return history

    @property
    def to_dict(self):
        """
        Returns a dict representation of the TransformedStructure.
        """
        d = self._structures[-1].to_dict
        d['history'] = self.history
        d['version'] = __version__
        d['other_parameters'] = self._other_parameters
        return d

    @staticmethod
    def from_cif_string(cif_string, transformations=[], primitive=True):
        """
        Generates TransformedStructure from a cif string.

        Args:
            cif_string:
                Input cif string. Should contain only one structure. For cifs
                containing multiple structures, please use CifTransmuter.
            transformations:
                Sequence of transformations to be applied to the input structure.
            primitive:
                Option to set if the primitive cell should be extracted. Defaults
                to True. However, there are certain instances where you might want
                to use a non-primitive cell, e.g., if you are trying to generate
                all possible orderings of partial removals or order a disordered
                structure.
        """
        parser = CifParser.from_string(cif_string)
        raw_string = re.sub("'", "\"", cif_string)
        cif_dict = parser.to_dict
        cif_keys = cif_dict.keys()
        s = parser.get_structures(primitive)[0]
        partial_cif = cif_dict[cif_keys[0]]
        if '_database_code_ICSD' in partial_cif:
            source = partial_cif['_database_code_ICSD'] + "-ICSD"
        else:
            source = 'uploaded cif'
        source_info = {'source':source, 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string, 'cif_data':cif_dict[cif_keys[0]]}
        return TransformedStructure(s, transformations, [source_info])

    @staticmethod
    def from_poscar_string(poscar_string, transformations=[]):
        """
        Generates TransformedStructure from a poscar string.

        Args:
            poscar_string:
                Input POSCAR string.
        """
        p = Poscar.from_string(poscar_string)
        if not p.true_names:
            raise ValueError("Transformation can be craeted only from POSCAR strings with proper VASP5 element symbols.")
        raw_string = re.sub("'", "\"", poscar_string)
        s = p.struct
        source_info = {'source': "uploaded POSCAR", 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string}
        return TransformedStructure(s, transformations, [source_info])

<commit_msg>Fix doc for materials module.<commit_after>#!/usr/bin/env python

'''
This module provides various representations of transformed structures. A
TransformedStructure is a structure that has been modified by undergoing a
series of transformations.
'''

from __future__ import division

__author__ = "Shyue Ping Ong, Will Richards"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Mar 2, 2012"

import os
import re
import json
import datetime

from pymatgen.core.structure import Structure
from pymatgen.transformations.standard_transformations import transformation_from_dict
from pymatgen.io.cifio import CifParser
from pymatgen.io.vaspio import Poscar
from copy import deepcopy

class TransformedStructure(object):
    """
    Container object for new structures that include history of transformations.

    Each transformed structure is made up of a sequence of structures with associated
    transformation history.
    """

    def __init__(self, structure, transformations, history=None,
                 other_parameters=None):
        """
        Standard constructor for a TransformedStructure

        Args:
            structure:
                input structure
            transformations:
                Sequence of transformations to be applied to the input
                structure.
            history:
                optional history for the input structure, which provides a way
                to track structures having undergone multiple series of
                transformations.
            other_parameters:
                optional parameters to store along with the
                TransformedStructure. This can include tags (a list) or author
                which will be parsed.
        """
        history = [] if history == None else history
        self._source = {}
        self._structures = []
        self._transformations = []
        self._transformation_parameters = []
        self._redo_trans = []
        self._other_parameters = {} if other_parameters == None else deepcopy(other_parameters)
        if len(history) > 0:
            self._source = history[0]
            for i in xrange(1, len(history)):
                self._structures.append(Structure.from_dict(history[i]['input_structure']))
                self._transformations.append(transformation_from_dict(history[i]))
                self._transformation_parameters.append(history[i].get('output_parameters', {}))

        self._structures.append(structure)
        for t in transformations:
            self.append_transformation(t)

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.

        Raises:
            IndexError if already at the oldest change.
        """
        if len(self._transformations) == 0:
            raise IndexError("Can't undo. Already at oldest change.")
        self._structures.pop()
        self._transformation_parameters.pop()
        self._redo_trans.append(self._transformations.pop())

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.

        Raises:
            IndexError if already at the latest change.
        """
        if len(self._redo_trans) == 0:
            raise IndexError("Can't undo. Already at latest change.")
        t = self._redo_trans.pop()
        self.append_transformation(t, clear_redo=False)

    def __getitem__(self, index):
        return (self._structures[index], self._transformations[0:index])

    def __getattr__(self, name):
        return getattr(self._structures[-1], name)

    def __len__(self):
        return len(self._structures)

    @staticmethod
    def _alternative_transformed_structures(unmodified_transformed_structure, transformation, structure_dicts):
        for x in structure_dicts:
            new_structure = deepcopy(unmodified_transformed_structure)
            new_structure._structures.append(x.pop('structure'))
            new_structure._transformation_parameters.append(x)
            new_structure._transformations.append(transformation)
            yield new_structure

    def append_transformation(self, transformation, return_alternatives=False, clear_redo=True):
        """
        Appends a transformation to the TransformedStructure.

        Args:
            transformation:
                Transformation to append
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the redo
                list should not be cleared to allow multiple redos.
        """

        if return_alternatives:
            structures_dict_list = transformation.apply_transformation(self._structures[-1], return_ranked_list=True)
            alternative_structures = self._alternative_transformed_structures(deepcopy(self), transformation, structures_dict_list[1:])

            new_s = structures_dict_list[0]
            self._structures.append(new_s.pop('structure'))
            self._transformations.append(transformation)
            self._transformation_parameters.append(new_s)
            return alternative_structures

        else:
            new_s = transformation.apply_transformation(self._structures[-1])
            self._structures.append(new_s)
            self._transformation_parameters.append({})
            self._transformations.append(transformation)

        if clear_redo:
            self._redo_trans = []

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.

        Args:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def get_vasp_input(self, vasp_input_set, generate_potcar=True):
        """
        Returns VASP input as a dict of vaspio objects.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a POTCAR,
                which contains the POTCAR labels but not the actual POTCAR. Defaults
                to True.
        """
        d = vasp_input_set.get_all_vasp_input(self._structures[-1], generate_potcar)
        d['transformations.json'] = json.dumps(self.to_dict)
        return d

    def write_vasp_input(self, vasp_input_set, output_dir, create_directory=True):
        """
        Writes VASP input to an output_dir.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
        """
        vasp_input_set.write_input(self._structures[-1], output_dir, make_dir_if_not_present=create_directory)
        with open(os.path.join(output_dir, 'transformations.json'), 'w') as fp:
            json.dump(self.to_dict, fp)

    def __str__(self):
        output = ["Current structure"]
        output.append("------------")
        output.append(str(self._structures[-1]))
        output.append("\nSource")
        output.append("------------")
        output.append(str(self._source))
        output.append("\nTransformation history")
        output.append("------------")
        for i, t in enumerate(self._transformations):
            output.append(str(t.to_dict) + ' ' + str(self._transformation_parameters[i]))
        output.append("\nOther parameters")
        output.append("------------")
        output.append(str(self._other_parameters))
        return "\n".join(output)

    def set_parameter(self, key, value):
        self._other_parameters[key] = value

    @property
    def other_parameters(self):
        return self._other_parameters

    @property
    def was_modified(self):
        """
        boolean describing whether the last transformation on the structure made any alterations to it
        one example of when this would return false is in the case of performing a substitution transformation
        on the structure when the specie to replace isn't in the structure.
        """
        return not self._structures[-1] == self._structures[-2]

    @property
    def structures(self):
        """
        Returns a copy of all structures in the TransformedStructure. A structure
        is stored after every single transformation.
        """
        return [s for s in self._structures]

    @property
    def transformations(self):
        """
        Returns a copy of all transformations in the TransformedStructure.
        """
        return [t for t in self._transformations]

    @property
    def final_structure(self):
        """
        Returns the final structure in the TransformedStructure.
        """
        return self._structures[-1]

    @staticmethod
    def from_dict(d):
        """
        Creates a TransformedStructure from a dict.
        """
        s = Structure.from_dict(d)
        return TransformedStructure(s, [], d['history'], d.get('other_parameters', None))

    @property
    def history(self):
        history = [self._source]
        for i, t in enumerate(self._transformations):
            tdict = t.to_dict
            tdict['input_structure'] = self._structures[i].to_dict
            tdict['output_parameters'] = self._transformation_parameters[i]
            history.append(tdict)
        return history

    @property
    def to_dict(self):
        """
        Returns a dict representation of the TransformedStructure.
        """
        d = self._structures[-1].to_dict
        d['history'] = self.history
        d['version'] = __version__
        d['other_parameters'] = self._other_parameters
        return d

    @staticmethod
    def from_cif_string(cif_string, transformations=[], primitive=True):
        """
        Generates TransformedStructure from a cif string.

        Args:
            cif_string:
                Input cif string. Should contain only one structure. For cifs
                containing multiple structures, please use CifTransmuter.
            transformations:
                Sequence of transformations to be applied to the input structure.
            primitive:
                Option to set if the primitive cell should be extracted. Defaults
                to True. However, there are certain instances where you might want
                to use a non-primitive cell, e.g., if you are trying to generate
                all possible orderings of partial removals or order a disordered
                structure.
        """
        parser = CifParser.from_string(cif_string)
        raw_string = re.sub("'", "\"", cif_string)
        cif_dict = parser.to_dict
        cif_keys = cif_dict.keys()
        s = parser.get_structures(primitive)[0]
        partial_cif = cif_dict[cif_keys[0]]
        if '_database_code_ICSD' in partial_cif:
            source = partial_cif['_database_code_ICSD'] + "-ICSD"
        else:
            source = 'uploaded cif'
        source_info = {'source':source, 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string, 'cif_data':cif_dict[cif_keys[0]]}
        return TransformedStructure(s, transformations, [source_info])

    @staticmethod
    def from_poscar_string(poscar_string, transformations=[]):
        """
        Generates TransformedStructure from a poscar string.

        Args:
            poscar_string:
                Input POSCAR string.
        """
        p = Poscar.from_string(poscar_string)
        if not p.true_names:
            raise ValueError("Transformation can be craeted only from POSCAR strings with proper VASP5 element symbols.")
        raw_string = re.sub("'", "\"", poscar_string)
        s = p.struct
        source_info = {'source': "uploaded POSCAR", 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string}
        return TransformedStructure(s, transformations, [source_info])

<|endoftext|>


<commit_before><commit_msg>Using Agg for plotting tests.<commit_after><|endoftext|>


<commit_before>import copy

from cornflake import fields
from cornflake import serializers
from cornflake.exceptions import ValidationError
from cornflake.sqlalchemy_orm import ModelSerializer, ReferenceField

from radar.api.serializers.common import MetaMixin, PatientMixin
from radar.marmoset.schema import Schema
from radar.marmoset.registry import Registry
from radar.models import Entry, Form

class FormSerializer(ModelSerializer):
    name = fields.StringField()
    data = fields.Field()

    class Meta(object):
        model_class = Form

class FormField(ReferenceField):
    model_class = Form
    serializer_class = FormSerializer

class DataField(fields.Field):
    def __init__(self, schema, **kwargs):
        super(DataField, self).__init__(**kwargs)
        self.schema = schema

    def __deepcopy__(self, memo):
        kwargs = copy.deepcopy(self._kwargs)
        return self.__class__(self.schema, **kwargs)

    def to_internal_value(self, data):
        data = self.schema.validate(data)
        data = self.schema.format(data)
        return data

    def to_representation(self, data):
        # Data is stored serialized so we need to deserialize it first
        # Note: currently we could just return the stored data here
        # We may want to use a different serialization format for the API and storage
        data = self.schema.parse(data)
        data = self.schema.format(data)
        return data

class BaseEntrySerializer(PatientMixin, MetaMixin, ModelSerializer):
    form = FormField()

    class Meta(object):
        model_class = Entry
        exclude = ['form_id']

    def _save(self, instance, data):
        instance.form = data['form']
        instance.patient = data['patient']

        # Serialize the data before saving (e.g. datetime to string)
        instance.data = self.data.to_representation(data['data'])

        # Metadata
        instance.created_user = data['created_user']
        instance.modified_user = data['modified_user']
        instance.created_date = data['created_date']
        instance.modified_date = data['modified_date']

    def create(self, data):
        instance = Entry()
        self._save(instance, data)
        return instance

    def update(self, instance, data):
        self._save(instance, data)
        return instance

class EntrySerializer(serializers.ProxySerializer):
    def __init__(self, *args, **kwargs):
        super(EntrySerializer, self).__init__(*args, **kwargs)

        # Entry must have a form so we know how to handle the data property/attribute
        self.field = FormField()
        self.field.bind(self, 'form')

    def create_serializer(self, form):
        registry = Registry()
        schema = Schema(registry, form.data)

        data_field = DataField(schema)

        serializer = type('EntrySerializer', (BaseEntrySerializer,), {
            'data': data_field
        })()

        return serializer

    def get_serializer(self, data):
        form = self.field.get_attribute(data)
        serializer = self.create_serializer(form)
        return serializer

    def get_deserializer(self, data):
        form_data = self.field.get_value(data)

        try:
            form = self.field.run_validation(form_data)
        except ValidationError as e:
            raise ValidationError({self.field.field_name: e.errors})

        serializer = self.create_serializer(form)

        return serializer
<commit_msg>Fix field name<commit_after>import copy

from cornflake import fields
from cornflake import serializers
from cornflake.exceptions import ValidationError
from cornflake.sqlalchemy_orm import ModelSerializer, ReferenceField

from radar.api.serializers.common import MetaMixin, PatientMixin
from radar.marmoset.schema import Schema
from radar.marmoset.registry import Registry
from radar.models import Entry, Form

class FormSerializer(ModelSerializer):
    name = fields.StringField()
    data_ = fields.Field(field_name='data', source='data')

    class Meta(object):
        model_class = Form

class FormField(ReferenceField):
    model_class = Form
    serializer_class = FormSerializer

class DataField(fields.Field):
    def __init__(self, schema, **kwargs):
        super(DataField, self).__init__(**kwargs)
        self.schema = schema

    def __deepcopy__(self, memo):
        kwargs = copy.deepcopy(self._kwargs)
        return self.__class__(self.schema, **kwargs)

    def to_internal_value(self, data):
        data = self.schema.validate(data)
        data = self.schema.format(data)
        return data

    def to_representation(self, data):
        # Data is stored serialized so we need to deserialize it first
        # Note: currently we could just return the stored data here
        # We may want to use a different serialization format for the API and storage
        data = self.schema.parse(data)
        data = self.schema.format(data)
        return data

class BaseEntrySerializer(PatientMixin, MetaMixin, ModelSerializer):
    form = FormField()

    class Meta(object):
        model_class = Entry
        exclude = ['form_id']

    def _save(self, instance, data):
        instance.form = data['form']
        instance.patient = data['patient']

        # Serialize the data before saving (e.g. datetime to string)
        instance.data = self.data.to_representation(data['data'])

        # Metadata
        instance.created_user = data['created_user']
        instance.modified_user = data['modified_user']
        instance.created_date = data['created_date']
        instance.modified_date = data['modified_date']

    def create(self, data):
        instance = Entry()
        self._save(instance, data)
        return instance

    def update(self, instance, data):
        self._save(instance, data)
        return instance

class EntrySerializer(serializers.ProxySerializer):
    def __init__(self, *args, **kwargs):
        super(EntrySerializer, self).__init__(*args, **kwargs)

        # Entry must have a form so we know how to handle the data property/attribute
        self.field = FormField()
        self.field.bind(self, 'form')

    def create_serializer(self, form):
        registry = Registry()
        schema = Schema(registry, form.data)

        data_field = DataField(schema)

        serializer = type('EntrySerializer', (BaseEntrySerializer,), {
            'data': data_field
        })()

        return serializer

    def get_serializer(self, data):
        form = self.field.get_attribute(data)
        serializer = self.create_serializer(form)
        return serializer

    def get_deserializer(self, data):
        form_data = self.field.get_value(data)

        try:
            form = self.field.run_validation(form_data)
        except ValidationError as e:
            raise ValidationError({self.field.field_name: e.errors})

        serializer = self.create_serializer(form)

        return serializer
<|endoftext|>


<commit_before><commit_msg>use @staticmethod<commit_after><|endoftext|>


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

import collectd

import base
import requests

NAME = 'elasticsearch_cluster'
HEALTH_MAP = {
    'green': 1,
    'yellow': 2,
    'red': 3,
}
METRICS = ['number_of_nodes', 'active_primary_shards', 'active_primary_shards',
           'active_shards', 'relocating_shards', 'unassigned_shards',
           'number_of_pending_tasks', 'initializing_shards']

HEALTH_ON_ERROR = {'type_instance': 'health', 'values': HEALTH_MAP['red']}

class ElasticsearchClusterHealthPlugin(base.Base):
    def __init__(self, *args, **kwargs):
        super(ElasticsearchClusterHealthPlugin, self).__init__(*args, **kwargs)
        self.plugin = NAME
        self.address = '127.0.0.1'
        self.port = 9200
        self.session = requests.Session()
        self.url = None
        self.session.mount(
            'http://',
            requests.adapters.HTTPAdapter(max_retries=self.max_retries)
        )

    def config_callback(self, conf):
        super(ElasticsearchClusterHealthPlugin, self).config_callback(conf)

        for node in conf.children:
            if node.key == 'Address':
                self.address = node.values[0]
            if node.key == 'Port':
                self.port = node.values[0]

        self.url = "http://{address}:{port}/_cluster/health".format(
            **{
                'address': self.address,
                'port': int(self.port),
            })

    def itermetrics(self):
        try:
            r = self.session.get(self.url)
        except Exception as e:
            self.logger.error("Got exception for '{}': {}".format(self.url, e))
            yield HEALTH_ON_ERROR
            return

        if r.status_code != 200:
            self.logger.error("{} responded with code {}".format(
                self.url, r.status_code))
            yield HEALTH_ON_ERROR
            return
        data = r.json()
        self.logger.debug("Got response from Elasticsearch: '%s'" % data)

        yield {
            'type_instance': 'health',
            'values': HEALTH_MAP[data['status']]
        }
        for metric in METRICS:
            yield {
                'type_instance': metric,
                'values': data[metric]
            }

plugin = ElasticsearchClusterHealthPlugin(collectd)

def init_callback():
    plugin.restore_sigchld()

def config_callback(conf):
    plugin.config_callback(conf)

def read_callback():
    plugin.read_callback()

collectd.register_init(init_callback)
collectd.register_config(config_callback)
collectd.register_read(read_callback)
<commit_msg>Fix Elasticsearch collectd plugin for old versions<commit_after>#!/usr/bin/python
# Copyright 2016 Mirantis, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import collectd

import base
import requests

NAME = 'elasticsearch_cluster'
HEALTH_MAP = {
    'green': 1,
    'yellow': 2,
    'red': 3,
}
METRICS = ['number_of_nodes', 'active_primary_shards', 'active_primary_shards',
           'active_shards', 'relocating_shards', 'unassigned_shards',
           'number_of_pending_tasks', 'initializing_shards']

HEALTH_ON_ERROR = {'type_instance': 'health', 'values': HEALTH_MAP['red']}

class ElasticsearchClusterHealthPlugin(base.Base):
    def __init__(self, *args, **kwargs):
        super(ElasticsearchClusterHealthPlugin, self).__init__(*args, **kwargs)
        self.plugin = NAME
        self.address = '127.0.0.1'
        self.port = 9200
        self.session = requests.Session()
        self.url = None
        self.session.mount(
            'http://',
            requests.adapters.HTTPAdapter(max_retries=self.max_retries)
        )

    def config_callback(self, conf):
        super(ElasticsearchClusterHealthPlugin, self).config_callback(conf)

        for node in conf.children:
            if node.key == 'Address':
                self.address = node.values[0]
            if node.key == 'Port':
                self.port = node.values[0]

        self.url = "http://{address}:{port}/_cluster/health".format(
            **{
                'address': self.address,
                'port': int(self.port),
            })

    def itermetrics(self):
        try:
            r = self.session.get(self.url)
        except Exception as e:
            self.logger.error("Got exception for '{}': {}".format(self.url, e))
            yield HEALTH_ON_ERROR
            return

        if r.status_code != 200:
            self.logger.error("{} responded with code {}".format(
                self.url, r.status_code))
            yield HEALTH_ON_ERROR
            return
        data = r.json()
        self.logger.debug("Got response from Elasticsearch: '%s'" % data)

        yield {
            'type_instance': 'health',
            'values': HEALTH_MAP[data['status']]
        }
        for metric in METRICS:
            value = data.get(metric)
            if not value:
                # Depending on the Elasticsearch version, not all metrics are
                # available
                self.logger.info("Couldn't find {} metric".format(metric))
                continue
            yield {
                'type_instance': metric,
                'values': value
            }

plugin = ElasticsearchClusterHealthPlugin(collectd)

def init_callback():
    plugin.restore_sigchld()

def config_callback(conf):
    plugin.config_callback(conf)

def read_callback():
    plugin.read_callback()

collectd.register_init(init_callback)
collectd.register_config(config_callback)
collectd.register_read(read_callback)
<|endoftext|>