<commit_before>#!/usr/bin/python
#
#  Copyright 2002-2018 Barcelona Supercomputing Center (www.bsc.es)
#
#  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.
#

# -*- coding: utf-8 -*-

"""
PyCOMPSs Utils: Data serializer/deserializer
This file implements the main serialization/deserialization functions.
All serialization/deserialization calls should be made using one of the following functions:

- serialize_to_file(obj, file_name) -> dumps the object "obj" to the file "file_name"
- serialize_to_string(obj) -> dumps the object "obj" to a string
- serialize_to_handler(obj, handler) -> writes the serialized object using the specified handler
                                        it also moves the handler's pointer to the end of the dump

- deserialize_from_file(file_name) -> loads the first object from the tile "file_name"
- deserialize_from_string(serialized_content) -> loads the first object from the given string
- deserialize_from_handler(handler) -> deserializes an object using the given handler, it also leaves the
                                       handler's pointer pointing to the end of the serialized object
"""

import types
import traceback

from pycompss.runtime.commons import IS_PYTHON3
from .serialization.extended_support import pickle_generator, convert_to_generator, GeneratorIndicator
from .object_properties import object_belongs_to_module

from io import BytesIO

if IS_PYTHON3:
    import pickle as pickle  # Uses _pickle if available
else:
    import cPickle as pickle

try:
    import dill
except ImportError:
    if IS_PYTHON3:
        import pickle as dill
    else:
        import cPickle as dill

try:
    import numpy
    NUMPY_AVAILABLE = True
except ImportError:
    if IS_PYTHON3:
        import pickle as numpy
    else:
        import cPickle as numpy
    NUMPY_AVAILABLE = False

lib2idx = {
    pickle: 0,
    numpy: 1,
    dill: 2
}

idx2lib = dict([(v, k) for (k, v) in lib2idx.items()])

class SerializerException(Exception):
    """
    Exception on serialization
    """

    pass

def get_serializer_priority(obj=()):
    """
    Computes the priority of the serializers.

    :param obj: Object to be analysed.
    :return: <List> The serializers sorted by priority in descending order
    """

    if object_belongs_to_module(obj, 'numpy'):
        return [numpy, pickle, dill]
    return [pickle, dill]

def get_serializers():
    """
    Returns a list with the available serializers in the most common order
    (i.e: the order that will work for almost the 90% of our objects).

    :return: <List> the available serializer modules
    """

    return get_serializer_priority()

def serialize_to_handler(obj, handler):
    """
    Serialize an object to a handler.

    :param obj: Object to be serialized.
    :param handler: A handler object. It must implement methods like write, writeline and similar stuff
    """

    # Get the serializer priority
    serializer_priority = get_serializer_priority(obj)
    i = 0
    success = False
    original_position = handler.tell()
    # Lets try the serializers in the given priority
    while i < len(serializer_priority) and not success:
        # Reset the handlers pointer to the first position
        handler.seek(original_position)
        serializer = serializer_priority[i]
        handler.write(bytearray('%04d' % lib2idx[serializer], 'utf8'))
        # Special case: obj is a generator
        if isinstance(obj, types.GeneratorType):
            try:
                pickle_generator(obj, handler, serializer)
                success = True
            except Exception:
                if __debug__:
                    traceback.print_exc()
        # General case
        else:
            try:
                # If it is a numpy object then use its saving mechanism
                if serializer is numpy and NUMPY_AVAILABLE:
                    serializer.save(handler, obj)
                else:
                    serializer.dump(obj, handler, protocol=serializer.HIGHEST_PROTOCOL)
                success = True
            except Exception:
                if __debug__:
                    traceback.print_exc()  # No need to print all serializer exceptions
                    print('WARNING! Serialization with %s failed.' % str(serializer))
        i += 1

    # if ret_value is None then all the serializers have failed
    if not success:
        try:
            traceback.print_exc()
        except AttributeError:
            # Bug fixed in 3.5 - issue10805
            pass
        raise SerializerException('Cannot serialize object %s' % obj)

def serialize_to_file(obj, file_name):
    """
    Serialize an object to a file.

    :param obj: Object to be serialized.
    :param file_name: File name where the object is going to be serialized.
    :return: Nothing, it just serializes the object
    """

    handler = open(file_name, 'wb')
    serialize_to_handler(obj, handler)
    handler.close()
    return file_name

def serialize_to_string(obj):
    """
    Serialize an object to a string.

    :param obj: Object to be serialized.
    :return: <String> the serialized content
    """

    handler = BytesIO()
    serialize_to_handler(obj, handler)
    ret = handler.getvalue()
    handler.close()
    return ret

def serialize_named_to_file(name, obj, file_name):
    '''Serialize a named object to a file
    :param name: Name of the object
    :param obj: Name of the object
    :param file_name: Name of the file name
    :return: Nothing, it just serializes the named object
    '''
    serialize_to_file((name, obj), file_name)

def serialize_named_to_string(name, obj):
    '''Serialize a named object to a string
    :param name: Name of the object
    :param obj: Object to serialize
    :return: <String> the serialized content
    '''
    return serialize_to_string((name, obj))

def get_numpy_dummy_obj():
    '''Return a numpy object or an empty list if numpy is not available
    '''
    try:
        return numpy.zeros(1)
    except:
        return []

def deserialize_from_handler(handler):
    '''Deserialize an object from a file.
    :param handler: File name from where the object is going to be deserialized.
    :return: The object deserialized.
    '''
    # Retrieve the used library
    serializer = idx2lib[int(handler.read(4))]
    try:
        ret = serializer.load(handler)
        # Special case: deserialized obj wraps a generator
        if isinstance(ret, tuple) and ret and isinstance(ret[0], GeneratorIndicator):
            ret = convert_to_generator(ret[1])
        return ret
    except Exception:
        if __debug__:
            print('WARNING! Deserialization with %s failed.' % str(serializer))
            traceback.print_exc()  # No need to print all deserialize exceptions
    # We are not able to deserialize the contents from file_name with any of our
    # serializers
    try:
        traceback.print_exc()
    except AttributeError:
        # Bug fixed in 3.5 - issue10805
        pass
    raise SerializerException('Cannot deserialize object')

def deserialize_from_file(file_name):
    '''Deserialize the contents in a given file.
    :param file_name: Name of the file with the contents to be deserialized
    :return: A deserialized object
    '''

    handler = open(file_name, 'rb')
    ret = deserialize_from_handler(handler)
    handler.close()
    return ret

def deserialize_from_string(serialized_content):
    '''
    Deserialize the contents in a given string.
    :param serialized_content: A string with serialized contents
    :return: A deserialized object
    '''
    handler = BytesIO(serialized_content)
    ret = deserialize_from_handler(handler)
    handler.close()
    return ret

def serialize_objects(to_serialize):
    '''
    Serialize a list of objects to file. If a single object fails to be serialized, then an Exception by
    serialize_to_file will be thrown (and not caught)
    The structure of the parameter is [(object1, file_name1), ... , (objectN, file_nameN)].
    :param to_serialize: List of lists to be serialized. Each sublist is a pair of the form ['object','file name']
    :return:
    '''
    for obj_and_file in to_serialize:
        serialize_to_file(*obj_and_file)
<commit_msg>Fixed KeyError Deserialization bug<commit_after>#!/usr/bin/python
#
#  Copyright 2002-2018 Barcelona Supercomputing Center (www.bsc.es)
#
#  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.
#

# -*- coding: utf-8 -*-

"""
PyCOMPSs Utils: Data serializer/deserializer
This file implements the main serialization/deserialization functions.
All serialization/deserialization calls should be made using one of the following functions:

- serialize_to_file(obj, file_name) -> dumps the object "obj" to the file "file_name"
- serialize_to_string(obj) -> dumps the object "obj" to a string
- serialize_to_handler(obj, handler) -> writes the serialized object using the specified handler
                                        it also moves the handler's pointer to the end of the dump

- deserialize_from_file(file_name) -> loads the first object from the tile "file_name"
- deserialize_from_string(serialized_content) -> loads the first object from the given string
- deserialize_from_handler(handler) -> deserializes an object using the given handler, it also leaves the
                                       handler's pointer pointing to the end of the serialized object
"""

import types
import traceback

from pycompss.runtime.commons import IS_PYTHON3
from .serialization.extended_support import pickle_generator, convert_to_generator, GeneratorIndicator
from .object_properties import object_belongs_to_module

from io import BytesIO

if IS_PYTHON3:
    import pickle as pickle  # Uses _pickle if available
else:
    import cPickle as pickle

try:
    import dill
except ImportError:
    if IS_PYTHON3:
        import pickle as dill
    else:
        import cPickle as dill

try:
    import numpy
    NUMPY_AVAILABLE = True
except ImportError:
    if IS_PYTHON3:
        import pickle as numpy
    else:
        import cPickle as numpy
    NUMPY_AVAILABLE = False

lib2idx = {
    pickle: 0,
    numpy: 1,
    dill: 2
}

idx2lib = dict([(v, k) for (k, v) in lib2idx.items()])

class SerializerException(Exception):
    """
    Exception on serialization
    """

    pass

def get_serializer_priority(obj=()):
    """
    Computes the priority of the serializers.

    :param obj: Object to be analysed.
    :return: <List> The serializers sorted by priority in descending order
    """

    if object_belongs_to_module(obj, 'numpy'):
        return [numpy, pickle, dill]
    return [pickle, dill]

def get_serializers():
    """
    Returns a list with the available serializers in the most common order
    (i.e: the order that will work for almost the 90% of our objects).

    :return: <List> the available serializer modules
    """

    return get_serializer_priority()

def serialize_to_handler(obj, handler):
    """
    Serialize an object to a handler.

    :param obj: Object to be serialized.
    :param handler: A handler object. It must implement methods like write, writeline and similar stuff
    """

    # Get the serializer priority
    serializer_priority = get_serializer_priority(obj)
    i = 0
    success = False
    original_position = handler.tell()
    # Lets try the serializers in the given priority
    while i < len(serializer_priority) and not success:
        # Reset the handlers pointer to the first position
        handler.seek(original_position)
        serializer = serializer_priority[i]
        handler.write(bytearray('%04d' % lib2idx[serializer], 'utf8'))
        # Special case: obj is a generator
        if isinstance(obj, types.GeneratorType):
            try:
                pickle_generator(obj, handler, serializer)
                success = True
            except Exception:
                if __debug__:
                    traceback.print_exc()
        # General case
        else:
            try:
                # If it is a numpy object then use its saving mechanism
                if serializer is numpy and NUMPY_AVAILABLE:
                    serializer.save(handler, obj)
                else:
                    serializer.dump(obj, handler, protocol=serializer.HIGHEST_PROTOCOL)
                success = True
            except Exception:
                if __debug__:
                    traceback.print_exc()  # No need to print all serializer exceptions
                    print('WARNING! Serialization with %s failed.' % str(serializer))
        i += 1

    # if ret_value is None then all the serializers have failed
    if not success:
        try:
            traceback.print_exc()
        except AttributeError:
            # Bug fixed in 3.5 - issue10805
            pass
        raise SerializerException('Cannot serialize object %s' % obj)

def serialize_to_file(obj, file_name):
    """
    Serialize an object to a file.

    :param obj: Object to be serialized.
    :param file_name: File name where the object is going to be serialized.
    :return: Nothing, it just serializes the object
    """

    handler = open(file_name, 'wb')
    serialize_to_handler(obj, handler)
    handler.close()
    return file_name

def serialize_to_string(obj):
    """
    Serialize an object to a string.

    :param obj: Object to be serialized.
    :return: <String> the serialized content
    """

    handler = BytesIO()
    serialize_to_handler(obj, handler)
    ret = handler.getvalue()
    handler.close()
    return ret

def serialize_named_to_file(name, obj, file_name):
    '''Serialize a named object to a file
    :param name: Name of the object
    :param obj: Name of the object
    :param file_name: Name of the file name
    :return: Nothing, it just serializes the named object
    '''
    serialize_to_file((name, obj), file_name)

def serialize_named_to_string(name, obj):
    '''Serialize a named object to a string
    :param name: Name of the object
    :param obj: Object to serialize
    :return: <String> the serialized content
    '''
    return serialize_to_string((name, obj))

def get_numpy_dummy_obj():
    '''Return a numpy object or an empty list if numpy is not available
    '''
    try:
        return numpy.zeros(1)
    except:
        return []

def deserialize_from_handler(handler):
    '''Deserialize an object from a file.
    :param handler: File name from where the object is going to be deserialized.
    :return: The object deserialized.
    '''
    # Retrieve the used library (if possible)
    try:
        original_position = handler.tell()
        serializer = idx2lib[int(handler.read(4))]
    except:
        handler.seek(original_position)
        raise SerializerException('Handler does not refer to a valid PyCOMPSs object')

    try:
        ret = serializer.load(handler)
        # Special case: deserialized obj wraps a generator
        if isinstance(ret, tuple) and ret and isinstance(ret[0], GeneratorIndicator):
            ret = convert_to_generator(ret[1])
        return ret
    except Exception:
        if __debug__:
            print('WARNING! Deserialization with %s failed.' % str(serializer))
            traceback.print_exc()  # No need to print all deserialize exceptions
    # We are not able to deserialize the contents from file_name with any of our
    # serializers
    try:
        traceback.print_exc()
    except AttributeError:
        # Bug fixed in 3.5 - issue10805
        pass
    raise SerializerException('Cannot deserialize object')

def deserialize_from_file(file_name):
    '''Deserialize the contents in a given file.
    :param file_name: Name of the file with the contents to be deserialized
    :return: A deserialized object
    '''

    handler = open(file_name, 'rb')
    ret = deserialize_from_handler(handler)
    handler.close()
    return ret

def deserialize_from_string(serialized_content):
    '''
    Deserialize the contents in a given string.
    :param serialized_content: A string with serialized contents
    :return: A deserialized object
    '''
    handler = BytesIO(serialized_content)
    ret = deserialize_from_handler(handler)
    handler.close()
    return ret

def serialize_objects(to_serialize):
    '''
    Serialize a list of objects to file. If a single object fails to be serialized, then an Exception by
    serialize_to_file will be thrown (and not caught)
    The structure of the parameter is [(object1, file_name1), ... , (objectN, file_nameN)].
    :param to_serialize: List of lists to be serialized. Each sublist is a pair of the form ['object','file name']
    :return:
    '''
    for obj_and_file in to_serialize:
        serialize_to_file(*obj_and_file)
<|endoftext|>


<commit_before>import flow.petri.safenet as sn

from test_helpers import RedisTest, FakeOrchestrator
import mock
import os
import redis
import sys
import unittest

class TestBase(RedisTest):
    def setUp(self):
        RedisTest.setUp(self)
        orch = FakeOrchestrator(self.conn)
        self.services = orch.services

class TestTransition(TestBase):
    def test_default_token_merger(self):
        tokens = []
        expected = {}
        for i in xrange(10):
            expected[str(i)] = i
            tokens.append(sn.Token.create(self.conn, data={i: i}))

        token_keys = [x.key for x in tokens]
        transition = sn._SafeTransition.create(self.conn,
                active_tokens=token_keys)
        self.assertEqual(expected, transition.input_data)

    def test_abstract_base(self):
        merger = sn.TokenMerger.create(self.conn)
        self.assertRaises(NotImplementedError, merger.merge, [])

    def test_custom_token_merger(self):
        class PrependingTokenMerger(sn.TokenMerger):
            def merge(self, tokens):
                data = {}
                for t in tokens:
                    data.update({"x" + k: v for k, v in t.data.iteritems()})
                return data

        tokens = []
        expected = {}
        for i in xrange(10):
            expected["x" + str(i)] = i
            tokens.append(sn.Token.create(self.conn, data={i: i}))

        token_keys = [x.key for x in tokens]
        merger = PrependingTokenMerger.create(self.conn)
        transition = sn._SafeTransition.create(self.conn,
                active_tokens=token_keys,
                token_merger_key=merger.key)
        self.assertEqual(expected, transition.input_data)

class TestSafeNet(TestBase):
    def test_no_connection(self):
        self.assertRaises(TypeError, sn.SafeNet.create, None)

    def test_abstract_transition_action(self):
        act = sn.TransitionAction.create(self.conn, name="boom")
        self.assertRaises(NotImplementedError, act.execute, input_data={},
                net=None, services=None)

    def test_constants(self):
        net = sn.SafeNet.create(connection=self.conn)
        env = {"PATH": "/bin:/usr/bin", "USER": "flow"}
        net.set_constant("environment", env)
        self.assertEqual(env, net.constant("environment"))
        self.assertRaises(TypeError, net.set_constant, "environment", 10)

        user_id = 1234
        no_user_id = net.constant("user_id")
        self.assertIsNone(no_user_id)
        net.set_constant("user_id", user_id)
        self.assertRaises(TypeError, net.set_constant, "user_id", 10)
        self.assertEqual(1234, net.constant("user_id"))

    def test_variables(self):
        net = sn.SafeNet.create(connection=self.conn)

        foo = net.variable("foo")
        self.assertIsNone(foo)
        net.set_variable("foo", 123)
        self.assertEqual(123, net.variable("foo"))

        bar = net.variable("bar")
        self.assertIsNone(bar)
        net.set_variable("bar", {"x": "y"})
        self.assertEqual({"x": "y"}, net.variable("bar"))

        baz = net.variable("baz")
        self.assertIsNone(baz)
        net.set_variable("bar", [1, 2, 3])
        self.assertEqual([1, 2, 3], net.variable("bar"))

    def test_places(self):
        action = sn.CounterAction.create(connection=self.conn, name="counter")
        net = sn.SafeNet.create(
                connection=self.conn,
                place_names=["a", "b", "c"],
                trans_actions = [action],
                place_arcs_out={0: [0]}, # a -> t1
                trans_arcs_out={0: [1, 2]}, # t1 -> b, c
                )

        self.assertEqual(3, net.num_places)
        self.assertEqual("a", str(net.place(0).name))
        self.assertEqual("b", str(net.place(1).name))
        self.assertEqual("c", str(net.place(2).name))

    def test_no_transition_action(self):
        trans = sn._SafeTransition.create(self.conn, name="t")
        self.assertTrue(trans.action is None)

    def test_notify_transition_args(self):
        net = sn.SafeNet.create(self.conn, place_names=[], trans_actions=[],
                place_arcs_out={}, trans_arcs_out={})
        self.assertRaises(TypeError, net.notify_transition)
        self.assertRaises(TypeError, net.notify_transition, trans_idx=0)
        self.assertRaises(TypeError, net.notify_transition, trans_idx=0,
            place_idx=0)

    def test_fire_transition(self):
        places = ["place %d" % i for i in xrange(5)]
        action = sn.CounterAction.create(connection=self.conn, name="counter")
        place_arcs_out = dict((i, [0]) for i in xrange(4))
        trans_arcs_out = {0: [4]}

        net = sn.SafeNet.create(
                connection=self.conn,
                place_names=places,
                trans_actions=[action],
                place_arcs_out=place_arcs_out,
                trans_arcs_out=trans_arcs_out)

        self.assertEqual("place 0", str(net.place(0).name))
        self.assertEqual("place 1", str(net.place(1).name))
        self.assertEqual("place 2", str(net.place(2).name))
        self.assertEqual("place 3", str(net.place(3).name))
        self.assertEqual("place 4", str(net.place(4).name))

        self.assertEqual(['0'], net.place(0).arcs_out.value)
        self.assertEqual(['0'], net.place(1).arcs_out.value)
        self.assertEqual(['0'], net.place(2).arcs_out.value)
        self.assertEqual(['0'], net.place(3).arcs_out.value)
        self.assertEqual([], net.place(4).arcs_out.value)

        self.assertEqual(['0', '1', '2', '3'], net.transition(0).arcs_in.value)
        self.assertEqual(set(['0', '1', '2', '3']), net.transition(0).state.value)
        self.assertEqual(['4'], net.transition(0).arcs_out.value)
        self.assertEqual(action.key, str(net.transition(0).action_key))

        token = sn.Token.create(self.conn)
        net.set_token(0, token.key, self.services)
        self.assertEqual(0, int(action.call_count))
        net.set_token(1, token.key, self.services)
        self.assertEqual(0, int(action.call_count))
        net.set_token(2, token.key, self.services)
        self.assertEqual(0, int(action.call_count))
        net.set_token(3, token.key, self.services)
        self.assertEqual(1, int(action.call_count))

        for i in xrange(0, 3):
            net.notify_transition(0, i, self.services)
            self.assertEqual(1, int(action.call_count))

    def test_place_capacity(self):
        net = sn.SafeNet.create(self.conn, place_names=["place_1"],
                trans_actions=[], place_arcs_out={}, trans_arcs_out={})

        place_idx = 0

        token1 = sn.Token.create(self.conn)
        token2 = sn.Token.create(self.conn)

        self.assertTrue(net.marking(place_idx) is None)
        self.assertEqual({}, net.marking())

        net.set_token(place_idx, token1.key, self.services)
        self.assertEqual(token1.key, net.marking(place_idx))
        net.set_token(place_idx, token1.key, self.services)
        self.assertEqual({"0": str(token1.key)}, net.marking())

        # setting the same token twice is not an error
        self.assertEqual(token1.key, net.marking(place_idx))
        self.assertEqual({"0": str(token1.key)}, net.marking())

        self.assertRaises(sn.PlaceCapacityError, net.set_token,
            place_idx, token2.key, self.services)

    def test_graph(self):
        net = sn.SafeNet.create(self.conn, place_names=["p1", "p2"],
                trans_actions=[None], place_arcs_out={0: set([0])},
                trans_arcs_out={0: set([1])})

        graph = net.graph()
        nodes = graph.nodes()
        edges = graph.edges()
        self.assertEqual(3, len(nodes))
        self.assertEqual(2, len(edges))

class TestTransitionActions(TestBase):
    def test_argument_encoding(self):
        arguments = {
            "integer": 7,
            "float": 1.23,
            "string": "hello",
            "list": [1, 2, "three"],
            "hash": {"x": "y"},
        }
        action = sn.CounterAction.create(connection=self.conn, args=arguments)

        self.assertEqual(action.args.value, arguments)

    def test_shell_command_action(self):
        success_place_id = 1
        failure_place_id = 2

        good_cmdline = [sys.executable, "-c", "import sys; sys.exit(0)"]
        fail_cmdline = [sys.executable, "-c", "import sys; sys.exit(1)"]

        action = sn.ShellCommandAction.create(
            connection=self.conn,
            name="TestAction",
            args={"command_line": good_cmdline},
            place_refs=[success_place_id, failure_place_id],
            )

        self.assertEqual(good_cmdline, action.args["command_line"])

        orchestrator = mock.MagicMock()
        services = {"orchestrator": orchestrator}
        net = mock.MagicMock()
        net.key = "netkey!"

        input_data = {}
        action.execute(input_data, net, services)
        orchestrator.set_token.assert_called_with(
                net.key, success_place_id, token_key=mock.ANY)

        action.args["command_line"] = fail_cmdline
        orchestrator.reset_mock()

        action.execute(input_data, net, services)
        orchestrator.set_token.assert_called_with(
                    net.key, failure_place_id, token_key=mock.ANY)

if __name__ == "__main__":
    unittest.main()
<commit_msg>whitespace<commit_after>import flow.petri.safenet as sn

from test_helpers import RedisTest, FakeOrchestrator
import mock
import os
import redis
import sys
import unittest

class TestBase(RedisTest):
    def setUp(self):
        RedisTest.setUp(self)
        orch = FakeOrchestrator(self.conn)
        self.services = orch.services

class TestTransition(TestBase):
    def test_default_token_merger(self):
        tokens = []
        expected = {}
        for i in xrange(10):
            expected[str(i)] = i
            tokens.append(sn.Token.create(self.conn, data={i: i}))

        token_keys = [x.key for x in tokens]
        transition = sn._SafeTransition.create(self.conn,
                active_tokens=token_keys)
        self.assertEqual(expected, transition.input_data)

    def test_abstract_base(self):
        merger = sn.TokenMerger.create(self.conn)
        self.assertRaises(NotImplementedError, merger.merge, [])

    def test_custom_token_merger(self):
        class PrependingTokenMerger(sn.TokenMerger):
            def merge(self, tokens):
                data = {}
                for t in tokens:
                    data.update({"x" + k: v for k, v in t.data.iteritems()})
                return data

        tokens = []
        expected = {}
        for i in xrange(10):
            expected["x" + str(i)] = i
            tokens.append(sn.Token.create(self.conn, data={i: i}))

        token_keys = [x.key for x in tokens]
        merger = PrependingTokenMerger.create(self.conn)
        transition = sn._SafeTransition.create(self.conn,
                active_tokens=token_keys,
                token_merger_key=merger.key)
        self.assertEqual(expected, transition.input_data)

class TestSafeNet(TestBase):
    def test_no_connection(self):
        self.assertRaises(TypeError, sn.SafeNet.create, None)

    def test_abstract_transition_action(self):
        act = sn.TransitionAction.create(self.conn, name="boom")
        self.assertRaises(NotImplementedError, act.execute, input_data={},
                net=None, services=None)

    def test_constants(self):
        net = sn.SafeNet.create(connection=self.conn)
        env = {"PATH": "/bin:/usr/bin", "USER": "flow"}
        net.set_constant("environment", env)
        self.assertEqual(env, net.constant("environment"))
        self.assertRaises(TypeError, net.set_constant, "environment", 10)

        user_id = 1234
        no_user_id = net.constant("user_id")
        self.assertIsNone(no_user_id)
        net.set_constant("user_id", user_id)
        self.assertRaises(TypeError, net.set_constant, "user_id", 10)
        self.assertEqual(1234, net.constant("user_id"))

    def test_variables(self):
        net = sn.SafeNet.create(connection=self.conn)

        foo = net.variable("foo")
        self.assertIsNone(foo)
        net.set_variable("foo", 123)
        self.assertEqual(123, net.variable("foo"))

        bar = net.variable("bar")
        self.assertIsNone(bar)
        net.set_variable("bar", {"x": "y"})
        self.assertEqual({"x": "y"}, net.variable("bar"))

        baz = net.variable("baz")
        self.assertIsNone(baz)
        net.set_variable("bar", [1, 2, 3])
        self.assertEqual([1, 2, 3], net.variable("bar"))

    def test_places(self):
        action = sn.CounterAction.create(connection=self.conn, name="counter")
        net = sn.SafeNet.create(
                connection=self.conn,
                place_names=["a", "b", "c"],
                trans_actions = [action],
                place_arcs_out={0: [0]}, # a -> t1
                trans_arcs_out={0: [1, 2]}, # t1 -> b, c
                )

        self.assertEqual(3, net.num_places)
        self.assertEqual("a", str(net.place(0).name))
        self.assertEqual("b", str(net.place(1).name))
        self.assertEqual("c", str(net.place(2).name))

    def test_no_transition_action(self):
        trans = sn._SafeTransition.create(self.conn, name="t")
        self.assertTrue(trans.action is None)

    def test_notify_transition_args(self):
        net = sn.SafeNet.create(self.conn, place_names=[], trans_actions=[],
                place_arcs_out={}, trans_arcs_out={})
        self.assertRaises(TypeError, net.notify_transition)
        self.assertRaises(TypeError, net.notify_transition, trans_idx=0)
        self.assertRaises(TypeError, net.notify_transition, trans_idx=0,
            place_idx=0)

    def test_fire_transition(self):
        places = ["place %d" % i for i in xrange(5)]
        action = sn.CounterAction.create(connection=self.conn, name="counter")
        place_arcs_out = dict((i, [0]) for i in xrange(4))
        trans_arcs_out = {0: [4]}

        net = sn.SafeNet.create(
                connection=self.conn,
                place_names=places,
                trans_actions=[action],
                place_arcs_out=place_arcs_out,
                trans_arcs_out=trans_arcs_out)

        self.assertEqual("place 0", str(net.place(0).name))
        self.assertEqual("place 1", str(net.place(1).name))
        self.assertEqual("place 2", str(net.place(2).name))
        self.assertEqual("place 3", str(net.place(3).name))
        self.assertEqual("place 4", str(net.place(4).name))

        self.assertEqual(['0'], net.place(0).arcs_out.value)
        self.assertEqual(['0'], net.place(1).arcs_out.value)
        self.assertEqual(['0'], net.place(2).arcs_out.value)
        self.assertEqual(['0'], net.place(3).arcs_out.value)
        self.assertEqual([], net.place(4).arcs_out.value)

        self.assertEqual(['0', '1', '2', '3'], net.transition(0).arcs_in.value)
        self.assertEqual(set(['0', '1', '2', '3']), net.transition(0).state.value)
        self.assertEqual(['4'], net.transition(0).arcs_out.value)
        self.assertEqual(action.key, str(net.transition(0).action_key))

        token = sn.Token.create(self.conn)
        net.set_token(0, token.key, self.services)
        self.assertEqual(0, int(action.call_count))
        net.set_token(1, token.key, self.services)
        self.assertEqual(0, int(action.call_count))
        net.set_token(2, token.key, self.services)
        self.assertEqual(0, int(action.call_count))
        net.set_token(3, token.key, self.services)
        self.assertEqual(1, int(action.call_count))

        for i in xrange(0, 3):
            net.notify_transition(0, i, self.services)
            self.assertEqual(1, int(action.call_count))

    def test_place_capacity(self):
        net = sn.SafeNet.create(self.conn, place_names=["place_1"],
                trans_actions=[], place_arcs_out={}, trans_arcs_out={})

        place_idx = 0

        token1 = sn.Token.create(self.conn)
        token2 = sn.Token.create(self.conn)

        self.assertTrue(net.marking(place_idx) is None)
        self.assertEqual({}, net.marking())

        net.set_token(place_idx, token1.key, self.services)
        self.assertEqual(token1.key, net.marking(place_idx))
        net.set_token(place_idx, token1.key, self.services)
        self.assertEqual({"0": str(token1.key)}, net.marking())

        # setting the same token twice is not an error
        self.assertEqual(token1.key, net.marking(place_idx))
        self.assertEqual({"0": str(token1.key)}, net.marking())

        self.assertRaises(sn.PlaceCapacityError, net.set_token,
            place_idx, token2.key, self.services)

    def test_graph(self):
        net = sn.SafeNet.create(self.conn, place_names=["p1", "p2"],
                trans_actions=[None], place_arcs_out={0: set([0])},
                trans_arcs_out={0: set([1])})

        graph = net.graph()
        nodes = graph.nodes()
        edges = graph.edges()
        self.assertEqual(3, len(nodes))
        self.assertEqual(2, len(edges))

class TestTransitionActions(TestBase):
    def test_argument_encoding(self):
        arguments = {
            "integer": 7,
            "float": 1.23,
            "string": "hello",
            "list": [1, 2, "three"],
            "hash": {"x": "y"},
        }
        action = sn.CounterAction.create(connection=self.conn, args=arguments)

        self.assertEqual(action.args.value, arguments)

    def test_shell_command_action(self):
        success_place_id = 1
        failure_place_id = 2

        good_cmdline = [sys.executable, "-c", "import sys; sys.exit(0)"]
        fail_cmdline = [sys.executable, "-c", "import sys; sys.exit(1)"]

        action = sn.ShellCommandAction.create(
            connection=self.conn,
            name="TestAction",
            args={"command_line": good_cmdline},
            place_refs=[success_place_id, failure_place_id],
            )

        self.assertEqual(good_cmdline, action.args["command_line"])

        orchestrator = mock.MagicMock()
        services = {"orchestrator": orchestrator}
        net = mock.MagicMock()
        net.key = "netkey!"

        input_data = {}
        action.execute(input_data, net, services)
        orchestrator.set_token.assert_called_with(
                net.key, success_place_id, token_key=mock.ANY)

        action.args["command_line"] = fail_cmdline
        orchestrator.reset_mock()

        action.execute(input_data, net, services)
        orchestrator.set_token.assert_called_with(
                    net.key, failure_place_id, token_key=mock.ANY)

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


<commit_before>"""
Implements base classes to derive plot classes from. The code in py:mod:`cea.plots.categories` uses
py:class:`cea.plots.base.PlotBase` to figure out the list of plots in a category.
"""
from __future__ import division
from __future__ import print_function
import os
import re
import jinja2
import plotly.graph_objs
import plotly.offline
import cea.inputlocator
from cea import MissingInputDataException
from cea.plots.variable_naming import LOGO, COLOR, NAMING

__author__ = "Daren Thomas"
__copyright__ = "Copyright 2018, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Daren Thomas"]
__license__ = "MIT"
__version__ = "0.1"
__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

class PlotBase(object):
    """A base class for plots containing helper methods used by all plots."""

    # override these in plot subclasses!
    name = None  # a label to name the plot
    category_name = None  # name of the category this plot belongs to (can be inherited from category base plot)
    category_path = None  # a relative path for outputting the plot to (FIXME: maybe we remove this later on)
    expected_parameters = {}  # maps parameter-name -> "section:parameter"

    @classmethod
    def id(cls):
        name = re.sub('\s+\(.*\)', '', cls.name)  # remove parenthesis
        return name.lower().replace(' ', '_').replace('/', '_')  # use for js/html etc.

    def __init__(self, project, parameters, cache):
        self.cache = cache  # a PlotCache implementation for reading cached data
        self.project = project  # full path to the project this plot belongs to
        self.category_path = None  # override this in the __init__.py subclasses for each category (see cea/plots/demand/__init__.py for an example)
        # self.analysis_fields = None  # override this in the plot subclasses! set it to a list of fields in self.data
        # self.input_files = []  # override this in the plot subclasses! set it to a list of tuples (locator.method, args)
        self.parameters = parameters
        self.buildings = self.process_buildings_parameter() if 'buildings' in self.expected_parameters else None

        for parameter_name in self.expected_parameters:
            assert parameter_name in parameters, "Missing parameter {}".format(parameter_name)

    def missing_input_files(self):
        """Return the list of missing input files for this plot"""
        result = []
        for locator_method, args in self.input_files:
            if not os.path.exists(locator_method(*args)):
                result.append((locator_method, args))
        return result

    @property
    def locator(self):
        """
        :return: cea.inputlocator.InputLocator
        """
        return cea.inputlocator.InputLocator(os.path.join(self.project, self.parameters['scenario-name']))

    @property
    def layout(self):
        # override this in the plot subclasses! set it to a plotly.graph_objs.Layout object
        return None

    @property
    def title(self):
        """Override the version in PlotBase"""
        if set(self.buildings) != set(self.locator.get_zone_building_names()):
            if len(self.buildings) == 1:
                return "%s for Building %s" % (self.name, self.buildings[0])
            else:
                return "%s for Selected Buildings" % self.name
        return "%s for District" % self.name

    def totals_bar_plot(self):
        """Creates a plot based on the totals data in percentages."""
        traces = []
        data = self.data
        data['total'] = data[self.analysis_fields].sum(axis=1)
        data = data.sort_values(by='total', ascending=False)  # this will get the maximum value to the left
        for field in self.analysis_fields:
            y = data[field]
            total_percent = (y / data['total'] * 100).round(2).values
            total_percent_txt = ["(%.2f %%)" % x for x in total_percent]
            name = NAMING[field]
            trace = plotly.graph_objs.Bar(x=data["Name"], y=y, name=name, marker=dict(color=COLOR[field]))
            traces.append(trace)
        return traces

    @property
    def output_path(self):
        """The output path to use for the solar-potential plots"""
        assert self.name, "Attribute 'name' not defined for this plot (%s)" % self.__class__
        assert self.category_path, "Attribute 'category_path' not defined for this plot(%s)" % self.__class__

        if len(self.buildings) == 1:
            prefix = 'Building_%s' % self.buildings[0]
        elif len(self.buildings) < len(self.locator.get_zone_building_names()):
            prefix = 'Selected_Buildings'
        else:
            prefix = 'District'
        file_name = "%s_%s" % (prefix, self.id())
        return self.locator.get_timeseries_plots_file(file_name, self.category_path)

    def remove_unused_fields(self, data, fields):
        """
        Helper method that, given a data frame and a list of fields in that data frame, returns the subset of those
        fields that actually have data.

        FIXME: what about columns with negative values?
        """
        return [field for field in fields if data[field].sum() > 0.0]

    def calc_graph(self):
        """Calculate a plotly Data object as to be passed to the data attribute of Figure"""
        raise NotImplementedError('Subclass needs to implement calc_graph for plot!')

    def calc_table(self):
        """Calculates a pandas.Dataframe to display as table."""
        raise NotImplementedError('This plot has no table')

    def plot(self, auto_open=False):
        """Plots the graphs to the filename (see output_path)"""
        if self.missing_input_files():
            raise MissingInputDataException("Dear developer: Run check_input_files() first, before plotting!")
        # PLOT
        template_path = os.path.join(os.path.dirname(__file__), 'plot.html')
        template = jinja2.Template(open(template_path, 'r').read())
        plot_html = template.render(plot_div=self.plot_div(), table_div=self.table_div(), title=self.title)
        with open(self.output_path, 'w') as f:
            f.write(plot_html)

        print("Plotted '%s' to %s" % (self.name, self.output_path))
        if auto_open:
            import webbrowser
            webbrowser.open(self.output_path)

    def plot_div(self):
        """Return the plot as an html <div/> for use in the dashboard. Override this method in subclasses"""
        if self.missing_input_files():
            raise MissingInputDataException("Dear developer: Run check_input_files() first, before plotting!")
        return self.cache.lookup_plot_div(self, self._plot_div_producer)

    def _plot_div_producer(self):
        fig = plotly.graph_objs.Figure(data=self.calc_graph(), layout=self.layout)
        fig['layout'] = dict(fig['layout'],**{'margin':dict(l=50, r=50, t=20, b=50),'hovermode':'closest'})
        fig['layout']['yaxis'] = dict(fig['layout']['yaxis'], **{'hoverformat': ".2f"})
        div = plotly.offline.plot(fig, output_type='div', include_plotlyjs=False, show_link=False)
        return div

    def table_div(self):
        """Returns the html div for a table, or an empty string if no table is to be produced"""
        if self.missing_input_files():
            raise MissingInputDataException("Dear developer: Run check_input_files() first, before plotting!")
        return self.cache.lookup_table_div(self, self._table_div_producer)

    def _table_div_producer(self):
        """Default producer for table divs (override if you need more control)"""
        try:
            table_df = self.calc_table()
            template_path = os.path.join(os.path.dirname(__file__), 'table.html')
            template = jinja2.Template(open(template_path, 'r').read())
            table_html = template.render(table_df=table_df)
            return table_html
        except NotImplementedError:
            return ''

    @classmethod
    def get_default_parameters(cls, config):
        """Return a dictionary of parameters taken by using the values in the config file"""
        return {
            k: config.get(v)
            for k, v in cls.expected_parameters.items()
        }

    def process_buildings_parameter(self):
        """
        Make sure the buildings parameter contains only buildings in the zone. Returns (and updates) the parameter.
        """
        # all plots in this category use the buildings parameter. make it easier to access
        # handle special case of buildings... (only allow buildings for the scenario in question)
        zone_building_names = self.locator.get_zone_building_names()

        if not self.parameters['buildings']:
            self.parameters['buildings'] = zone_building_names
        self.parameters['buildings'] = ([b for b in self.parameters['buildings'] if
                                         b in zone_building_names]
                                        or zone_building_names)
        return self.parameters['buildings']
<commit_msg>Allow loading of default values for missing plot parameters<commit_after>"""
Implements base classes to derive plot classes from. The code in py:mod:`cea.plots.categories` uses
py:class:`cea.plots.base.PlotBase` to figure out the list of plots in a category.
"""
from __future__ import division
from __future__ import print_function
import os
import re
import jinja2
import plotly.graph_objs
import plotly.offline
import cea.inputlocator
import cea.config
from cea import MissingInputDataException
from cea.plots.variable_naming import LOGO, COLOR, NAMING

__author__ = "Daren Thomas"
__copyright__ = "Copyright 2018, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Daren Thomas"]
__license__ = "MIT"
__version__ = "0.1"
__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

class PlotBase(object):
    """A base class for plots containing helper methods used by all plots."""

    # override these in plot subclasses!
    name = None  # a label to name the plot
    category_name = None  # name of the category this plot belongs to (can be inherited from category base plot)
    category_path = None  # a relative path for outputting the plot to (FIXME: maybe we remove this later on)
    expected_parameters = {}  # maps parameter-name -> "section:parameter"

    @classmethod
    def id(cls):
        name = re.sub('\s+\(.*\)', '', cls.name)  # remove parenthesis
        return name.lower().replace(' ', '_').replace('/', '_')  # use for js/html etc.

    def __init__(self, project, parameters, cache):
        self.cache = cache  # a PlotCache implementation for reading cached data
        self.project = project  # full path to the project this plot belongs to
        self.category_path = None  # override this in the __init__.py subclasses for each category (see cea/plots/demand/__init__.py for an example)
        # self.analysis_fields = None  # override this in the plot subclasses! set it to a list of fields in self.data
        # self.input_files = []  # override this in the plot subclasses! set it to a list of tuples (locator.method, args)
        self.parameters = parameters
        self.buildings = self.process_buildings_parameter() if 'buildings' in self.expected_parameters else None

        for parameter_name in self.expected_parameters:
            # Try to load missing parameters with default values
            if parameter_name not in parameters:
                try:
                    self.parameters[parameter_name] = cea.config.Configuration(cea.config.DEFAULT_CONFIG).get('plots:{}'.format(parameter_name))
                except Exception:
                    assert parameter_name in parameters, "Missing parameter {}".format(parameter_name)

    def missing_input_files(self):
        """Return the list of missing input files for this plot"""
        result = []
        for locator_method, args in self.input_files:
            if not os.path.exists(locator_method(*args)):
                result.append((locator_method, args))
        return result

    @property
    def locator(self):
        """
        :return: cea.inputlocator.InputLocator
        """
        return cea.inputlocator.InputLocator(os.path.join(self.project, self.parameters['scenario-name']))

    @property
    def layout(self):
        # override this in the plot subclasses! set it to a plotly.graph_objs.Layout object
        return None

    @property
    def title(self):
        """Override the version in PlotBase"""
        if set(self.buildings) != set(self.locator.get_zone_building_names()):
            if len(self.buildings) == 1:
                return "%s for Building %s" % (self.name, self.buildings[0])
            else:
                return "%s for Selected Buildings" % self.name
        return "%s for District" % self.name

    def totals_bar_plot(self):
        """Creates a plot based on the totals data in percentages."""
        traces = []
        data = self.data
        data['total'] = data[self.analysis_fields].sum(axis=1)
        data = data.sort_values(by='total', ascending=False)  # this will get the maximum value to the left
        for field in self.analysis_fields:
            y = data[field]
            total_percent = (y / data['total'] * 100).round(2).values
            total_percent_txt = ["(%.2f %%)" % x for x in total_percent]
            name = NAMING[field]
            trace = plotly.graph_objs.Bar(x=data["Name"], y=y, name=name, marker=dict(color=COLOR[field]))
            traces.append(trace)
        return traces

    @property
    def output_path(self):
        """The output path to use for the solar-potential plots"""
        assert self.name, "Attribute 'name' not defined for this plot (%s)" % self.__class__
        assert self.category_path, "Attribute 'category_path' not defined for this plot(%s)" % self.__class__

        if len(self.buildings) == 1:
            prefix = 'Building_%s' % self.buildings[0]
        elif len(self.buildings) < len(self.locator.get_zone_building_names()):
            prefix = 'Selected_Buildings'
        else:
            prefix = 'District'
        file_name = "%s_%s" % (prefix, self.id())
        return self.locator.get_timeseries_plots_file(file_name, self.category_path)

    def remove_unused_fields(self, data, fields):
        """
        Helper method that, given a data frame and a list of fields in that data frame, returns the subset of those
        fields that actually have data.

        FIXME: what about columns with negative values?
        """
        return [field for field in fields if data[field].sum() > 0.0]

    def calc_graph(self):
        """Calculate a plotly Data object as to be passed to the data attribute of Figure"""
        raise NotImplementedError('Subclass needs to implement calc_graph for plot!')

    def calc_table(self):
        """Calculates a pandas.Dataframe to display as table."""
        raise NotImplementedError('This plot has no table')

    def plot(self, auto_open=False):
        """Plots the graphs to the filename (see output_path)"""
        if self.missing_input_files():
            raise MissingInputDataException("Dear developer: Run check_input_files() first, before plotting!")
        # PLOT
        template_path = os.path.join(os.path.dirname(__file__), 'plot.html')
        template = jinja2.Template(open(template_path, 'r').read())
        plot_html = template.render(plot_div=self.plot_div(), table_div=self.table_div(), title=self.title)
        with open(self.output_path, 'w') as f:
            f.write(plot_html)

        print("Plotted '%s' to %s" % (self.name, self.output_path))
        if auto_open:
            import webbrowser
            webbrowser.open(self.output_path)

    def plot_div(self):
        """Return the plot as an html <div/> for use in the dashboard. Override this method in subclasses"""
        if self.missing_input_files():
            raise MissingInputDataException("Dear developer: Run check_input_files() first, before plotting!")
        return self.cache.lookup_plot_div(self, self._plot_div_producer)

    def _plot_div_producer(self):
        fig = plotly.graph_objs.Figure(data=self.calc_graph(), layout=self.layout)
        fig['layout'] = dict(fig['layout'],**{'margin':dict(l=50, r=50, t=20, b=50),'hovermode':'closest'})
        fig['layout']['yaxis'] = dict(fig['layout']['yaxis'], **{'hoverformat': ".2f"})
        div = plotly.offline.plot(fig, output_type='div', include_plotlyjs=False, show_link=False)
        return div

    def table_div(self):
        """Returns the html div for a table, or an empty string if no table is to be produced"""
        if self.missing_input_files():
            raise MissingInputDataException("Dear developer: Run check_input_files() first, before plotting!")
        return self.cache.lookup_table_div(self, self._table_div_producer)

    def _table_div_producer(self):
        """Default producer for table divs (override if you need more control)"""
        try:
            table_df = self.calc_table()
            template_path = os.path.join(os.path.dirname(__file__), 'table.html')
            template = jinja2.Template(open(template_path, 'r').read())
            table_html = template.render(table_df=table_df)
            return table_html
        except NotImplementedError:
            return ''

    @classmethod
    def get_default_parameters(cls, config):
        """Return a dictionary of parameters taken by using the values in the config file"""
        return {
            k: config.get(v)
            for k, v in cls.expected_parameters.items()
        }

    def process_buildings_parameter(self):
        """
        Make sure the buildings parameter contains only buildings in the zone. Returns (and updates) the parameter.
        """
        # all plots in this category use the buildings parameter. make it easier to access
        # handle special case of buildings... (only allow buildings for the scenario in question)
        zone_building_names = self.locator.get_zone_building_names()

        if not self.parameters['buildings']:
            self.parameters['buildings'] = zone_building_names
        self.parameters['buildings'] = ([b for b in self.parameters['buildings'] if
                                         b in zone_building_names]
                                        or zone_building_names)
        return self.parameters['buildings']
<|endoftext|>


<commit_before>"""
"Makefile" for doit to test the whole package with Jenkins.

This file gets run with the ``cea test`` command as well as by the Jenkins continuous integration server. It runs all
the unit tests in the ``cea/tests/`` folder as well as some of the CEA scripts, to make sure they at least run through.

In order to run reference cases besides the one called "open", you will need to set up authentication to the private
GitHub repository. The easiest way to do this is with ``cea test --save --user USERNAME --token PERSONAL_ACCESS_TOKEN``,
adding your own user name and a GitHub personal access token.

The reference cases can be found here: https://github.com/architecture-building-systems/cea-reference-case/archive/master.zip
"""
import os
import shutil
import sys
import zipfile
import tempfile
import cea.inputlocator
import cea.config

import requests

__author__ = "Daren Thomas"
__copyright__ = "Copyright 2016, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Daren Thomas"]
__license__ = "MIT"
__version__ = "0.1"

__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

REPOSITORY_URL = "https://github.com/architecture-building-systems/cea-reference-case/archive/%s.zip"
REPOSITORY_NAME = "master"

if 'JOB_NAME' in os.environ:
    # this script is being run as part of a Jenkins job
    ARCHIVE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), '%JOB_NAME%/cea-reference-case.zip'))
    REFERENCE_CASE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), '%JOB_NAME%/cea-reference-case'))
else:
    ARCHIVE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), 'cea-reference-case.zip'))
    REFERENCE_CASE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), 'cea-reference-case'))

REFERENCE_CASES = {
    'open': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME, "reference-case-open",
                         "baseline"),
    'zug/baseline': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME, "reference-case-zug",
                                 "baseline"),
    'zurich/baseline': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME,
                                    "reference-case-zurich",
                                    "baseline"),
    'zurich/masterplan': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME,
                                      "reference-case-zurich",
                                      "masterplan")}

REFERENCE_CASES_DATA = {
    'open': {'weather': 'Zug', 'latitude': 47.1628017306431, 'longitude': 8.31,
             'radiation': 'open.baseline.radiation.csv',
             'properties_surfaces': 'open.baseline.properties_surfaces.csv'},
    'zug/baseline': {'weather': 'Zug', 'latitude': 47.1628017306431, 'longitude': 8.31,
                     'radiation': 'zug.baseline.radiation.csv',
                     'properties_surfaces': 'zug.baseline.properties_surfaces.csv'},
    'zurich/baseline': {'weather': 'Zurich', 'latitude': 46.9524055556, 'longitude': 7.43958333333,
                        'radiation': 'hq.baseline.radiation.csv',
                        'properties_surfaces': 'hq.baseline.properties_surfaces.csv'},
    'zurich/masterplan': {'weather': 'Zurich', 'latitude': 46.9524055556, 'longitude': 7.43958333333,
                          'radiation': 'hq.masterplan.radiation.csv',
                          'properties_surfaces': 'hq.masterplan.properties_surfaces.csv'}}

# set to github user and personal access token in main
_user = None
_token = None
_reference_cases = []

def get_github_auth():
    """
    get the username / token for github from a file in the home directory
    called "cea_github.auth". The first line contains the user, the second line
    the personal access token.

    :return: (user, token)
    """
    if _user and _token:
        return _user, _token

    with open(os.path.expanduser('~/cea_github.auth')) as f:
        user, token = map(str.strip, f.readlines())
    return user, token

def task_run_unit_tests():
    """run the unittests"""
    def run_unit_tests():
        import unittest
        import os
        testsuite = unittest.defaultTestLoader.discover(os.path.dirname(__file__))
        result = unittest.TextTestRunner(verbosity=1).run(testsuite)
        return result.wasSuccessful()
    return {
        'actions': [run_unit_tests],
        'task_dep': [],
        'verbosity': 1
    }

def task_download_reference_cases():
    """Download the (current) state of the reference-case-zug"""
    def download_reference_cases():
        if os.path.exists(REFERENCE_CASE_PATH):
            shutil.rmtree(REFERENCE_CASE_PATH)
        # extract the bundled reference case (we will use this anyways
        import cea.examples
        archive = zipfile.ZipFile(os.path.join(os.path.dirname(cea.examples.__file__), 'reference-case-open.zip'))
        archive.extractall(os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME))

        if len([rc for rc in _reference_cases if rc.lower() != 'open']):
            if os.path.exists(ARCHIVE_PATH):
                os.remove(ARCHIVE_PATH)
            r = requests.get(REPOSITORY_URL % REPOSITORY_NAME, auth=get_github_auth())
            assert r.ok, 'could not download the reference case'
            with open(ARCHIVE_PATH, 'wb') as f:
                f.write(r.content)
            # extract the reference cases to the temp folder
            archive = zipfile.ZipFile(ARCHIVE_PATH)
            archive.extractall(REFERENCE_CASE_PATH)

    return {
        'actions': [download_reference_cases],
        'verbosity': 1,
    }

def task_run_data_helper():
    """Run the data helper for each reference case"""
    import cea.demand.preprocessing.data_helper
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': reference_case,
            'actions': [
                (cea.demand.preprocessing.data_helper.main, [], {
                    'config': config})],
            'verbosity': 1,
        }

def task_download_radiation():
    """For some reference cases, the radiation and properties_surfaces.csv files are too big for git and are stored
    with git lfs... For these cases we download a known good version from a url"""
    def download_radiation(scenario_path, reference_case):
        locator = cea.inputlocator.InputLocator(scenario_path)
        data = REFERENCE_CASES_DATA[reference_case]
        properties_surfaces_csv = os.path.join(os.path.dirname(__file__), 'radiation_data', data['properties_surfaces'])
        radiation_csv = os.path.join(os.path.dirname(__file__), 'radiation_data', data['radiation'])
        shutil.copyfile(properties_surfaces_csv, locator.get_surface_properties())
        shutil.copyfile(radiation_csv, locator.get_radiation())

    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        yield {
            'name': reference_case,
            'actions': [(download_radiation, [], {
                'scenario_path': scenario_path,
                'reference_case': reference_case})]
        }

def task_run_demand():
    """run the demand script for each reference cases and weather file"""
    import cea.demand.demand_main
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        locator = cea.inputlocator.InputLocator(scenario_path)
        weather = REFERENCE_CASES_DATA[reference_case]['weather']

        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        config.weather = weather
        config.demand.use_daysim_radiation = False

        yield {
            'name': '%(reference_case)s@%(weather)s' % locals(),
            'actions': [(cea.demand.demand_main.main, [], {
                'config': config,
            })],
            'verbosity': 1,
        }

def task_run_embodied_energy():
    """Run the embodied energy script for each reference case"""
    import cea.analysis.lca.embodied
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        config.embodied_energy.year_to_calculate = 2050
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.analysis.lca.embodied.main, [], {
                'config': config,
            })],
            'verbosity': 1,
        }

def task_run_emissions_operation():
    """run the emissions operation script for each reference case"""
    import cea.analysis.lca.operation
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.analysis.lca.operation.main, [], {
                'config': config
            })],
            'verbosity': 1,
        }

def task_run_emissions_mobility():
    """run the emissions mobility script for each reference case"""
    import cea.analysis.lca.mobility
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.analysis.lca.mobility.main, [], {
                'config': config
            })],
            'verbosity': 1,
        }

def task_run_heatmaps():
    """run the heat maps script for each reference case"""
    try:
        from cea.interfaces.arcgis.modules import arcpy
    except ImportError:
        # do not require ArcGIS to be installed, but skip testing the heatmaps
        # module if it isn't installed.
        return
    import cea.plots.heatmaps
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.plots.heatmaps.main, [], {
                'config': config
            })],
            'verbosity': 1,
        }

def task_run_scenario_plots():
    """run the scenario plots script for each reference case"""
    import cea.plots.scenario_plots
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        project = os.path.normpath(os.path.join(scenario_path, '..'))
        output_file = os.path.join(project, 'scenarios.pdf')

        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario_plots.project = project
        config.scenario_plots.scenarios = ['baseline']
        config.scenario_plots.output_file = output_file
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.plots.scenario_plots.main, [], {
                'config': config,
            })],
            'verbosity': 1,
        }

def task_run_calibration():
    """run the calibration_sampling for each reference case"""
    def run_calibration():
        import cea.demand.calibration.bayesian_calibrator.calibration_sampling as calibration_sampling
        import cea.demand.calibration.bayesian_calibrator.calibration_gaussian_emulator as calibration_gaussian_emulator
        import cea.demand.calibration.bayesian_calibrator.calibration_main as calibration_main

        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        locator = cea.inputlocator.ReferenceCaseOpenLocator()

        config.scenario = locator.scenario
        config.single_calibration.building = 'B01'
        config.single_calibration.variables = ['U_win', 'U_wall', 'U_roof', 'n50', 'Tcs_set_C', 'Hs']
        config.single_calibration.load = 'Qcsf'
        config.single_calibration.samples = 10
        config.single_calibration.show_plots = False

        # run calibration_sampling
        calibration_sampling.sampling_main(locator=locator, config=config)

        # run calibration_gaussian_emulator
        calibration_gaussian_emulator.gaussian_emulator(locator=locator, config=config)

        # run calibration_main
        calibration_main.calibration_main(locator=locator, config=config)

        # make sure the files were created
        # FIXME: @JIMENOFONSECA - what files do I need to check? (this has changed)

    return {
        'actions': [(run_calibration, [], {})],
        'verbosity': 1,
    }

def main(config):
    from doit.api import DoitMain
    from doit.api import ModuleTaskLoader

    global _reference_cases
    if not config.test.reference_cases:
        _reference_cases = ['open']
    elif 'all' in config.test.reference_cases:
        _reference_cases = REFERENCE_CASES.keys()
    else:
        _reference_cases = [rc for rc in config.test.reference_cases if rc in REFERENCE_CASES.keys()]
    sys.exit(DoitMain(ModuleTaskLoader(globals())).run([]))

if __name__ == '__main__':
    main(cea.config.Configuration())
<commit_msg>fixed documentation in dodo script, trigger jenkins<commit_after>"""
"Makefile" for doit to test the whole package with Jenkins.

This file gets run with the ``cea test`` command as well as by the Jenkins continuous integration server. It runs all
the unit tests in the ``cea/tests/`` folder as well as some of the CEA scripts, to make sure they at least run through.

In order to run reference cases besides the one called "open", you will need to set up authentication to the private
GitHub repository. To do this, you need to create a GitHub [authentication token](https://help.github.com/articles/creating-a-personal-access-token-for-the-command-line/).
Then, create a text file called ``cea_github.auth`` in your home directory (e.g. ``C:\Users\your-user-name`` for Windows systems or equivalent on
POSIX systems, ask your administrator if you don't know what this is). The file should contain two lines, the first being
your GitHub user name the second the authentication token.

The reference cases can be found here: https://github.com/architecture-building-systems/cea-reference-case/archive/master.zip
"""
import os
import shutil
import sys
import zipfile
import tempfile
import cea.inputlocator
import cea.config

import requests

__author__ = "Daren Thomas"
__copyright__ = "Copyright 2016, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Daren Thomas"]
__license__ = "MIT"
__version__ = "0.1"

__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

REPOSITORY_URL = "https://github.com/architecture-building-systems/cea-reference-case/archive/%s.zip"
REPOSITORY_NAME = "master"

if 'JOB_NAME' in os.environ:
    # this script is being run as part of a Jenkins job
    ARCHIVE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), '%JOB_NAME%/cea-reference-case.zip'))
    REFERENCE_CASE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), '%JOB_NAME%/cea-reference-case'))
else:
    ARCHIVE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), 'cea-reference-case.zip'))
    REFERENCE_CASE_PATH = os.path.expandvars(os.path.join(tempfile.gettempdir(), 'cea-reference-case'))

REFERENCE_CASES = {
    'open': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME, "reference-case-open",
                         "baseline"),
    'zug/baseline': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME, "reference-case-zug",
                                 "baseline"),
    'zurich/baseline': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME,
                                    "reference-case-zurich",
                                    "baseline"),
    'zurich/masterplan': os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME,
                                      "reference-case-zurich",
                                      "masterplan")}

REFERENCE_CASES_DATA = {
    'open': {'weather': 'Zug', 'latitude': 47.1628017306431, 'longitude': 8.31,
             'radiation': 'open.baseline.radiation.csv',
             'properties_surfaces': 'open.baseline.properties_surfaces.csv'},
    'zug/baseline': {'weather': 'Zug', 'latitude': 47.1628017306431, 'longitude': 8.31,
                     'radiation': 'zug.baseline.radiation.csv',
                     'properties_surfaces': 'zug.baseline.properties_surfaces.csv'},
    'zurich/baseline': {'weather': 'Zurich', 'latitude': 46.9524055556, 'longitude': 7.43958333333,
                        'radiation': 'hq.baseline.radiation.csv',
                        'properties_surfaces': 'hq.baseline.properties_surfaces.csv'},
    'zurich/masterplan': {'weather': 'Zurich', 'latitude': 46.9524055556, 'longitude': 7.43958333333,
                          'radiation': 'hq.masterplan.radiation.csv',
                          'properties_surfaces': 'hq.masterplan.properties_surfaces.csv'}}

# set to github user and personal access token in main
_user = None
_token = None
_reference_cases = []

def get_github_auth():
    """
    get the username / token for github from a file in the home directory
    called "cea_github.auth". The first line contains the user, the second line
    the personal access token.

    :return: (user, token)
    """
    if _user and _token:
        return _user, _token

    with open(os.path.expanduser('~/cea_github.auth')) as f:
        user, token = map(str.strip, f.readlines())
    return user, token

def task_run_unit_tests():
    """run the unittests"""
    def run_unit_tests():
        import unittest
        import os
        testsuite = unittest.defaultTestLoader.discover(os.path.dirname(__file__))
        result = unittest.TextTestRunner(verbosity=1).run(testsuite)
        return result.wasSuccessful()
    return {
        'actions': [run_unit_tests],
        'task_dep': [],
        'verbosity': 1
    }

def task_download_reference_cases():
    """Download the (current) state of the reference-case-zug"""
    def download_reference_cases():
        if os.path.exists(REFERENCE_CASE_PATH):
            shutil.rmtree(REFERENCE_CASE_PATH)
        # extract the bundled reference case (we will use this anyways
        import cea.examples
        archive = zipfile.ZipFile(os.path.join(os.path.dirname(cea.examples.__file__), 'reference-case-open.zip'))
        archive.extractall(os.path.join(REFERENCE_CASE_PATH, "cea-reference-case-%s" % REPOSITORY_NAME))

        if len([rc for rc in _reference_cases if rc.lower() != 'open']):
            if os.path.exists(ARCHIVE_PATH):
                os.remove(ARCHIVE_PATH)
            r = requests.get(REPOSITORY_URL % REPOSITORY_NAME, auth=get_github_auth())
            assert r.ok, 'could not download the reference case'
            with open(ARCHIVE_PATH, 'wb') as f:
                f.write(r.content)
            # extract the reference cases to the temp folder
            archive = zipfile.ZipFile(ARCHIVE_PATH)
            archive.extractall(REFERENCE_CASE_PATH)

    return {
        'actions': [download_reference_cases],
        'verbosity': 1,
    }

def task_run_data_helper():
    """Run the data helper for each reference case"""
    import cea.demand.preprocessing.data_helper
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': reference_case,
            'actions': [
                (cea.demand.preprocessing.data_helper.main, [], {
                    'config': config})],
            'verbosity': 1,
        }

def task_download_radiation():
    """For some reference cases, the radiation and properties_surfaces.csv files are too big for git and are stored
    with git lfs... For these cases we download a known good version from a url"""
    def download_radiation(scenario_path, reference_case):
        locator = cea.inputlocator.InputLocator(scenario_path)
        data = REFERENCE_CASES_DATA[reference_case]
        properties_surfaces_csv = os.path.join(os.path.dirname(__file__), 'radiation_data', data['properties_surfaces'])
        radiation_csv = os.path.join(os.path.dirname(__file__), 'radiation_data', data['radiation'])
        shutil.copyfile(properties_surfaces_csv, locator.get_surface_properties())
        shutil.copyfile(radiation_csv, locator.get_radiation())

    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        yield {
            'name': reference_case,
            'actions': [(download_radiation, [], {
                'scenario_path': scenario_path,
                'reference_case': reference_case})]
        }

def task_run_demand():
    """run the demand script for each reference cases and weather file"""
    import cea.demand.demand_main
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        locator = cea.inputlocator.InputLocator(scenario_path)
        weather = REFERENCE_CASES_DATA[reference_case]['weather']

        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        config.weather = weather
        config.demand.use_daysim_radiation = False

        yield {
            'name': '%(reference_case)s@%(weather)s' % locals(),
            'actions': [(cea.demand.demand_main.main, [], {
                'config': config,
            })],
            'verbosity': 1,
        }

def task_run_embodied_energy():
    """Run the embodied energy script for each reference case"""
    import cea.analysis.lca.embodied
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        config.embodied_energy.year_to_calculate = 2050
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.analysis.lca.embodied.main, [], {
                'config': config,
            })],
            'verbosity': 1,
        }

def task_run_emissions_operation():
    """run the emissions operation script for each reference case"""
    import cea.analysis.lca.operation
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.analysis.lca.operation.main, [], {
                'config': config
            })],
            'verbosity': 1,
        }

def task_run_emissions_mobility():
    """run the emissions mobility script for each reference case"""
    import cea.analysis.lca.mobility
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.analysis.lca.mobility.main, [], {
                'config': config
            })],
            'verbosity': 1,
        }

def task_run_heatmaps():
    """run the heat maps script for each reference case"""
    try:
        from cea.interfaces.arcgis.modules import arcpy
    except ImportError:
        # do not require ArcGIS to be installed, but skip testing the heatmaps
        # module if it isn't installed.
        return
    import cea.plots.heatmaps
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario = scenario_path
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.plots.heatmaps.main, [], {
                'config': config
            })],
            'verbosity': 1,
        }

def task_run_scenario_plots():
    """run the scenario plots script for each reference case"""
    import cea.plots.scenario_plots
    for reference_case, scenario_path in REFERENCE_CASES.items():
        if _reference_cases and reference_case not in _reference_cases:
            continue
        project = os.path.normpath(os.path.join(scenario_path, '..'))
        output_file = os.path.join(project, 'scenarios.pdf')

        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        config.scenario_plots.project = project
        config.scenario_plots.scenarios = ['baseline']
        config.scenario_plots.output_file = output_file
        yield {
            'name': '%(reference_case)s' % locals(),
            'actions': [(cea.plots.scenario_plots.main, [], {
                'config': config,
            })],
            'verbosity': 1,
        }

def task_run_calibration():
    """run the calibration_sampling for each reference case"""
    def run_calibration():
        import cea.demand.calibration.bayesian_calibrator.calibration_sampling as calibration_sampling
        import cea.demand.calibration.bayesian_calibrator.calibration_gaussian_emulator as calibration_gaussian_emulator
        import cea.demand.calibration.bayesian_calibrator.calibration_main as calibration_main

        config = cea.config.Configuration(cea.config.DEFAULT_CONFIG)
        locator = cea.inputlocator.ReferenceCaseOpenLocator()

        config.scenario = locator.scenario
        config.single_calibration.building = 'B01'
        config.single_calibration.variables = ['U_win', 'U_wall', 'U_roof', 'n50', 'Tcs_set_C', 'Hs']
        config.single_calibration.load = 'Qcsf'
        config.single_calibration.samples = 10
        config.single_calibration.show_plots = False

        # run calibration_sampling
        calibration_sampling.sampling_main(locator=locator, config=config)

        # run calibration_gaussian_emulator
        calibration_gaussian_emulator.gaussian_emulator(locator=locator, config=config)

        # run calibration_main
        calibration_main.calibration_main(locator=locator, config=config)

        # make sure the files were created
        # FIXME: @JIMENOFONSECA - what files do I need to check? (this has changed)

    return {
        'actions': [(run_calibration, [], {})],
        'verbosity': 1,
    }

def main(config):
    from doit.api import DoitMain
    from doit.api import ModuleTaskLoader

    global _reference_cases
    if not config.test.reference_cases:
        _reference_cases = ['open']
    elif 'all' in config.test.reference_cases:
        _reference_cases = REFERENCE_CASES.keys()
    else:
        _reference_cases = [rc for rc in config.test.reference_cases if rc in REFERENCE_CASES.keys()]
    sys.exit(DoitMain(ModuleTaskLoader(globals())).run([]))

if __name__ == '__main__':
    main(cea.config.Configuration())
<|endoftext|>


<commit_before>import json
import os
import time
import re
import requests
import yaml
from flask import Flask, g
from peewee import *
from flask_restful import Api
from flask_debugtoolbar import DebugToolbarExtension
from datetime import date
from flask_cache import Cache
from gcloud import datastore

# Caching
app = Flask(__name__, static_url_path='/static')
app.config['TEMPLATES_AUTO_RELOAD'] = True
dbname = "cegwas_v2" # don't remove, imported elsewhere.

releases = ["20170531",
            "20160408"]

app.config['CURRENT_RELEASE'] = releases[0]

def get_vcf(release = releases[0]):
    return "http://storage.googleapis.com/elegansvariation.org/releases/{release}/WI.{release}.vcf.gz".format(release=release)

def get_ds():
    with app.app_context():
        if not hasattr(g, 'ds'):
            g.ds = datastore.Client(project="andersen-lab")
        return g.ds

ds = get_ds()

credentials = dict(ds.get(ds.key("credential", 'cegwas-data')))
db = MySQLDatabase(dbname, **credentials)

@app.before_request
def _db_connect():
    try:
        db.connect()
    except:
        time.sleep(1)
        db.connect()

@app.teardown_request
def _db_close(exc):
    if not db.is_closed():
        db.close()

biotypes = {
    "miRNA" : "microRNA",
    "piRNA" : "piwi-interacting RNA",
    "rRNA"  : "ribosomal RNA",
    "siRNA" : "small interfering RNA",
    "snRNA" : "small nuclear RNA",
    "snoRNA": "small nucleolar RNA",
    "tRNA"  : "transfer RNA",
    "vaultRNA" : "Short non-coding RNA genes that form part of the vault ribonucleoprotein complex",
    "lncRNA" : "Long non-coding RNA",
    "lincRNA" : "Long interspersed ncRNA",
    "pseudogene" : "non-functional gene.",
    "asRNA" : "Anti-sense RNA",
    "ncRNA" : "Non-coding RNA",
    "scRNA" : "Small cytoplasmic RNA"
}

def json_serial(obj):
    """JSON serializer for objects not serializable by default json code"""

    if isinstance(obj, date):
        serial = obj.isoformat()
        return serial
    raise TypeError("Type not serializable")

def autoconvert(s):
    for fn in (int, float):
        try:
            return fn(s)
        except ValueError:
            pass
    return s

class CustomEncoder(json.JSONEncoder):
    def default(self, o):
        if type(o) == decimal.Decimal:
            return float(o)
        if isinstance(o, datetime.date):
            return str(o)
        return super(CustomEncoder, self).default(o)

if os.getenv('HOME') == "/root":
    # Running on server
    # cache = Cache(app, config={'CACHE_TYPE': 'gaememcached'})
    # Cache service not yet available - use simple for now.
    cache = Cache(app, config={'CACHE_TYPE': 'simple'})
    app.debug = False
    app.config["debug"] = False
    from flask_sslify import SSLify
    # Ignore leading slash of urls; skips must use start of path
    sslify = SSLify(app, skips=['strains/global-strain-map',
                                '.well-known',
                                'cronmapping'])
else:
    # Running locally
    cache = Cache(app, config={'CACHE_TYPE': 'simple'})
    app.debug = True
    app.config["debug"] = True
    app.config['SECRET_KEY'] = "test"
    toolbar = DebugToolbarExtension(app)

version = re.search("VERSION=version-(.*)\n", open(".travis.yml", 'r').read()) \
            .group(1) \
            .replace('-', '.')

app.config['VERSION'] = version
app.config['TEMPLATE_AUTO_RELOAD'] = True
api = Api(app)
app.config['SECRET_KEY'] = os.getenv("SECRET_KEY")
app.config['DEBUG_TB_INTERCEPT_REDIRECTS'] = False

def get_google_order_sheet():
    with app.app_context():
        if not hasattr(g, 'gc'):
            import gspread
            from oauth2client.service_account import ServiceAccountCredentials
            ds = get_ds()
            sa = ds.get(ds.key("credential", "service_account"))
            sa = json.loads(sa["json"])
            scope = ['https://spreadsheets.google.com/feeds']
            credentials = ServiceAccountCredentials.from_json_keyfile_dict(sa, scope)
            gc = gspread.authorize(credentials)
            g.gc = gc
        return gc.open_by_key("1BCnmdJNRjQR3Bx8fMjD_IlTzmh3o7yj8ZQXTkk6tTXM").worksheet("orders")

def add_to_order_ws(row):
    """
        Stores order info in a google sheet.
    """
    ws = get_google_order_sheet()
    index = sum([1 for x in ws.col_values(1) if x]) + 1

    header_row = filter(len, ws.row_values(1))
    values = []
    for x in header_row:
        if x in row.keys():
            values.append(row[x])
        else:
            values.append("")

    row = map(str, row)
    ws.insert_row(values, index)

def lookup_order(invoice_hash):
    ws = get_google_order_sheet()
    find_row = ws.findall(invoice_hash)
    if len(find_row) > 0:
        row = ws.row_values(find_row[0].row)
        header_row = ws.row_values(1)
        result = dict(zip(header_row, row))
        return {k:v for k,v in result.items() if v}
    else:
        return None

def send_mail(data):
    ds = get_ds()
    api_key = ds.get(ds.key("credential", 'mailgun'))['apiKey']
    return requests.post(
        "https://api.mailgun.net/v3/mail.elegansvariation.org/messages",
        auth=("api", api_key),
        data=data)

#
# Custom Filters
#

@app.template_filter('comma')
def comma_filter(value):
    return "{:,.0f}".format(value)

@app.template_filter('format_release')
def format_release_filter(value):
    return datetime.strptime(str(value), '%Y%m%d').strftime('%Y-%m-%d')

from task import *
from views import *
from cegwas import *
<commit_msg>Fix json encoder<commit_after>import json
import os
import time
import re
import requests
import yaml
import decimal
from flask import Flask, g
from peewee import *
from flask_restful import Api
from flask_debugtoolbar import DebugToolbarExtension
from datetime import date
from flask_cache import Cache
from gcloud import datastore

# Caching
app = Flask(__name__, static_url_path='/static')
app.config['TEMPLATES_AUTO_RELOAD'] = True
dbname = "cegwas_v2" # don't remove, imported elsewhere.

releases = ["20170531",
            "20160408"]

app.config['CURRENT_RELEASE'] = releases[0]

def get_vcf(release = releases[0]):
    return "http://storage.googleapis.com/elegansvariation.org/releases/{release}/WI.{release}.vcf.gz".format(release=release)

def get_ds():
    with app.app_context():
        if not hasattr(g, 'ds'):
            g.ds = datastore.Client(project="andersen-lab")
        return g.ds

ds = get_ds()

credentials = dict(ds.get(ds.key("credential", 'cegwas-data')))
db = MySQLDatabase(dbname, **credentials)

@app.before_request
def _db_connect():
    try:
        db.connect()
    except:
        time.sleep(1)
        db.connect()

@app.teardown_request
def _db_close(exc):
    if not db.is_closed():
        db.close()

biotypes = {
    "miRNA" : "microRNA",
    "piRNA" : "piwi-interacting RNA",
    "rRNA"  : "ribosomal RNA",
    "siRNA" : "small interfering RNA",
    "snRNA" : "small nuclear RNA",
    "snoRNA": "small nucleolar RNA",
    "tRNA"  : "transfer RNA",
    "vaultRNA" : "Short non-coding RNA genes that form part of the vault ribonucleoprotein complex",
    "lncRNA" : "Long non-coding RNA",
    "lincRNA" : "Long interspersed ncRNA",
    "pseudogene" : "non-functional gene.",
    "asRNA" : "Anti-sense RNA",
    "ncRNA" : "Non-coding RNA",
    "scRNA" : "Small cytoplasmic RNA"
}

def json_serial(obj):
    """JSON serializer for objects not serializable by default json code"""

    if isinstance(obj, date):
        serial = obj.isoformat()
        return serial
    if type(obj) == decimal.Decimal:
        return float(obj)
    raise TypeError("Type not serializable")

def autoconvert(s):
    for fn in (int, float):
        try:
            return fn(s)
        except ValueError:
            pass
    return s

class CustomEncoder(json.JSONEncoder):
    def default(self, o):
        if type(o) == decimal.Decimal:
            return float(o)
        if isinstance(o, datetime.date):
            return str(o)
        return super(CustomEncoder, self).default(o)

if os.getenv('HOME') == "/root":
    # Running on server
    # cache = Cache(app, config={'CACHE_TYPE': 'gaememcached'})
    # Cache service not yet available - use simple for now.
    cache = Cache(app, config={'CACHE_TYPE': 'simple'})
    app.debug = False
    app.config["debug"] = False
    from flask_sslify import SSLify
    # Ignore leading slash of urls; skips must use start of path
    sslify = SSLify(app, skips=['strains/global-strain-map',
                                '.well-known',
                                'cronmapping'])
else:
    # Running locally
    cache = Cache(app, config={'CACHE_TYPE': 'simple'})
    app.debug = True
    app.config["debug"] = True
    app.config['SECRET_KEY'] = "test"
    toolbar = DebugToolbarExtension(app)

version = re.search("VERSION=version-(.*)\n", open(".travis.yml", 'r').read()) \
            .group(1) \
            .replace('-', '.')

app.config['VERSION'] = version
app.config['TEMPLATE_AUTO_RELOAD'] = True
api = Api(app)
app.config['SECRET_KEY'] = os.getenv("SECRET_KEY")
app.config['DEBUG_TB_INTERCEPT_REDIRECTS'] = False

def get_google_order_sheet():
    with app.app_context():
        if not hasattr(g, 'gc'):
            import gspread
            from oauth2client.service_account import ServiceAccountCredentials
            ds = get_ds()
            sa = ds.get(ds.key("credential", "service_account"))
            sa = json.loads(sa["json"])
            scope = ['https://spreadsheets.google.com/feeds']
            credentials = ServiceAccountCredentials.from_json_keyfile_dict(sa, scope)
            gc = gspread.authorize(credentials)
            g.gc = gc
        return gc.open_by_key("1BCnmdJNRjQR3Bx8fMjD_IlTzmh3o7yj8ZQXTkk6tTXM").worksheet("orders")

def add_to_order_ws(row):
    """
        Stores order info in a google sheet.
    """
    ws = get_google_order_sheet()
    index = sum([1 for x in ws.col_values(1) if x]) + 1

    header_row = filter(len, ws.row_values(1))
    values = []
    for x in header_row:
        if x in row.keys():
            values.append(row[x])
        else:
            values.append("")

    row = map(str, row)
    ws.insert_row(values, index)

def lookup_order(invoice_hash):
    ws = get_google_order_sheet()
    find_row = ws.findall(invoice_hash)
    if len(find_row) > 0:
        row = ws.row_values(find_row[0].row)
        header_row = ws.row_values(1)
        result = dict(zip(header_row, row))
        return {k:v for k,v in result.items() if v}
    else:
        return None

def send_mail(data):
    ds = get_ds()
    api_key = ds.get(ds.key("credential", 'mailgun'))['apiKey']
    return requests.post(
        "https://api.mailgun.net/v3/mail.elegansvariation.org/messages",
        auth=("api", api_key),
        data=data)

#
# Custom Filters
#

@app.template_filter('comma')
def comma_filter(value):
    return "{:,.0f}".format(value)

@app.template_filter('format_release')
def format_release_filter(value):
    return datetime.strptime(str(value), '%Y%m%d').strftime('%Y-%m-%d')

from task import *
from views import *
from cegwas import *
<|endoftext|>