<commit_before>import unittest

from pyjs import ast as js
from pyjs import serializer
from l20n.format.lol import ast as l20n
from l20n.compiler.js import Compiler, CompilerError

compile = Compiler.compile

class L20nJsCompilerTestCase(unittest.TestCase):
    def test_empty_lol(self):
        lol = l20n.LOL()
        prog = compile(lol)
        self.assertEqual(len(prog.body), 0)

    def test_empty_entity(self):
        lol = l20n.LOL(
            (l20n.Entity(l20n.Identifier('id')), )
        )
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(exp.right.value, None)

    def test_string_value(self):
        lol = l20n.LOL(
            (l20n.Entity(id=l20n.Identifier('id'),
                         value=l20n.String('foo')),)
        )
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(exp.right.value, "foo")

    def test_array_value(self):
        lol = l20n.LOL(
            (l20n.Entity(id=l20n.Identifier('id'),
                         value=l20n.Array((
                             l20n.String('foo'),
                             l20n.String('foo2'),))),)
        )
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(len(exp.right.elements), 2)
        self.assertEqual(exp.right.elements[0].value, 'foo')
        self.assertEqual(exp.right.elements[1].value, 'foo2')

    def test_hash_value(self):
        entity = l20n.Entity(id=l20n.Identifier('id'),
                         value=l20n.Hash((
                             l20n.HashItem(
                                 l20n.Identifier('foo'),
                                 l20n.String('foo value')),
                             l20n.HashItem(
                                 l20n.Identifier('foo2'),
                                 l20n.String('foo2 value'))
                         )))
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        s = serializer.Serializer()
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(len(exp.right.properties), 2)
        self.assertEqual(exp.right.properties[0].value.value, 'foo value')
        self.assertEqual(exp.right.properties[1].value.value, 'foo2 value')

    def test_attributes(self):
        attrs = {'a': l20n.Attribute(
                         l20n.Identifier('a'),
                         l20n.String('foo')),

                 'b': l20n.Attribute(
                         l20n.Identifier('b'),
                         l20n.String('foo2')
                 )}
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=l20n.String('foo'),
                             attrs=attrs)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 2)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)

        self.assertTrue(isinstance(exp.right, js.NewExpression))
        self.assertEqual(exp.right.callee.name, "String")
        self.assertEqual(exp.right.arguments[0].value, "foo")

        exp = prog.body[1].expression
        self.assertEqual(exp.operator.token, '=')
        self.assertTrue(isinstance(exp.left.obj, js.MemberExpression))
        self.assertEqual(exp.left.obj.prop.name, 'id')
        self.assertEqual(exp.left.prop.name, "_attrs")
        self.assertEqual(exp.left.computed, False)

    def test_property_expression(self):
        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.Identifier('foo'),
                l20n.Identifier('a'),
                False),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[1].name, "sys")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 1)
        self.assertEqual(getent_arguments[3].elements[0].value, 'a')

        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.PropertyExpression(
                    l20n.Identifier('foo'),
                    l20n.Identifier('a'),
                    False),
                l20n.Identifier('b'),
                False),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[1].name, "sys")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 2)
        self.assertEqual(getent_arguments[3].elements[0].value, 'a')
        self.assertEqual(getent_arguments[3].elements[1].value, 'b')

        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.PropertyExpression(
                    l20n.PropertyExpression(
                        l20n.Identifier('foo'),
                        l20n.String('a'),
                        True),
                    l20n.Identifier('b'),
                    False),
                l20n.String('c'),
                True),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 3)
        self.assertEqual(getent_arguments[3].elements[0].value, 'a')
        self.assertEqual(getent_arguments[3].elements[1].value, 'b')
        self.assertEqual(getent_arguments[3].elements[2].value, 'c')

    def test_attribute_expression(self):
        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.PropertyExpression(
                    l20n.AttributeExpression(
                        l20n.Identifier('foo'),
                        l20n.String('a'),
                        True),
                    l20n.Identifier('b'),
                    False),
                l20n.String('c'),
                True),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getattr")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(getent_arguments[3].value, "a")
        self.assertEqual(len(getent_arguments[4].elements), 2)
        self.assertEqual(getent_arguments[4].elements[0].value, 'b')
        self.assertEqual(getent_arguments[4].elements[1].value, 'c')

    def test_computed_property_expression(self):
        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.Identifier('foo'),
                l20n.Identifier('a'),
                True),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 1)
        subget = getent_arguments[3].elements[0]
        self.assertEqual(subget.prop.callee.name, 'getent')
        self.assertEqual(subget.prop.arguments[0].name, 'env')
        self.assertEqual(subget.prop.arguments[2].value, 'a')

    def test_macro(self):
        m = l20n.Identifier('foo')
        args = [
             l20n.VariableExpression(l20n.Identifier('n')),
        ]
        exp = l20n.VariableExpression(l20n.Identifier('n'))
        exp = l20n.BinaryExpression(
            l20n.BinaryOperator('+'),
            l20n.VariableExpression(l20n.Identifier('n')),
            l20n.Identifier('n')
        )
        macro = l20n.Macro(id=m, args=args, expression=exp)

        lol = l20n.LOL((macro,))
        prog = compile(lol)
        s = serializer.Serializer()
        print(s.dump_program(prog))

        macro = prog.body[0].expression
        value = macro.right
        val_body = value.body
        val_return = val_body.body[0]
        return_exp = val_return.argument
        data = return_exp.left
        self.assertEqual(data.obj.name, 'data')
        self.assertEqual(data.prop.value, 0)
        getent = return_exp.right
        self.assertEqual(getent.obj.name, 'sys')
        self.assertEqual(getent.prop.callee.name, 'getent')
        self.assertEqual(getent.prop.arguments[1].name, 'sys')
        self.assertEqual(getent.prop.arguments[0].name, 'env')
        self.assertEqual(getent.prop.arguments[2].value, 'n')

if __name__ == '__main__':

    unittest.main()

<commit_msg>remove debug print from tests<commit_after>import unittest

from pyjs import ast as js
from pyjs import serializer
from l20n.format.lol import ast as l20n
from l20n.compiler.js import Compiler, CompilerError

compile = Compiler.compile

class L20nJsCompilerTestCase(unittest.TestCase):
    def test_empty_lol(self):
        lol = l20n.LOL()
        prog = compile(lol)
        self.assertEqual(len(prog.body), 0)

    def test_empty_entity(self):
        lol = l20n.LOL(
            (l20n.Entity(l20n.Identifier('id')), )
        )
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(exp.right.value, None)

    def test_string_value(self):
        lol = l20n.LOL(
            (l20n.Entity(id=l20n.Identifier('id'),
                         value=l20n.String('foo')),)
        )
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(exp.right.value, "foo")

    def test_array_value(self):
        lol = l20n.LOL(
            (l20n.Entity(id=l20n.Identifier('id'),
                         value=l20n.Array((
                             l20n.String('foo'),
                             l20n.String('foo2'),))),)
        )
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(len(exp.right.elements), 2)
        self.assertEqual(exp.right.elements[0].value, 'foo')
        self.assertEqual(exp.right.elements[1].value, 'foo2')

    def test_hash_value(self):
        entity = l20n.Entity(id=l20n.Identifier('id'),
                         value=l20n.Hash((
                             l20n.HashItem(
                                 l20n.Identifier('foo'),
                                 l20n.String('foo value')),
                             l20n.HashItem(
                                 l20n.Identifier('foo2'),
                                 l20n.String('foo2 value'))
                         )))
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        s = serializer.Serializer()
        self.assertEqual(len(prog.body), 1)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)
        self.assertEqual(len(exp.right.properties), 2)
        self.assertEqual(exp.right.properties[0].value.value, 'foo value')
        self.assertEqual(exp.right.properties[1].value.value, 'foo2 value')

    def test_attributes(self):
        attrs = {'a': l20n.Attribute(
                         l20n.Identifier('a'),
                         l20n.String('foo')),

                 'b': l20n.Attribute(
                         l20n.Identifier('b'),
                         l20n.String('foo2')
                 )}
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=l20n.String('foo'),
                             attrs=attrs)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 2)
        exp = prog.body[0].expression
        self.assertEqual(exp.operator.token, "=")
        self.assertTrue(isinstance(exp.left.obj, js.ThisExpression))
        self.assertEqual(exp.left.prop.name, "id")
        self.assertEqual(exp.left.computed, False)

        self.assertTrue(isinstance(exp.right, js.NewExpression))
        self.assertEqual(exp.right.callee.name, "String")
        self.assertEqual(exp.right.arguments[0].value, "foo")

        exp = prog.body[1].expression
        self.assertEqual(exp.operator.token, '=')
        self.assertTrue(isinstance(exp.left.obj, js.MemberExpression))
        self.assertEqual(exp.left.obj.prop.name, 'id')
        self.assertEqual(exp.left.prop.name, "_attrs")
        self.assertEqual(exp.left.computed, False)

    def test_property_expression(self):
        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.Identifier('foo'),
                l20n.Identifier('a'),
                False),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[1].name, "sys")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 1)
        self.assertEqual(getent_arguments[3].elements[0].value, 'a')

        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.PropertyExpression(
                    l20n.Identifier('foo'),
                    l20n.Identifier('a'),
                    False),
                l20n.Identifier('b'),
                False),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[1].name, "sys")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 2)
        self.assertEqual(getent_arguments[3].elements[0].value, 'a')
        self.assertEqual(getent_arguments[3].elements[1].value, 'b')

        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.PropertyExpression(
                    l20n.PropertyExpression(
                        l20n.Identifier('foo'),
                        l20n.String('a'),
                        True),
                    l20n.Identifier('b'),
                    False),
                l20n.String('c'),
                True),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 3)
        self.assertEqual(getent_arguments[3].elements[0].value, 'a')
        self.assertEqual(getent_arguments[3].elements[1].value, 'b')
        self.assertEqual(getent_arguments[3].elements[2].value, 'c')

    def test_attribute_expression(self):
        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.PropertyExpression(
                    l20n.AttributeExpression(
                        l20n.Identifier('foo'),
                        l20n.String('a'),
                        True),
                    l20n.Identifier('b'),
                    False),
                l20n.String('c'),
                True),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getattr")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(getent_arguments[3].value, "a")
        self.assertEqual(len(getent_arguments[4].elements), 2)
        self.assertEqual(getent_arguments[4].elements[0].value, 'b')
        self.assertEqual(getent_arguments[4].elements[1].value, 'c')

    def test_computed_property_expression(self):
        string = l20n.ComplexString([
            l20n.String("word1 "),
            l20n.PropertyExpression(
                l20n.Identifier('foo'),
                l20n.Identifier('a'),
                True),
            l20n.String(" word2")
        ])
        entity = l20n.Entity(id=l20n.Identifier('id'),
                             value=string)
        lol = l20n.LOL((entity,))
        prog = compile(lol)
        self.assertEqual(len(prog.body), 1)
        entity = prog.body[0].expression
        value = entity.right
        val_body = value.body
        val_return = val_body.body[0]
        complex_string = val_return.argument
        first_part = complex_string.left
        getent_mexp = complex_string.right.left
        third_part = complex_string.right.right
        getent_id = getent_mexp.obj
        self.assertEqual(getent_id.name, "sys")
        getent_call = getent_mexp.prop
        getent_arguments = getent_call.arguments
        self.assertEqual(getent_call.callee.name, "getent")
        self.assertEqual(getent_arguments[0].name, "env")
        self.assertEqual(getent_arguments[2].value, "foo")
        self.assertEqual(len(getent_arguments[3].elements), 1)
        subget = getent_arguments[3].elements[0]
        self.assertEqual(subget.prop.callee.name, 'getent')
        self.assertEqual(subget.prop.arguments[0].name, 'env')
        self.assertEqual(subget.prop.arguments[2].value, 'a')

    def test_macro(self):
        m = l20n.Identifier('foo')
        args = [
             l20n.VariableExpression(l20n.Identifier('n')),
        ]
        exp = l20n.VariableExpression(l20n.Identifier('n'))
        exp = l20n.BinaryExpression(
            l20n.BinaryOperator('+'),
            l20n.VariableExpression(l20n.Identifier('n')),
            l20n.Identifier('n')
        )
        macro = l20n.Macro(id=m, args=args, expression=exp)

        lol = l20n.LOL((macro,))
        prog = compile(lol)
        macro = prog.body[0].expression
        value = macro.right
        val_body = value.body
        val_return = val_body.body[0]
        return_exp = val_return.argument
        data = return_exp.left
        self.assertEqual(data.obj.name, 'data')
        self.assertEqual(data.prop.value, 0)
        getent = return_exp.right
        self.assertEqual(getent.obj.name, 'sys')
        self.assertEqual(getent.prop.callee.name, 'getent')
        self.assertEqual(getent.prop.arguments[1].name, 'sys')
        self.assertEqual(getent.prop.arguments[0].name, 'env')
        self.assertEqual(getent.prop.arguments[2].value, 'n')

if __name__ == '__main__':

    unittest.main()

<|endoftext|>


<commit_before>import sys
import argparse
import subprocess
import os

queue = 'staff'
__version__ = '2.5.0'
prefix = 'trozamon_testing'
input_unstructured = '/'.join([prefix, 'input_unstructured'])
input_structured = '/'.join([prefix, 'input_structured'])
output_prefix = '/'.join([prefix, 'output'])
mrjob_conf = 'python/mrjob.conf'

big_description = """Hadoop and related software cluster tester. Given the
fast-paced nature of Hadoop, Spark, Hive, and other technologies, running a
test suite after an upgrade can ensure that all expected functionality
continues to work.  This script will generate and upload data to
$HOME/trozamon_testing in HDFS."""

def hdfs_mkdir(dir):
    if subprocess.call('hdfs dfs -test -d ' + dir, shell=True) != 0:
        if subprocess.call('hdfs dfs -mkdir ' + dir, shell=True) != 0:
            print('Failed to create ' + dir + ' in HDFS')
            return 1
    return 0

def hdfs_put(local, remote):
    if subprocess.call('hdfs dfs -test -f ' + remote, shell=True) != 0:
        if subprocess.call('hdfs dfs -put ' + local + ' ' + remote,
                shell=True) != 0:
            print('Failed to upload ' + local + ' to ' + remote)
            return 1
    return 0

def hdfs_rmdir(dir):
    if subprocess.call('hdfs dfs -test -d ' + dir, shell=True) == 0:
        return subprocess.call('hdfs dfs -rm -r ' + dir, shell=True)
    return 0

def test_hadoop_java():
    outdir = '/'.join([output_prefix, 'hadoop_java'])
    cmd = ' '.join([
        'yarn jar hadoop/target/hadoop-examples-hadoop-' + __version__ + '.jar',
        'com.alectenharmsel.research.LineCount',
        '-Dmapreduce.job.queuename=' + queue,
        input_unstructured,
        outdir
        ])

    if hdfs_rmdir(outdir) != 0:
        return 1

    print('Testing Hadoop Java by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_mrjob():
    cmd = ' '.join([
        'python2.7',
        'python/mrjob_test.py',
        'structured.data',
        '-r',
        'hadoop'
        ])

    print('Testing MRJob by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_hadoop_streaming():
    outdir = '/'.join([output_prefix, 'hadoop_streaming'])
    cmd = ' '.join([
        'yarn jar /usr/lib/hadoop-mapreduce/hadoop-streaming.jar',
        '-Dmapreduce.job.queuename=' + queue,
        '-input ' + input_unstructured,
        '-output ' + outdir,
        '-mapper python/srctok-map.py -reducer python/sum.py',
        '-file python/srctok-map.py -file python/sum.py'
        ])

    if hdfs_rmdir(outdir) != 0:
        return 1

    print('Testing Hadoop Streaming by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_pig():
    cmd = ' '.join([
        'pig',
        '-Dmapreduce.job.queuename=' + queue,
        '-f pig/cluster_test.pig'
        ])

    print('Testing Pig by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_spark():
    outdir = '/'.join([output_prefix, 'spark'])
    cmd = ' '.join([
        'spark-submit',
        '--master yarn-client',
        '--queue ' + queue,
        '--class com.alectenharmsel.research.spark.LineCount',
        'spark/target/hadoop-examples-spark-' + __version__ + '.jar',
        input_unstructured,
        outdir
        ])

    if hdfs_rmdir(outdir) != 0:
        return 1

    print('Testing Spark by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_pyspark():
    cmd = ' '.join([
        'spark-submit',
        '--master yarn-client',
        '--queue ' + queue,
        'python/spark/lc.py',
        input_unstructured
        ])

    print('Testing PySpark by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_hive():
    cmd = ' '.join([
        'hive',
        '--hiveconf mapreduce.job.queuename=' + queue,
        '-f hive/cluster_test.sql'
        ])

    print('Testing Hive by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def run():
    parser = argparse.ArgumentParser(description=big_description)

    parser.add_argument('--all',
            action='store_true',
            help='Run all codes that would be run by using all possible flags')
    parser.add_argument('--hadoop-java',
            action='store_true',
            help='Run a selection of Java codes written for Hadoop')
    parser.add_argument('--hadoop-streaming',
            action='store_true',
            help='Run a selection of Python codes written for Hadoop')
    parser.add_argument('--pig',
            action='store_true',
            help='Run a selection of Pig scripts')
    parser.add_argument('--spark',
            action='store_true',
            help='Run a selection of Scala codes written for Spark')
    parser.add_argument('--pyspark',
            action='store_true',
            help='Run a selection of Python codes written for Spark')
    parser.add_argument('--hive',
            action='store_true',
            help='Run a selection of Hive queries')
    parser.add_argument('--mrjob',
            action='store_true',
            help='Run a mrjob script')

    parsed_args = parser.parse_args()

    tests = list()
    if parsed_args.all:
        parsed_args.hadoop_java = True
        parsed_args.hadoop_streaming = True
        parsed_args.pig = True
        parsed_args.hive = True
        parsed_args.spark = True
        parsed_args.pyspark = True
        parsed_args.mrjob = True

    if parsed_args.hadoop_java:
        tests.append(test_hadoop_java)
    if parsed_args.hadoop_streaming:
        tests.append(test_hadoop_streaming)
    if parsed_args.pig:
        tests.append(test_pig)
    if parsed_args.pyspark:
        tests.append(test_pyspark)
    if parsed_args.spark:
        tests.append(test_spark)
    if parsed_args.mrjob:
        os.environ['MRJOB_CONF'] = '/'.join([os.environ['PWD'], mrjob_conf])
        tests.append(test_mrjob)
    # Hive should always be last - it munges data
    if parsed_args.hive:
        tests.append(test_hive)

    print('Running "mvn package -DskipTests"...')
    if subprocess.call('mvn package -DskipTests', shell=True) != 0:
        print('Maven build failed. Are you running this in the root ' +
                'directory of the repo?')
        return 1

    print('Creating directory structure...')
    err = hdfs_mkdir(prefix)
    if err != 0:
        return err
    err = hdfs_mkdir(input_unstructured)
    if err != 0:
        return err
    err = hdfs_mkdir(input_structured)
    if err != 0:
        return err
    err = hdfs_mkdir(output_prefix)
    if err != 0:
        return err
    print('Uploading unstructured data...')
    err = hdfs_put('pom.xml', '/'.join([input_unstructured, 'pom.xml']))
    if err != 0:
        return err
    print('Uploading structured data...')
    err = hdfs_put('structured.data', '/'.join([input_structured, 'data']))
    if err != 0:
        return err

    print('Starting to run tests...')
    results = {}
    ret = 0
    for test in tests:
        res = test()
        if res == 0:
            results[test.__name__] = True
        else:
            ret = 1
            results[test.__name__] = False

    for test in sorted(results.keys()):
        if results[test]:
            print(test + ': SUCCESS')
        else:
            print(test + ': FAILURE')

    return ret

if __name__ == '__main__':
    sys.exit(run())
<commit_msg>Update version in cluster_test.py<commit_after>import sys
import argparse
import subprocess
import os

queue = 'staff'
__version__ = '2.5.0-cdh5.3.2'
prefix = 'trozamon_testing'
input_unstructured = '/'.join([prefix, 'input_unstructured'])
input_structured = '/'.join([prefix, 'input_structured'])
output_prefix = '/'.join([prefix, 'output'])
mrjob_conf = 'python/mrjob.conf'

big_description = """Hadoop and related software cluster tester. Given the
fast-paced nature of Hadoop, Spark, Hive, and other technologies, running a
test suite after an upgrade can ensure that all expected functionality
continues to work.  This script will generate and upload data to
$HOME/trozamon_testing in HDFS."""

def hdfs_mkdir(dir):
    if subprocess.call('hdfs dfs -test -d ' + dir, shell=True) != 0:
        if subprocess.call('hdfs dfs -mkdir ' + dir, shell=True) != 0:
            print('Failed to create ' + dir + ' in HDFS')
            return 1
    return 0

def hdfs_put(local, remote):
    if subprocess.call('hdfs dfs -test -f ' + remote, shell=True) != 0:
        if subprocess.call('hdfs dfs -put ' + local + ' ' + remote,
                shell=True) != 0:
            print('Failed to upload ' + local + ' to ' + remote)
            return 1
    return 0

def hdfs_rmdir(dir):
    if subprocess.call('hdfs dfs -test -d ' + dir, shell=True) == 0:
        return subprocess.call('hdfs dfs -rm -r ' + dir, shell=True)
    return 0

def test_hadoop_java():
    outdir = '/'.join([output_prefix, 'hadoop_java'])
    cmd = ' '.join([
        'yarn jar hadoop/target/hadoop-examples-hadoop-' + __version__ + '.jar',
        'com.alectenharmsel.research.LineCount',
        '-Dmapreduce.job.queuename=' + queue,
        input_unstructured,
        outdir
        ])

    if hdfs_rmdir(outdir) != 0:
        return 1

    print('Testing Hadoop Java by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_mrjob():
    cmd = ' '.join([
        'python2.7',
        'python/mrjob_test.py',
        'structured.data',
        '-r',
        'hadoop'
        ])

    print('Testing MRJob by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_hadoop_streaming():
    outdir = '/'.join([output_prefix, 'hadoop_streaming'])
    cmd = ' '.join([
        'yarn jar /usr/lib/hadoop-mapreduce/hadoop-streaming.jar',
        '-Dmapreduce.job.queuename=' + queue,
        '-input ' + input_unstructured,
        '-output ' + outdir,
        '-mapper python/srctok-map.py -reducer python/sum.py',
        '-file python/srctok-map.py -file python/sum.py'
        ])

    if hdfs_rmdir(outdir) != 0:
        return 1

    print('Testing Hadoop Streaming by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_pig():
    cmd = ' '.join([
        'pig',
        '-Dmapreduce.job.queuename=' + queue,
        '-f pig/cluster_test.pig'
        ])

    print('Testing Pig by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_spark():
    outdir = '/'.join([output_prefix, 'spark'])
    cmd = ' '.join([
        'spark-submit',
        '--master yarn-client',
        '--queue ' + queue,
        '--class com.alectenharmsel.research.spark.LineCount',
        'spark/target/hadoop-examples-spark-' + __version__ + '.jar',
        input_unstructured,
        outdir
        ])

    if hdfs_rmdir(outdir) != 0:
        return 1

    print('Testing Spark by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_pyspark():
    cmd = ' '.join([
        'spark-submit',
        '--master yarn-client',
        '--queue ' + queue,
        'python/spark/lc.py',
        input_unstructured
        ])

    print('Testing PySpark by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def test_hive():
    cmd = ' '.join([
        'hive',
        '--hiveconf mapreduce.job.queuename=' + queue,
        '-f hive/cluster_test.sql'
        ])

    print('Testing Hive by running:')
    print(cmd)

    return subprocess.call(cmd, shell=True)

def run():
    parser = argparse.ArgumentParser(description=big_description)

    parser.add_argument('--all',
            action='store_true',
            help='Run all codes that would be run by using all possible flags')
    parser.add_argument('--hadoop-java',
            action='store_true',
            help='Run a selection of Java codes written for Hadoop')
    parser.add_argument('--hadoop-streaming',
            action='store_true',
            help='Run a selection of Python codes written for Hadoop')
    parser.add_argument('--pig',
            action='store_true',
            help='Run a selection of Pig scripts')
    parser.add_argument('--spark',
            action='store_true',
            help='Run a selection of Scala codes written for Spark')
    parser.add_argument('--pyspark',
            action='store_true',
            help='Run a selection of Python codes written for Spark')
    parser.add_argument('--hive',
            action='store_true',
            help='Run a selection of Hive queries')
    parser.add_argument('--mrjob',
            action='store_true',
            help='Run a mrjob script')

    parsed_args = parser.parse_args()

    tests = list()
    if parsed_args.all:
        parsed_args.hadoop_java = True
        parsed_args.hadoop_streaming = True
        parsed_args.pig = True
        parsed_args.hive = True
        parsed_args.spark = True
        parsed_args.pyspark = True
        parsed_args.mrjob = True

    if parsed_args.hadoop_java:
        tests.append(test_hadoop_java)
    if parsed_args.hadoop_streaming:
        tests.append(test_hadoop_streaming)
    if parsed_args.pig:
        tests.append(test_pig)
    if parsed_args.pyspark:
        tests.append(test_pyspark)
    if parsed_args.spark:
        tests.append(test_spark)
    if parsed_args.mrjob:
        os.environ['MRJOB_CONF'] = '/'.join([os.environ['PWD'], mrjob_conf])
        tests.append(test_mrjob)
    # Hive should always be last - it munges data
    if parsed_args.hive:
        tests.append(test_hive)

    print('Running "mvn package -DskipTests"...')
    if subprocess.call('mvn package -DskipTests', shell=True) != 0:
        print('Maven build failed. Are you running this in the root ' +
                'directory of the repo?')
        return 1

    print('Creating directory structure...')
    err = hdfs_mkdir(prefix)
    if err != 0:
        return err
    err = hdfs_mkdir(input_unstructured)
    if err != 0:
        return err
    err = hdfs_mkdir(input_structured)
    if err != 0:
        return err
    err = hdfs_mkdir(output_prefix)
    if err != 0:
        return err
    print('Uploading unstructured data...')
    err = hdfs_put('pom.xml', '/'.join([input_unstructured, 'pom.xml']))
    if err != 0:
        return err
    print('Uploading structured data...')
    err = hdfs_put('structured.data', '/'.join([input_structured, 'data']))
    if err != 0:
        return err

    print('Starting to run tests...')
    results = {}
    ret = 0
    for test in tests:
        res = test()
        if res == 0:
            results[test.__name__] = True
        else:
            ret = 1
            results[test.__name__] = False

    for test in sorted(results.keys()):
        if results[test]:
            print(test + ': SUCCESS')
        else:
            print(test + ': FAILURE')

    return ret

if __name__ == '__main__':
    sys.exit(run())
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from runner.koan import *

class AboutSets(Koan):
    def test_sets_make_keep_lists_unique(self):
        highlanders = ['MacLeod', 'Ramirez', 'MacLeod', 'Matunas', 'MacLeod', 'Malcolm', 'MacLeod']

        there_can_only_be_only_one = set(highlanders)

        self.assertEqual(__, there_can_only_be_only_one)

    def test_empty_sets_have_different_syntax_to_populated_sets(self):
        self.assertEqual(__, {1, 2, 3})
        self.assertEqual(__, set())

    def test_dictionaries_and_sets_use_same_curly_braces(self):
        # Note: Sets have only started using braces since Python 3

        self.assertEqual(__, {1, 2, 3}.__class__)
        self.assertEqual(__, {'one': 1, 'two': 2}.__class__)

        self.assertEqual(__, {}.__class__)

    def test_creating_sets_using_strings(self):
        self.assertEqual(__, {'12345'})
        self.assertEqual(__, set('12345'))

    def test_convert_the_set_into_a_list_to_sort_it(self):
        self.assertEqual(__, sorted(set('12345')))

    # ------------------------------------------------------------------

    def test_set_have_arithmetic_operators(self):
        scotsmen = {'MacLeod', 'Wallace', 'Willie'}
        warriors = {'MacLeod', 'Wallace', 'Leonidas'}

        self.assertEqual(__, scotsmen - warriors)
        self.assertEqual(__, scotsmen | warriors)
        self.assertEqual(__, scotsmen & warriors)
        self.assertEqual(__, scotsmen ^ warriors)

    # ------------------------------------------------------------------

    def test_we_can_query_set_membership(self):
        self.assertEqual(__, 127 in {127, 0, 0, 1} )
        self.assertEqual(__, 'cow' not in set('apocalypse now') )

    def test_we_can_compare_subsets(self):
        self.assertEqual(__, set('cake') <= set('cherry cake'))
        self.assertEqual(__, set('cake').issubset(set('cherry cake')) )

        self.assertEqual(__, set('cake') > set('pie'))
<commit_msg>Koan done: sets<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from runner.koan import *

class AboutSets(Koan):
    def test_sets_make_keep_lists_unique(self):
        highlanders = ['MacLeod', 'Ramirez', 'MacLeod', 'Matunas', 'MacLeod', 'Malcolm', 'MacLeod']

        there_can_only_be_only_one = set(highlanders)

        self.assertEqual({'MacLeod', 'Ramirez', 'Malcolm', 'Matunas'}, there_can_only_be_only_one)

    def test_empty_sets_have_different_syntax_to_populated_sets(self):
        self.assertEqual(set((1,2,3)), {1, 2, 3})
        self.assertEqual(set(), set())

    def test_dictionaries_and_sets_use_same_curly_braces(self):
        # Note: Sets have only started using braces since Python 3

        self.assertEqual(set, {1, 2, 3}.__class__)
        self.assertEqual(dict, {'one': 1, 'two': 2}.__class__)

        self.assertEqual(dict, {}.__class__)

    def test_creating_sets_using_strings(self):
        self.assertEqual({'12345'}, {'12345'})
        self.assertEqual({'1','2','3','4','5'}, set('12345'))

    def test_convert_the_set_into_a_list_to_sort_it(self):
        self.assertEqual(['1','2','3','4','5'], sorted(set('12345')))

    # ------------------------------------------------------------------

    def test_set_have_arithmetic_operators(self):
        scotsmen = {'MacLeod', 'Wallace', 'Willie'}
        warriors = {'MacLeod', 'Wallace', 'Leonidas'}

        self.assertEqual({"Willie",}, scotsmen - warriors)
        self.assertEqual({"MacLeod","Wallace","Willie","Leonidas"}, scotsmen | warriors)
        self.assertEqual({"MacLeod","Wallace"}, scotsmen & warriors)
        self.assertEqual({"Willie",'Leonidas'}, scotsmen ^ warriors)

    # ------------------------------------------------------------------

    def test_we_can_query_set_membership(self):
        self.assertEqual(True, 127 in {127, 0, 0, 1} )
        self.assertEqual(True, 'cow' not in set('apocalypse now') )

    def test_we_can_compare_subsets(self):
        self.assertEqual(True, set('cake') <= set('cherry cake'))
        self.assertEqual(True, set('cake').issubset(set('cherry cake')) )

        self.assertEqual(False, set('cake') > set('pie'))
<|endoftext|>


<commit_before>#!python
"""
rcping - ping a remote robotcomm instance with messages, commands and rt-commands
Author: JMJ
"""
import argparse
import logging
import pprint
import re
import string
import sys

from robotutils import logging_helper
from robotutils.comm_helper import EchoClient

_LOGGER = logging.getLogger('rcping')
_TRACE = logging_helper.LevelSpecificLogger(logging_helper.TRACELEVEL, _LOGGER)

def generate_argparser():
    """Generate the argument parser for this utility"""
    parser = argparse.ArgumentParser(description='Robotcomm ping utility',
                                     allow_abbrev=False)

    # Just one positional argument
    parser.add_argument('address', help='hostname[:port] of destination')

    # Number of messages - mutually exclusive options
    numgroup = parser.add_mutually_exclusive_group()
    numgroup.add_argument('-n', dest='count', metavar='count', action='store',
                          type=int, default=4, help='count of things to send')
    numgroup.add_argument('-t', dest='count', action='store_const', const=None,
                          help='send indifinitely')

    # Message size or explicit message - mutually exclusive options
    sizegroup = parser.add_mutually_exclusive_group()
    sizegroup.add_argument('-l', dest='size', metavar='size', action='store',
                           type=int, help='size of auto payload')
    payload_help = "send PAYLOAD, which has the form mytype[::mybody] or ::mybody"
    sizegroup.add_argument('-payload', help=payload_help)

    # Type of command - mutually exclusive, and exactly must be specified
    cmdgroup = parser.add_mutually_exclusive_group(required=True)
    cmdgroup.add_argument('-msg', action='store_true', help='send messages')
    cmdgroup.add_argument('-cmd', action='store_true', help='send commands')
    cmdgroup.add_argument('-rtcmd', action='store_true', help='send rt-commands')

    parser.add_argument('-c', dest='channel', metavar='CHANNEL',
                        help='name of channel')
    parser.add_argument('-rate', type=float, default=1.0,
                        help='rate in sends per second')
    choices_text = "TRACE DEBUG INFO WARNING ERROR CRITICAL"
    choices_list = choices_text.split()
    loglevel_help = 'sets logging level. LOGLEVEL is one of: ' + choices_text
    parser.add_argument('-loglevel', default='ERROR', choices=choices_list,
                        help=loglevel_help)

    return parser

def parse_args(args):
    """Parse input args, including all error handling.
    Returns a parameters object on success. Exits program
    on failure.
    """
    parser = generate_argparser()
    params = parser.parse_args(args)
    try:
        hostname, port = parse_address(params.address)
        params.hostname = hostname
        params.port = port
        msgtype, msgbody = parse_payload(params.payload)
        params.msgtype = msgtype
        params.msgbody = msgbody
    except ValueError as exp:
        parser.error(str(exp))
    return params

# Note: IPV6 addresses are not supported - it would
# be extra work to distinguish a port from a ':' inside
# the IPv6 address.
_HOSTNAME_REGEX = re.compile(r'(\w|\.)+')
_PORTRANGE = 41 # Range 41000-41999
def parse_address(address):
    """Parse address of the form hostname[:port]"""
    errmsg = '\n'.join(("Invalid address '{}'".format(address),
                        "Hostname should have the form NAME_OR_IP[:PORT]"))
    hostname, *rest = address.split(':')
    if not _HOSTNAME_REGEX.fullmatch(hostname) or len(rest) > 1:
        raise ValueError(errmsg)
    port = int(rest[0]) if rest else None
    if port is not None and port//1000 != _PORTRANGE:
        msg = "Port must be in the range {} to {}"
        minport = _PORTRANGE*1000
        raise ValueError(msg.format(minport, minport+999))
    return (hostname, port)

_BAD_MSGTYPE_CHARS = string.whitespace + ','
_PAYLOAD_SEPARATOR = '::'
def parse_payload(payload):
    """Parse payload, which has the form type or ::body or type::body"""
    msgtype, msgbody = None, None
    if payload:
        errmsg = '\n'.join(("Invalid payload '{}'".format(payload),
                            "Payload should have the form TYPE or ::TYPE or TYPE::BODY"))
        # We pick the first instance of '::'
        try:
            index = payload.index(_PAYLOAD_SEPARATOR)
            msgtype = payload[:index] or None
            msgbody = payload[index+len(_PAYLOAD_SEPARATOR):] # skip past '::'
            msgbody = msgbody or None # convert '' to None
        except ValueError:
            msgtype = payload
        if msgtype and containschars(msgtype, _BAD_MSGTYPE_CHARS):
            raise ValueError(errmsg)
    return (msgtype, msgbody)

def send_messages(client, count):
    """Send messages using an instance of echo client"""
    receive_count = 0

    def send_handler(resptype, respbody):
        msg = "Sending: '{}::{}'".format(resptype, respbody)
        _TRACE(msg)
        print(msg)

    def response_handler(resptype, respbody):
        msg = "Response: '{}::{}'".format(resptype, respbody)
        _TRACE(msg)
        print(msg)
        nonlocal receive_count
        receive_count += 1 # assume call to hander is serialized

    # send_messages will block until done...
    _TRACE("GOING TO SEND MESSAGES")
    client.send_messages(count, send_handler=send_handler,
                         response_handler=response_handler)
    _TRACE("DONE SENDING MESSAGES")
    print("Received = {}".format(receive_count))

def set_loglevel(strloglevel):
    """Sets up logging with the specified logging level"""
    if strloglevel == 'TRACE':
        # Special case - not defined in logging module
        level = logging_helper.TRACELEVEL
    else:
        choices = "DEBUG INFO WARNING ERROR CRITICAL".split()
        index = choices.index(strloglevel)
        level = getattr(logging, choices[index])
    logging.basicConfig(level=level)

def main(args):
    """Main entry point"""
    params = parse_args(args)
    set_loglevel(params.loglevel)
    _LOGGER.info("parameters:\n%s", pprint.pformat(vars(params)))

    client = EchoClient('localhost')
    try:
        client.set_parameters(rate=params.rate)
        if params.msg:
            send_messages(client, params.count)
        elif params.cmd:
            print('send_commands(params)')
        elif params.rtcmd:
            print('send_rtcommands(params)')
    except Exception: # pylint: disable=broad-except
        _LOGGER.exception()
    finally:
        client.close()

def containschars(str_, charset) -> bool:
    """Returns if {str} contains any chars in {chars}"""
    for char in str_:
        if char in charset:
            return True
    return False

main(sys.argv[1:])
<commit_msg>rcping: Setting all supported EchoClient parameters<commit_after>#!python
"""
rcping - ping a remote robotcomm instance with messages, commands and rt-commands
Author: JMJ
"""
import argparse
import logging
import pprint
import re
import string
import sys

from robotutils import logging_helper
from robotutils.comm_helper import EchoClient

_LOGGER = logging.getLogger('rcping')
_TRACE = logging_helper.LevelSpecificLogger(logging_helper.TRACELEVEL, _LOGGER)

def generate_argparser():
    """Generate the argument parser for this utility"""
    parser = argparse.ArgumentParser(description='Robotcomm ping utility',
                                     allow_abbrev=False)

    # Just one positional argument
    parser.add_argument('address', help='hostname[:port] of destination')

    # Number of messages - mutually exclusive options
    numgroup = parser.add_mutually_exclusive_group()
    numgroup.add_argument('-n', dest='count', metavar='count', action='store',
                          type=int, default=4, help='count of things to send')
    numgroup.add_argument('-t', dest='count', action='store_const', const=None,
                          help='send indifinitely')

    # Message size or explicit message - mutually exclusive options
    sizegroup = parser.add_mutually_exclusive_group()
    sizegroup.add_argument('-l', dest='size', metavar='size', action='store',
                           type=int, help='size of auto payload')
    payload_help = "send PAYLOAD, which has the form mytype[::mybody] or ::mybody"
    sizegroup.add_argument('-payload', help=payload_help)

    # Type of command - mutually exclusive, and exactly must be specified
    cmdgroup = parser.add_mutually_exclusive_group(required=True)
    cmdgroup.add_argument('-msg', action='store_true', help='send messages')
    cmdgroup.add_argument('-cmd', action='store_true', help='send commands')
    cmdgroup.add_argument('-rtcmd', action='store_true', help='send rt-commands')

    parser.add_argument('-c', dest='channel', metavar='CHANNEL',
                        help='name of channel')
    parser.add_argument('-rate', type=float, default=1.0,
                        help='rate in sends per second')
    choices_text = "TRACE DEBUG INFO WARNING ERROR CRITICAL"
    choices_list = choices_text.split()
    loglevel_help = 'sets logging level. LOGLEVEL is one of: ' + choices_text
    parser.add_argument('-loglevel', default='ERROR', choices=choices_list,
                        help=loglevel_help)

    return parser

def parse_args(args):
    """Parse input args, including all error handling.
    Returns a parameters object on success. Exits program
    on failure.
    """
    parser = generate_argparser()
    params = parser.parse_args(args)
    try:
        hostname, port = parse_address(params.address)
        params.hostname = hostname
        params.port = port
        msgtype, msgbody = parse_payload(params.payload)
        params.msgtype = msgtype
        params.msgbody = msgbody
    except ValueError as exp:
        parser.error(str(exp))
    return params

# Note: IPV6 addresses are not supported - it would
# be extra work to distinguish a port from a ':' inside
# the IPv6 address.
_HOSTNAME_REGEX = re.compile(r'(\w|\.)+')
_PORTRANGE = 41 # Range 41000-41999
def parse_address(address):
    """Parse address of the form hostname[:port]"""
    errmsg = '\n'.join(("Invalid address '{}'".format(address),
                        "Hostname should have the form NAME_OR_IP[:PORT]"))
    hostname, *rest = address.split(':')
    if not _HOSTNAME_REGEX.fullmatch(hostname) or len(rest) > 1:
        raise ValueError(errmsg)
    port = int(rest[0]) if rest else None
    if port is not None and port//1000 != _PORTRANGE:
        msg = "Port must be in the range {} to {}"
        minport = _PORTRANGE*1000
        raise ValueError(msg.format(minport, minport+999))
    return (hostname, port)

_BAD_MSGTYPE_CHARS = string.whitespace + ','
_PAYLOAD_SEPARATOR = '::'
def parse_payload(payload):
    """Parse payload, which has the form type or ::body or type::body"""
    msgtype, msgbody = None, None
    if payload:
        errmsg = '\n'.join(("Invalid payload '{}'".format(payload),
                            "Payload should have the form TYPE or ::TYPE or TYPE::BODY"))
        # We pick the first instance of '::'
        try:
            index = payload.index(_PAYLOAD_SEPARATOR)
            msgtype = payload[:index] or None
            msgbody = payload[index+len(_PAYLOAD_SEPARATOR):] # skip past '::'
            msgbody = msgbody or None # convert '' to None
        except ValueError:
            msgtype = payload
        if msgtype and containschars(msgtype, _BAD_MSGTYPE_CHARS):
            raise ValueError(errmsg)
    return (msgtype, msgbody)

def send_messages(client, count):
    """Send messages using an instance of echo client"""
    receive_count = 0

    def send_handler(resptype, respbody):
        msg = "Sending: '{}::{}'".format(resptype, respbody)
        _TRACE(msg)
        print(msg)

    def response_handler(resptype, respbody):
        msg = "Response: '{}::{}'".format(resptype, respbody)
        _TRACE(msg)
        print(msg)
        nonlocal receive_count
        receive_count += 1 # assume call to hander is serialized

    # send_messages will block until done...
    _TRACE("GOING TO SEND MESSAGES")
    client.send_messages(count, send_handler=send_handler,
                         response_handler=response_handler)
    _TRACE("DONE SENDING MESSAGES")
    print("Received = {}".format(receive_count))

def set_loglevel(strloglevel):
    """Sets up logging with the specified logging level"""
    if strloglevel == 'TRACE':
        # Special case - not defined in logging module
        level = logging_helper.TRACELEVEL
    else:
        choices = "DEBUG INFO WARNING ERROR CRITICAL".split()
        index = choices.index(strloglevel)
        level = getattr(logging, choices[index])
    logging.basicConfig(level=level)

def main(args):
    """Main entry point"""
    params = parse_args(args)
    set_loglevel(params.loglevel)
    _LOGGER.info("parameters:\n%s", pprint.pformat(vars(params)))

    client = EchoClient('localhost')
    try:
        client.set_parameters(size=params.size, rate=params.rate,
                              bodytype=params.msgtype, body=params.msgbody)
        if params.msg:
            send_messages(client, params.count)
        elif params.cmd:
            print('send_commands(params)')
        elif params.rtcmd:
            print('send_rtcommands(params)')
    except Exception: # pylint: disable=broad-except
        _LOGGER.exception()
    finally:
        client.close()

def containschars(str_, charset) -> bool:
    """Returns if {str} contains any chars in {chars}"""
    for char in str_:
        if char in charset:
            return True
    return False

main(sys.argv[1:])
<|endoftext|>


<commit_before>#!/usr/bin/python

import yaml
import os.path

_config_file_base_name = 'mediawords.yml'
_config_file_name = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'mediawords.yml'))

def read_config():
    yml_file = open(_config_file_name, 'rb')
    config_file = yaml.load( yml_file )
    return config_file

<commit_msg>Read the defaults config file and merge it with mediawords.yml<commit_after>#!/usr/bin/python

import yaml
import os.path

_config_file_base_name = 'mediawords.yml'
_config_file_name = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'mediawords.yml'))

_defaults_config_file_base_name = 'defaults.yml'
_defaults_config_file_name = os.path.abspath(os.path.join(os.path.dirname(__file__), '../config', _defaults_config_file_base_name))

def _load_yml( file_path ):
    yml_file = open(file_path, 'rb')
    config_file = yaml.load( yml_file )

    return config_file

def deep_merge( original, update ):
    for key, value in update.iteritems():
        if not key in original:
            original[ key ] = value
        elif isinstance( value, dict) and isinstance( original[key], dict):
            deep_merge( original[ key ], value )

    return original

def read_config():

    config_file   = _load_yml( _config_file_name )
    defaults_file = _load_yml( _defaults_config_file_name )

    # print "config_file"
    # print config_file
    # print "defaults_file"
    # print defaults_file

    config_file = deep_merge( config_file, defaults_file )

    # print "Merged"
    # print config_file

    return config_file

<|endoftext|>


<commit_before>from collections import Mapping

from hestia.list_utils import to_list
from kubernetes import client

from polyaxon_k8s import constants as k8s_constants
from scheduler.spawners.templates.env_vars import get_pod_env_from, get_resources_env_vars
from scheduler.spawners.templates.gpu_volumes import get_gpu_volumes_def
from scheduler.spawners.templates.resources import get_resources
from scheduler.spawners.templates.sidecars import get_sidecar_args, get_sidecar_container
from scheduler.spawners.templates.tpu import get_tpu_annotations, requests_tpu
from schemas.exceptions import PolyaxonConfigurationError

class BaseResourceManager(object):
    def __init__(self,
                 namespace,
                 project_name,
                 project_uuid,
                 job_container_name,
                 job_docker_image,
                 job_docker_image_pull_policy,
                 sidecar_container_name,
                 sidecar_docker_image,
                 sidecar_docker_image_pull_policy,
                 init_container_name,
                 init_docker_image,
                 init_docker_image_pull_policy,
                 role_label,
                 type_label,
                 app_label,
                 health_check_url,
                 use_sidecar,
                 sidecar_config,
                 log_level):
        self.namespace = namespace
        self.project_name = project_name
        self.project_uuid = project_uuid
        self.job_container_name = job_container_name
        self.job_docker_image = job_docker_image
        self.job_docker_image_pull_policy = job_docker_image_pull_policy
        self.sidecar_container_name = sidecar_container_name
        self.sidecar_docker_image = sidecar_docker_image
        self.sidecar_docker_image_pull_policy = sidecar_docker_image_pull_policy
        self.init_container_name = init_container_name
        self.init_docker_image = init_docker_image
        self.init_docker_image_pull_policy = init_docker_image_pull_policy
        self.role_label = role_label
        self.type_label = type_label
        self.app_label = app_label
        self.use_sidecar = use_sidecar
        if use_sidecar and not sidecar_config:
            sidecar_config = {}
        if use_sidecar and not isinstance(sidecar_config, Mapping):
            raise PolyaxonConfigurationError(
                'In order to use a `sidecar_config` is required. '
                'The `sidecar_config` must correspond to the sidecar docker image used.')
        self.sidecar_config = sidecar_config
        self.health_check_url = health_check_url
        self.log_level = log_level

    def get_resource_name(self):
        raise NotImplementedError()

    def get_labels(self):
        raise NotImplementedError()

    def _pod_container_checks(self):
        pass

    def _get_logs_path(self, persistence_logs='default'):
        raise NotImplementedError()

    def _get_outputs_path(self, persistence_outputs):
        raise NotImplementedError()

    def _get_container_pod_env_vars(self,
                                    persistence_outputs,
                                    persistence_data,
                                    outputs_refs_jobs,
                                    outputs_refs_experiments,
                                    ephemeral_token):
        raise NotImplementedError()

    def get_pod_container(self,
                          volume_mounts,
                          persistence_outputs=None,
                          persistence_data=None,
                          outputs_refs_jobs=None,
                          outputs_refs_experiments=None,
                          secret_refs=None,
                          configmap_refs=None,
                          env_vars=None,
                          command=None,
                          args=None,
                          resources=None,
                          ports=None,
                          ephemeral_token=None):
        """Pod job container for task."""
        self._pod_container_checks()

        # Env vars preparations
        env_vars = to_list(env_vars, check_none=True)
        env_vars += self._get_container_pod_env_vars(
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            ephemeral_token=ephemeral_token
        )
        env_vars += get_resources_env_vars(resources=resources)

        # Env from configmap and secret refs
        env_from = get_pod_env_from(secret_refs=secret_refs, configmap_refs=configmap_refs)

        def get_ports():
            _ports = to_list(ports) if ports else []
            return [client.V1ContainerPort(container_port=port) for port in _ports] or None

        return client.V1Container(name=self.job_container_name,
                                  image=self.job_docker_image,
                                  command=command,
                                  args=args,
                                  ports=get_ports(),
                                  env=env_vars,
                                  env_from=env_from,
                                  resources=get_resources(resources),
                                  volume_mounts=volume_mounts)

    def get_sidecar_volume_mounts(self, context_mounts, persistence_outputs, persistence_data):
        return context_mounts

    def get_sidecar_container(self, resource_name, volume_mounts):
        """Pod sidecar container for task logs."""
        return get_sidecar_container(
            resource_name=resource_name,
            job_container_name=self.job_container_name,
            sidecar_container_name=self.sidecar_container_name,
            sidecar_docker_image=self.sidecar_docker_image,
            sidecar_docker_image_pull_policy=self.sidecar_docker_image_pull_policy,
            namespace=self.namespace,
            sidecar_config=self.sidecar_config,
            sidecar_args=get_sidecar_args(pod_id=resource_name,
                                          container_id=self.job_container_name,
                                          app_label=self.app_label),
            internal_health_check_url=self.health_check_url,
            volume_mounts=volume_mounts)

    def get_init_container(self,
                           init_command,
                           init_args,
                           env_vars,
                           context_mounts,
                           persistence_outputs,
                           persistence_data):
        """Pod init container for setting outputs path."""
        raise NotImplementedError()

    def get_task_pod_spec(self,
                          volume_mounts,
                          volumes,
                          resource_name,
                          persistence_outputs=None,
                          persistence_data=None,
                          outputs_refs_jobs=None,
                          outputs_refs_experiments=None,
                          env_vars=None,
                          command=None,
                          args=None,
                          init_command=None,
                          init_args=None,
                          init_env_vars=None,
                          resources=None,
                          ports=None,
                          secret_refs=None,
                          configmap_refs=None,
                          ephemeral_token=None,
                          node_selector=None,
                          affinity=None,
                          tolerations=None,
                          sidecar_context_mounts=None,
                          init_context_mounts=None,
                          restart_policy='OnFailure'):
        """Pod spec to be used to create pods for tasks: master, worker, ps."""
        sidecar_context_mounts = to_list(sidecar_context_mounts, check_none=True)
        init_context_mounts = to_list(init_context_mounts, check_none=True)
        volume_mounts = to_list(volume_mounts, check_none=True)
        volumes = to_list(volumes, check_none=True)

        gpu_volume_mounts, gpu_volumes = get_gpu_volumes_def(resources)
        volume_mounts += gpu_volume_mounts
        volumes += gpu_volumes

        pod_container = self.get_pod_container(volume_mounts=volume_mounts,
                                               persistence_outputs=persistence_outputs,
                                               persistence_data=persistence_data,
                                               outputs_refs_jobs=outputs_refs_jobs,
                                               outputs_refs_experiments=outputs_refs_experiments,
                                               secret_refs=secret_refs,
                                               configmap_refs=configmap_refs,
                                               resources=resources,
                                               env_vars=env_vars,
                                               command=command,
                                               args=args,
                                               ports=ports,
                                               ephemeral_token=ephemeral_token)

        containers = [pod_container]
        if self.use_sidecar:
            sidecar_volume_mounts = self.get_sidecar_volume_mounts(
                persistence_outputs=persistence_outputs,
                persistence_data=persistence_data,
                context_mounts=sidecar_context_mounts)
            sidecar_container = self.get_sidecar_container(resource_name=resource_name,
                                                           volume_mounts=sidecar_volume_mounts)
            containers.append(sidecar_container)

        init_container = self.get_init_container(init_command=init_command,
                                                 init_args=init_args,
                                                 env_vars=init_env_vars,
                                                 context_mounts=init_context_mounts,
                                                 persistence_outputs=persistence_outputs,
                                                 persistence_data=persistence_data)
        init_containers = to_list(init_container, check_none=True)

        node_selector = self._get_node_selector(node_selector=node_selector)
        affinity = self._get_affinity(affinity=affinity)
        tolerations = self._get_tolerations(tolerations=tolerations)
        service_account_name = self._get_service_account_name()
        return client.V1PodSpec(
            restart_policy=restart_policy,
            service_account_name=service_account_name,
            init_containers=init_containers,
            containers=containers,
            volumes=volumes,
            node_selector=node_selector,
            tolerations=tolerations,
            affinity=affinity)

    def _get_node_selector(self, node_selector):
        raise NotImplementedError()

    def _get_affinity(self, affinity):
        raise NotImplementedError()

    def _get_tolerations(self, tolerations):
        raise NotImplementedError()

    def _get_service_account_name(self):
        raise NotImplementedError()

    def get_pod(self,
                resource_name,
                volume_mounts,
                volumes,
                labels,
                env_vars=None,
                command=None,
                args=None,
                init_command=None,
                init_args=None,
                init_env_vars=None,
                ports=None,
                persistence_outputs=None,
                persistence_data=None,
                outputs_refs_jobs=None,
                outputs_refs_experiments=None,
                secret_refs=None,
                configmap_refs=None,
                resources=None,
                ephemeral_token=None,
                node_selector=None,
                affinity=None,
                tolerations=None,
                sidecar_context_mounts=None,
                init_context_mounts=None,
                restart_policy=None):
        annotations = None
        if requests_tpu(resources):
            annotations = get_tpu_annotations()
        metadata = client.V1ObjectMeta(name=resource_name,
                                       labels=labels,
                                       namespace=self.namespace,
                                       annotations=annotations)

        pod_spec = self.get_task_pod_spec(
            resource_name=resource_name,
            volume_mounts=volume_mounts,
            volumes=volumes,
            env_vars=env_vars,
            command=command,
            args=args,
            init_command=init_command,
            init_args=init_args,
            init_env_vars=init_env_vars,
            ports=ports,
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            secret_refs=secret_refs,
            configmap_refs=configmap_refs,
            resources=resources,
            ephemeral_token=ephemeral_token,
            node_selector=node_selector,
            affinity=affinity,
            tolerations=tolerations,
            init_context_mounts=init_context_mounts,
            sidecar_context_mounts=sidecar_context_mounts,
            restart_policy=restart_policy)
        return client.V1Pod(api_version=k8s_constants.K8S_API_VERSION_V1,
                            kind=k8s_constants.K8S_POD_KIND,
                            metadata=metadata,
                            spec=pod_spec)

    def get_deployment_spec(self,
                            resource_name,
                            volume_mounts,
                            volumes,
                            labels,
                            env_vars=None,
                            command=None,
                            args=None,
                            init_command=None,
                            init_args=None,
                            init_env_vars=None,
                            ports=None,
                            persistence_outputs=None,
                            persistence_data=None,
                            outputs_refs_jobs=None,
                            outputs_refs_experiments=None,
                            secret_refs=None,
                            configmap_refs=None,
                            resources=None,
                            ephemeral_token=None,
                            node_selector=None,
                            affinity=None,
                            tolerations=None,
                            restart_policy=None,
                            init_context_mounts=None,
                            sidecar_context_mounts=None,
                            replicas=1):
        metadata = client.V1ObjectMeta(name=resource_name, labels=labels, namespace=self.namespace)

        pod_spec = self.get_task_pod_spec(
            resource_name=resource_name,
            volume_mounts=volume_mounts,
            volumes=volumes,
            env_vars=env_vars,
            command=command,
            args=args,
            init_command=init_command,
            init_args=init_args,
            init_env_vars=init_env_vars,
            ports=ports,
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            secret_refs=secret_refs,
            configmap_refs=configmap_refs,
            resources=resources,
            ephemeral_token=ephemeral_token,
            node_selector=node_selector,
            affinity=affinity,
            tolerations=tolerations,
            init_context_mounts=init_context_mounts,
            sidecar_context_mounts=sidecar_context_mounts,
            restart_policy=restart_policy)
        template_spec = client.V1PodTemplateSpec(metadata=metadata, spec=pod_spec)
        return client.AppsV1beta1DeploymentSpec(replicas=replicas, template=template_spec)

    def get_deployment(self,
                       resource_name,
                       volume_mounts,
                       volumes,
                       labels,
                       env_vars=None,
                       command=None,
                       args=None,
                       init_command=None,
                       init_args=None,
                       init_env_vars=None,
                       ports=None,
                       persistence_outputs=None,
                       persistence_data=None,
                       outputs_refs_jobs=None,
                       outputs_refs_experiments=None,
                       secret_refs=None,
                       configmap_refs=None,
                       resources=None,
                       ephemeral_token=None,
                       node_selector=None,
                       affinity=None,
                       tolerations=None,
                       restart_policy=None,
                       init_context_mounts=None,
                       sidecar_context_mounts=None,
                       replicas=1):
        deployment_spec = self.get_deployment_spec(
            resource_name=resource_name,
            volume_mounts=volume_mounts,
            volumes=volumes,
            labels=labels,
            env_vars=env_vars,
            command=command,
            args=args,
            init_command=init_command,
            init_args=init_args,
            init_env_vars=init_env_vars,
            ports=ports,
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            secret_refs=secret_refs,
            configmap_refs=configmap_refs,
            resources=resources,
            ephemeral_token=ephemeral_token,
            node_selector=node_selector,
            affinity=affinity,
            tolerations=tolerations,
            restart_policy=restart_policy,
            init_context_mounts=init_context_mounts,
            sidecar_context_mounts=sidecar_context_mounts,
            replicas=replicas,
        )
        metadata = client.V1ObjectMeta(name=resource_name, labels=labels, namespace=self.namespace)
        return client.AppsV1beta1Deployment(api_version=k8s_constants.K8S_API_VERSION_V1_BETA1,
                                            kind=k8s_constants.K8S_DEPLOYMENT_KIND,
                                            metadata=metadata,
                                            spec=deployment_spec)
<commit_msg>Add missing TPU annotations for deployments. fixes #379<commit_after>from collections import Mapping

from hestia.list_utils import to_list
from kubernetes import client

from polyaxon_k8s import constants as k8s_constants
from scheduler.spawners.templates.env_vars import get_pod_env_from, get_resources_env_vars
from scheduler.spawners.templates.gpu_volumes import get_gpu_volumes_def
from scheduler.spawners.templates.resources import get_resources
from scheduler.spawners.templates.sidecars import get_sidecar_args, get_sidecar_container
from scheduler.spawners.templates.tpu import get_tpu_annotations, requests_tpu
from schemas.exceptions import PolyaxonConfigurationError

class BaseResourceManager(object):
    def __init__(self,
                 namespace,
                 project_name,
                 project_uuid,
                 job_container_name,
                 job_docker_image,
                 job_docker_image_pull_policy,
                 sidecar_container_name,
                 sidecar_docker_image,
                 sidecar_docker_image_pull_policy,
                 init_container_name,
                 init_docker_image,
                 init_docker_image_pull_policy,
                 role_label,
                 type_label,
                 app_label,
                 health_check_url,
                 use_sidecar,
                 sidecar_config,
                 log_level):
        self.namespace = namespace
        self.project_name = project_name
        self.project_uuid = project_uuid
        self.job_container_name = job_container_name
        self.job_docker_image = job_docker_image
        self.job_docker_image_pull_policy = job_docker_image_pull_policy
        self.sidecar_container_name = sidecar_container_name
        self.sidecar_docker_image = sidecar_docker_image
        self.sidecar_docker_image_pull_policy = sidecar_docker_image_pull_policy
        self.init_container_name = init_container_name
        self.init_docker_image = init_docker_image
        self.init_docker_image_pull_policy = init_docker_image_pull_policy
        self.role_label = role_label
        self.type_label = type_label
        self.app_label = app_label
        self.use_sidecar = use_sidecar
        if use_sidecar and not sidecar_config:
            sidecar_config = {}
        if use_sidecar and not isinstance(sidecar_config, Mapping):
            raise PolyaxonConfigurationError(
                'In order to use a `sidecar_config` is required. '
                'The `sidecar_config` must correspond to the sidecar docker image used.')
        self.sidecar_config = sidecar_config
        self.health_check_url = health_check_url
        self.log_level = log_level

    def get_resource_name(self):
        raise NotImplementedError()

    def get_labels(self):
        raise NotImplementedError()

    def _pod_container_checks(self):
        pass

    def _get_logs_path(self, persistence_logs='default'):
        raise NotImplementedError()

    def _get_outputs_path(self, persistence_outputs):
        raise NotImplementedError()

    def _get_container_pod_env_vars(self,
                                    persistence_outputs,
                                    persistence_data,
                                    outputs_refs_jobs,
                                    outputs_refs_experiments,
                                    ephemeral_token):
        raise NotImplementedError()

    def get_pod_container(self,
                          volume_mounts,
                          persistence_outputs=None,
                          persistence_data=None,
                          outputs_refs_jobs=None,
                          outputs_refs_experiments=None,
                          secret_refs=None,
                          configmap_refs=None,
                          env_vars=None,
                          command=None,
                          args=None,
                          resources=None,
                          ports=None,
                          ephemeral_token=None):
        """Pod job container for task."""
        self._pod_container_checks()

        # Env vars preparations
        env_vars = to_list(env_vars, check_none=True)
        env_vars += self._get_container_pod_env_vars(
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            ephemeral_token=ephemeral_token
        )
        env_vars += get_resources_env_vars(resources=resources)

        # Env from configmap and secret refs
        env_from = get_pod_env_from(secret_refs=secret_refs, configmap_refs=configmap_refs)

        def get_ports():
            _ports = to_list(ports) if ports else []
            return [client.V1ContainerPort(container_port=port) for port in _ports] or None

        return client.V1Container(name=self.job_container_name,
                                  image=self.job_docker_image,
                                  command=command,
                                  args=args,
                                  ports=get_ports(),
                                  env=env_vars,
                                  env_from=env_from,
                                  resources=get_resources(resources),
                                  volume_mounts=volume_mounts)

    def get_sidecar_volume_mounts(self, context_mounts, persistence_outputs, persistence_data):
        return context_mounts

    def get_sidecar_container(self, resource_name, volume_mounts):
        """Pod sidecar container for task logs."""
        return get_sidecar_container(
            resource_name=resource_name,
            job_container_name=self.job_container_name,
            sidecar_container_name=self.sidecar_container_name,
            sidecar_docker_image=self.sidecar_docker_image,
            sidecar_docker_image_pull_policy=self.sidecar_docker_image_pull_policy,
            namespace=self.namespace,
            sidecar_config=self.sidecar_config,
            sidecar_args=get_sidecar_args(pod_id=resource_name,
                                          container_id=self.job_container_name,
                                          app_label=self.app_label),
            internal_health_check_url=self.health_check_url,
            volume_mounts=volume_mounts)

    def get_init_container(self,
                           init_command,
                           init_args,
                           env_vars,
                           context_mounts,
                           persistence_outputs,
                           persistence_data):
        """Pod init container for setting outputs path."""
        raise NotImplementedError()

    def get_task_pod_spec(self,
                          volume_mounts,
                          volumes,
                          resource_name,
                          persistence_outputs=None,
                          persistence_data=None,
                          outputs_refs_jobs=None,
                          outputs_refs_experiments=None,
                          env_vars=None,
                          command=None,
                          args=None,
                          init_command=None,
                          init_args=None,
                          init_env_vars=None,
                          resources=None,
                          ports=None,
                          secret_refs=None,
                          configmap_refs=None,
                          ephemeral_token=None,
                          node_selector=None,
                          affinity=None,
                          tolerations=None,
                          sidecar_context_mounts=None,
                          init_context_mounts=None,
                          restart_policy='OnFailure'):
        """Pod spec to be used to create pods for tasks: master, worker, ps."""
        sidecar_context_mounts = to_list(sidecar_context_mounts, check_none=True)
        init_context_mounts = to_list(init_context_mounts, check_none=True)
        volume_mounts = to_list(volume_mounts, check_none=True)
        volumes = to_list(volumes, check_none=True)

        gpu_volume_mounts, gpu_volumes = get_gpu_volumes_def(resources)
        volume_mounts += gpu_volume_mounts
        volumes += gpu_volumes

        pod_container = self.get_pod_container(volume_mounts=volume_mounts,
                                               persistence_outputs=persistence_outputs,
                                               persistence_data=persistence_data,
                                               outputs_refs_jobs=outputs_refs_jobs,
                                               outputs_refs_experiments=outputs_refs_experiments,
                                               secret_refs=secret_refs,
                                               configmap_refs=configmap_refs,
                                               resources=resources,
                                               env_vars=env_vars,
                                               command=command,
                                               args=args,
                                               ports=ports,
                                               ephemeral_token=ephemeral_token)

        containers = [pod_container]
        if self.use_sidecar:
            sidecar_volume_mounts = self.get_sidecar_volume_mounts(
                persistence_outputs=persistence_outputs,
                persistence_data=persistence_data,
                context_mounts=sidecar_context_mounts)
            sidecar_container = self.get_sidecar_container(resource_name=resource_name,
                                                           volume_mounts=sidecar_volume_mounts)
            containers.append(sidecar_container)

        init_container = self.get_init_container(init_command=init_command,
                                                 init_args=init_args,
                                                 env_vars=init_env_vars,
                                                 context_mounts=init_context_mounts,
                                                 persistence_outputs=persistence_outputs,
                                                 persistence_data=persistence_data)
        init_containers = to_list(init_container, check_none=True)

        node_selector = self._get_node_selector(node_selector=node_selector)
        affinity = self._get_affinity(affinity=affinity)
        tolerations = self._get_tolerations(tolerations=tolerations)
        service_account_name = self._get_service_account_name()
        return client.V1PodSpec(
            restart_policy=restart_policy,
            service_account_name=service_account_name,
            init_containers=init_containers,
            containers=containers,
            volumes=volumes,
            node_selector=node_selector,
            tolerations=tolerations,
            affinity=affinity)

    def _get_node_selector(self, node_selector):
        raise NotImplementedError()

    def _get_affinity(self, affinity):
        raise NotImplementedError()

    def _get_tolerations(self, tolerations):
        raise NotImplementedError()

    def _get_service_account_name(self):
        raise NotImplementedError()

    def get_pod(self,
                resource_name,
                volume_mounts,
                volumes,
                labels,
                env_vars=None,
                command=None,
                args=None,
                init_command=None,
                init_args=None,
                init_env_vars=None,
                ports=None,
                persistence_outputs=None,
                persistence_data=None,
                outputs_refs_jobs=None,
                outputs_refs_experiments=None,
                secret_refs=None,
                configmap_refs=None,
                resources=None,
                ephemeral_token=None,
                node_selector=None,
                affinity=None,
                tolerations=None,
                sidecar_context_mounts=None,
                init_context_mounts=None,
                restart_policy=None):
        annotations = None
        if requests_tpu(resources):
            annotations = get_tpu_annotations()
        metadata = client.V1ObjectMeta(name=resource_name,
                                       labels=labels,
                                       namespace=self.namespace,
                                       annotations=annotations)

        pod_spec = self.get_task_pod_spec(
            resource_name=resource_name,
            volume_mounts=volume_mounts,
            volumes=volumes,
            env_vars=env_vars,
            command=command,
            args=args,
            init_command=init_command,
            init_args=init_args,
            init_env_vars=init_env_vars,
            ports=ports,
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            secret_refs=secret_refs,
            configmap_refs=configmap_refs,
            resources=resources,
            ephemeral_token=ephemeral_token,
            node_selector=node_selector,
            affinity=affinity,
            tolerations=tolerations,
            init_context_mounts=init_context_mounts,
            sidecar_context_mounts=sidecar_context_mounts,
            restart_policy=restart_policy)
        return client.V1Pod(api_version=k8s_constants.K8S_API_VERSION_V1,
                            kind=k8s_constants.K8S_POD_KIND,
                            metadata=metadata,
                            spec=pod_spec)

    def get_deployment_spec(self,
                            resource_name,
                            volume_mounts,
                            volumes,
                            labels,
                            env_vars=None,
                            command=None,
                            args=None,
                            init_command=None,
                            init_args=None,
                            init_env_vars=None,
                            ports=None,
                            persistence_outputs=None,
                            persistence_data=None,
                            outputs_refs_jobs=None,
                            outputs_refs_experiments=None,
                            secret_refs=None,
                            configmap_refs=None,
                            resources=None,
                            ephemeral_token=None,
                            node_selector=None,
                            affinity=None,
                            tolerations=None,
                            restart_policy=None,
                            init_context_mounts=None,
                            sidecar_context_mounts=None,
                            replicas=1):
        annotations = None
        if requests_tpu(resources):
            annotations = get_tpu_annotations()
        metadata = client.V1ObjectMeta(name=resource_name,
                                       labels=labels,
                                       namespace=self.namespace,
                                       annotations=annotations)

        pod_spec = self.get_task_pod_spec(
            resource_name=resource_name,
            volume_mounts=volume_mounts,
            volumes=volumes,
            env_vars=env_vars,
            command=command,
            args=args,
            init_command=init_command,
            init_args=init_args,
            init_env_vars=init_env_vars,
            ports=ports,
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            secret_refs=secret_refs,
            configmap_refs=configmap_refs,
            resources=resources,
            ephemeral_token=ephemeral_token,
            node_selector=node_selector,
            affinity=affinity,
            tolerations=tolerations,
            init_context_mounts=init_context_mounts,
            sidecar_context_mounts=sidecar_context_mounts,
            restart_policy=restart_policy)
        template_spec = client.V1PodTemplateSpec(metadata=metadata, spec=pod_spec)
        return client.AppsV1beta1DeploymentSpec(replicas=replicas, template=template_spec)

    def get_deployment(self,
                       resource_name,
                       volume_mounts,
                       volumes,
                       labels,
                       env_vars=None,
                       command=None,
                       args=None,
                       init_command=None,
                       init_args=None,
                       init_env_vars=None,
                       ports=None,
                       persistence_outputs=None,
                       persistence_data=None,
                       outputs_refs_jobs=None,
                       outputs_refs_experiments=None,
                       secret_refs=None,
                       configmap_refs=None,
                       resources=None,
                       ephemeral_token=None,
                       node_selector=None,
                       affinity=None,
                       tolerations=None,
                       restart_policy=None,
                       init_context_mounts=None,
                       sidecar_context_mounts=None,
                       replicas=1):
        deployment_spec = self.get_deployment_spec(
            resource_name=resource_name,
            volume_mounts=volume_mounts,
            volumes=volumes,
            labels=labels,
            env_vars=env_vars,
            command=command,
            args=args,
            init_command=init_command,
            init_args=init_args,
            init_env_vars=init_env_vars,
            ports=ports,
            persistence_outputs=persistence_outputs,
            persistence_data=persistence_data,
            outputs_refs_jobs=outputs_refs_jobs,
            outputs_refs_experiments=outputs_refs_experiments,
            secret_refs=secret_refs,
            configmap_refs=configmap_refs,
            resources=resources,
            ephemeral_token=ephemeral_token,
            node_selector=node_selector,
            affinity=affinity,
            tolerations=tolerations,
            restart_policy=restart_policy,
            init_context_mounts=init_context_mounts,
            sidecar_context_mounts=sidecar_context_mounts,
            replicas=replicas,
        )
        metadata = client.V1ObjectMeta(name=resource_name, labels=labels, namespace=self.namespace)
        return client.AppsV1beta1Deployment(api_version=k8s_constants.K8S_API_VERSION_V1_BETA1,
                                            kind=k8s_constants.K8S_DEPLOYMENT_KIND,
                                            metadata=metadata,
                                            spec=deployment_spec)
<|endoftext|>