<commit_before>#!/usr/bin/env python
"""Tests for HDF5 stuff"""
import tempfile
from os.path import join, dirname

import numpy as np
import tables as tb

from km3pipe import Blob, Module, Pipeline, Pump, version
from km3pipe.dataclasses import Table
from km3pipe.io.hdf5 import HDF5Pump, HDF5Sink, FORMAT_VERSION    # noqa
from km3pipe.tools import insert_prefix_to_dtype
from km3pipe.testing import TestCase

DATA_DIR = join(dirname(__file__), '../../kp-data/test_data/')

class TestMultiTable(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ],
                            dtype=[
                                ('a', '<f8'),
                                ('b', '<f8'),
                                ('c', '<f8'),
                            ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ],
                            dtype=[
                                ('aa', '<f8'),
                                ('bb', '<f8'),
                                ('cc', '<f8'),
                            ])
        self.tabs = {'foo': self.foo, 'bar': self.bar}
        self.where = '/lala'
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.h5name = self.fobj.name
        self.h5file = tb.open_file(
        # create the file in memory only
            self.h5name,
            'w',
            driver="H5FD_CORE",
            driver_core_backing_store=0
        )
        for name, tab in self.tabs.items():
            self.h5file.create_table(
                self.where, name=name, obj=tab, createparents=True
            )

    def tearDown(self):
        self.h5file.close()
        self.fobj.close()

    def test_name_insert(self):
        exp_foo = ('foo_a', 'foo_b', 'foo_c')
        exp_bar = ('bar_aa', 'bar_bb', 'bar_cc')
        pref_foo = insert_prefix_to_dtype(self.tabs['foo'], 'foo')
        pref_bar = insert_prefix_to_dtype(self.tabs['bar'], 'bar')
        self.assertEqual(exp_foo, pref_foo.dtype.names)
        self.assertEqual(exp_bar, pref_bar.dtype.names)

    # def test_group_read(self):
    #    tabs = read_group(self.h5file.root)
    #    exp_cols = (
    #        'bar_aa', 'bar_bb', 'bar_cc',
    #        'foo_a', 'foo_b', 'foo_c',
    #    )
    #    exp_shape = (2, 6)
    #    res_shape = tabs.shape
    #    res_cols = tuple(tabs.columns)
    #    print(exp_cols)
    #    print(res_cols)
    #    self.assertEqual(exp_shape, res_shape)
    #    self.assertEqual(exp_cols, res_cols)

class TestH5Pump(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR, 'numu_cc_test.h5')

    def test_init_sets_filename_if_no_keyword_arg_is_passed(self):
        p = HDF5Pump(self.fname)
        self.assertEqual(self.fname, p.filename)
        p.finish()

    def test_context(self):
        with HDF5Pump(self.fname) as h5:
            self.assertEqual(self.fname, h5.filename)
            assert h5[0] is not None
            for blob in h5:
                assert blob is not None
                break

    def test_standalone(self):
        pump = HDF5Pump(filename=self.fname)
        next(pump)
        pump.finish()

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.drain()

    def test_event_info_is_not_empty(self):
        self.fname = join(DATA_DIR, 'test_event_info.h5')

        class Printer(Module):
            def process(self, blob):
                assert blob['EventInfo'].size != 0
                return blob

        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(Printer)
        p.drain()

    def test_event_info_has_correct_group_id(self):
        self.fname = join(DATA_DIR, 'test_event_info.h5')

        class Printer(Module):
            def configure(self):
                self.index = 0

            def process(self, blob):
                assert blob['EventInfo'][0].group_id == self.index
                self.index += 1
                return blob

        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(Printer)
        p.drain()

class TestH5Sink(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR, 'numu_cc_test.h5')
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.out = tb.open_file(
            self.fobj.name,
            "w",
            driver="H5FD_CORE",
            driver_core_backing_store=0
        )

    def tearDown(self):
        self.out.close()
        self.fobj.close()

    # def test_init_has_to_be_explicit(self):
    #     with self.assertRaises(TypeError):
    #         HDF5Sink(self.out)

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(HDF5Sink, h5file=self.out)
        p.drain()

    def test_scalars(self):
        out = tb.open_file(
            'foobar_scalar',
            "a",
            driver="H5FD_CORE",
            driver_core_backing_store=0
        )

        def pu(blob):
            return {'foo': 42.0}

        p = Pipeline()
        p.attach(pu)
        p.attach(HDF5Sink, h5file=out, keep_open=True)
        p.drain(2)
        node = out.root.misc
        assert node is not None
        assert node.cols is not None
        assert 'foo' in set(node.cols._v_colnames)
        out.close()

    def test_h5info(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def process(self, blob):
                return Blob()

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.open_file(fname, 'r') as h5file:
            assert version == h5file.root._v_attrs.km3pipe.decode()
            assert tb.__version__ == h5file.root._v_attrs.pytables.decode()
            assert FORMAT_VERSION == h5file.root._v_attrs.format_version
            assert 'None' == h5file.root._v_attrs.jpp.decode()

        fobj.close()

class TestH5SinkConsistency(TestCase):
    def test_h5_consistency_for_tables_without_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 10
                tab = Table({'a': self.count, 'b': 1}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            a = f.get_node("/tab")[:]['a']
            b = f.get_node("/tab")[:]['b']
            group_id = f.get_node("/tab")[:]['group_id']
        assert np.allclose([10, 20, 30, 40, 50], a)
        assert np.allclose([1, 1, 1, 1, 1], b)
        assert np.allclose([0, 1, 2, 3, 4], group_id)
        fobj.close()

    def test_h5_consistency_for_tables_without_group_id_and_multiple_keys(
            self
    ):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 10
                tab1 = Table({'a': self.count, 'b': 1}, h5loc='tab1')
                tab2 = Table({'c': self.count + 1, 'd': 2}, h5loc='tab2')
                return Blob({'tab1': tab1, 'tab2': tab2})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            a = f.get_node("/tab1")[:]['a']
            b = f.get_node("/tab1")[:]['b']
            c = f.get_node("/tab2")[:]['c']
            d = f.get_node("/tab2")[:]['d']
            group_id_1 = f.get_node("/tab1")[:]['group_id']
            group_id_2 = f.get_node("/tab1")[:]['group_id']
        assert np.allclose([10, 20, 30, 40, 50], a)
        assert np.allclose([1, 1, 1, 1, 1], b)
        assert np.allclose([0, 1, 2, 3, 4], group_id_1)
        assert np.allclose([11, 21, 31, 41, 51], c)
        assert np.allclose([2, 2, 2, 2, 2], d)
        assert np.allclose([0, 1, 2, 3, 4], group_id_2)
        fobj.close()

    def test_h5_consistency_for_tables_with_custom_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def process(self, blob):
                tab = Table({'group_id': 2}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            group_id = f.get_node("/tab")[:]['group_id']

        assert np.allclose([2, 2, 2, 2, 2], group_id)

        fobj.close()

class TestHDF5PumpConsistency(TestCase):
    def test_hdf5_readout(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 1
                tab = Table({'a': self.count * 10, 'b': 1}, h5loc='tab')
                tab2 = Table({'a': np.arange(self.count)}, h5loc='tab2')
                blob['Tab'] = tab
                blob['Tab2'] = tab2
                return blob

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        class BlobTester(Module):
            def configure(self):
                self.index = 0

            def process(self, blob):
                self.index += 1
                assert 'GroupInfo' in blob
                assert 'Tab' in blob
                print(self.index)
                print(blob['Tab'])
                print(blob['Tab']['a'])
                assert self.index - 1 == blob['GroupInfo'].group_id
                assert self.index * 10 == blob['Tab']['a']
                assert 1 == blob['Tab']['b'] == 1
                assert np.allclose(np.arange(self.index), blob['Tab2']['a'])
                return blob

        pipe = Pipeline()
        pipe.attach(HDF5Pump, filename=fname)
        pipe.attach(BlobTester)
        pipe.drain()

        fobj.close()

    def test_hdf5_readout_split_tables(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 1
                tab = Table({
                    'a': self.count * 10,
                    'b': 1
                },
                            h5loc='/tab',
                            split_h5=True)
                blob['Tab'] = tab
                return blob

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        class BlobTester(Module):
            def configure(self):
                self.index = 0

            def process(self, blob):
                self.index += 1
                assert 'GroupInfo' in blob
                assert 'Tab' in blob
                assert self.index - 1 == blob['GroupInfo'].group_id
                assert self.index * 10 == blob['Tab']['a']
                assert 1 == blob['Tab']['b'] == 1
                return blob

        pipe = Pipeline()
        pipe.attach(HDF5Pump, filename=fname)
        pipe.attach(BlobTester)
        pipe.drain()

        fobj.close()
<commit_msg>Don't check for jpp version<commit_after>#!/usr/bin/env python
"""Tests for HDF5 stuff"""
import tempfile
from os.path import join, dirname

import numpy as np
import tables as tb

from km3pipe import Blob, Module, Pipeline, Pump, version
from km3pipe.dataclasses import Table
from km3pipe.io.hdf5 import HDF5Pump, HDF5Sink, FORMAT_VERSION    # noqa
from km3pipe.tools import insert_prefix_to_dtype
from km3pipe.testing import TestCase

DATA_DIR = join(dirname(__file__), '../../kp-data/test_data/')

class TestMultiTable(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ],
                            dtype=[
                                ('a', '<f8'),
                                ('b', '<f8'),
                                ('c', '<f8'),
                            ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ],
                            dtype=[
                                ('aa', '<f8'),
                                ('bb', '<f8'),
                                ('cc', '<f8'),
                            ])
        self.tabs = {'foo': self.foo, 'bar': self.bar}
        self.where = '/lala'
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.h5name = self.fobj.name
        self.h5file = tb.open_file(
        # create the file in memory only
            self.h5name,
            'w',
            driver="H5FD_CORE",
            driver_core_backing_store=0
        )
        for name, tab in self.tabs.items():
            self.h5file.create_table(
                self.where, name=name, obj=tab, createparents=True
            )

    def tearDown(self):
        self.h5file.close()
        self.fobj.close()

    def test_name_insert(self):
        exp_foo = ('foo_a', 'foo_b', 'foo_c')
        exp_bar = ('bar_aa', 'bar_bb', 'bar_cc')
        pref_foo = insert_prefix_to_dtype(self.tabs['foo'], 'foo')
        pref_bar = insert_prefix_to_dtype(self.tabs['bar'], 'bar')
        self.assertEqual(exp_foo, pref_foo.dtype.names)
        self.assertEqual(exp_bar, pref_bar.dtype.names)

    # def test_group_read(self):
    #    tabs = read_group(self.h5file.root)
    #    exp_cols = (
    #        'bar_aa', 'bar_bb', 'bar_cc',
    #        'foo_a', 'foo_b', 'foo_c',
    #    )
    #    exp_shape = (2, 6)
    #    res_shape = tabs.shape
    #    res_cols = tuple(tabs.columns)
    #    print(exp_cols)
    #    print(res_cols)
    #    self.assertEqual(exp_shape, res_shape)
    #    self.assertEqual(exp_cols, res_cols)

class TestH5Pump(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR, 'numu_cc_test.h5')

    def test_init_sets_filename_if_no_keyword_arg_is_passed(self):
        p = HDF5Pump(self.fname)
        self.assertEqual(self.fname, p.filename)
        p.finish()

    def test_context(self):
        with HDF5Pump(self.fname) as h5:
            self.assertEqual(self.fname, h5.filename)
            assert h5[0] is not None
            for blob in h5:
                assert blob is not None
                break

    def test_standalone(self):
        pump = HDF5Pump(filename=self.fname)
        next(pump)
        pump.finish()

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.drain()

    def test_event_info_is_not_empty(self):
        self.fname = join(DATA_DIR, 'test_event_info.h5')

        class Printer(Module):
            def process(self, blob):
                assert blob['EventInfo'].size != 0
                return blob

        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(Printer)
        p.drain()

    def test_event_info_has_correct_group_id(self):
        self.fname = join(DATA_DIR, 'test_event_info.h5')

        class Printer(Module):
            def configure(self):
                self.index = 0

            def process(self, blob):
                assert blob['EventInfo'][0].group_id == self.index
                self.index += 1
                return blob

        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(Printer)
        p.drain()

class TestH5Sink(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR, 'numu_cc_test.h5')
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.out = tb.open_file(
            self.fobj.name,
            "w",
            driver="H5FD_CORE",
            driver_core_backing_store=0
        )

    def tearDown(self):
        self.out.close()
        self.fobj.close()

    # def test_init_has_to_be_explicit(self):
    #     with self.assertRaises(TypeError):
    #         HDF5Sink(self.out)

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(HDF5Sink, h5file=self.out)
        p.drain()

    def test_scalars(self):
        out = tb.open_file(
            'foobar_scalar',
            "a",
            driver="H5FD_CORE",
            driver_core_backing_store=0
        )

        def pu(blob):
            return {'foo': 42.0}

        p = Pipeline()
        p.attach(pu)
        p.attach(HDF5Sink, h5file=out, keep_open=True)
        p.drain(2)
        node = out.root.misc
        assert node is not None
        assert node.cols is not None
        assert 'foo' in set(node.cols._v_colnames)
        out.close()

    def test_h5info(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def process(self, blob):
                return Blob()

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.open_file(fname, 'r') as h5file:
            assert version == h5file.root._v_attrs.km3pipe.decode()
            assert tb.__version__ == h5file.root._v_attrs.pytables.decode()
            assert FORMAT_VERSION == h5file.root._v_attrs.format_version

        fobj.close()

class TestH5SinkConsistency(TestCase):
    def test_h5_consistency_for_tables_without_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 10
                tab = Table({'a': self.count, 'b': 1}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            a = f.get_node("/tab")[:]['a']
            b = f.get_node("/tab")[:]['b']
            group_id = f.get_node("/tab")[:]['group_id']
        assert np.allclose([10, 20, 30, 40, 50], a)
        assert np.allclose([1, 1, 1, 1, 1], b)
        assert np.allclose([0, 1, 2, 3, 4], group_id)
        fobj.close()

    def test_h5_consistency_for_tables_without_group_id_and_multiple_keys(
            self
    ):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 10
                tab1 = Table({'a': self.count, 'b': 1}, h5loc='tab1')
                tab2 = Table({'c': self.count + 1, 'd': 2}, h5loc='tab2')
                return Blob({'tab1': tab1, 'tab2': tab2})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            a = f.get_node("/tab1")[:]['a']
            b = f.get_node("/tab1")[:]['b']
            c = f.get_node("/tab2")[:]['c']
            d = f.get_node("/tab2")[:]['d']
            group_id_1 = f.get_node("/tab1")[:]['group_id']
            group_id_2 = f.get_node("/tab1")[:]['group_id']
        assert np.allclose([10, 20, 30, 40, 50], a)
        assert np.allclose([1, 1, 1, 1, 1], b)
        assert np.allclose([0, 1, 2, 3, 4], group_id_1)
        assert np.allclose([11, 21, 31, 41, 51], c)
        assert np.allclose([2, 2, 2, 2, 2], d)
        assert np.allclose([0, 1, 2, 3, 4], group_id_2)
        fobj.close()

    def test_h5_consistency_for_tables_with_custom_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def process(self, blob):
                tab = Table({'group_id': 2}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            group_id = f.get_node("/tab")[:]['group_id']

        assert np.allclose([2, 2, 2, 2, 2], group_id)

        fobj.close()

class TestHDF5PumpConsistency(TestCase):
    def test_hdf5_readout(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 1
                tab = Table({'a': self.count * 10, 'b': 1}, h5loc='tab')
                tab2 = Table({'a': np.arange(self.count)}, h5loc='tab2')
                blob['Tab'] = tab
                blob['Tab2'] = tab2
                return blob

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        class BlobTester(Module):
            def configure(self):
                self.index = 0

            def process(self, blob):
                self.index += 1
                assert 'GroupInfo' in blob
                assert 'Tab' in blob
                print(self.index)
                print(blob['Tab'])
                print(blob['Tab']['a'])
                assert self.index - 1 == blob['GroupInfo'].group_id
                assert self.index * 10 == blob['Tab']['a']
                assert 1 == blob['Tab']['b'] == 1
                assert np.allclose(np.arange(self.index), blob['Tab2']['a'])
                return blob

        pipe = Pipeline()
        pipe.attach(HDF5Pump, filename=fname)
        pipe.attach(BlobTester)
        pipe.drain()

        fobj.close()

    def test_hdf5_readout_split_tables(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 1
                tab = Table({
                    'a': self.count * 10,
                    'b': 1
                },
                            h5loc='/tab',
                            split_h5=True)
                blob['Tab'] = tab
                return blob

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        class BlobTester(Module):
            def configure(self):
                self.index = 0

            def process(self, blob):
                self.index += 1
                assert 'GroupInfo' in blob
                assert 'Tab' in blob
                assert self.index - 1 == blob['GroupInfo'].group_id
                assert self.index * 10 == blob['Tab']['a']
                assert 1 == blob['Tab']['b'] == 1
                return blob

        pipe = Pipeline()
        pipe.attach(HDF5Pump, filename=fname)
        pipe.attach(BlobTester)
        pipe.drain()

        fobj.close()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
VERSION = (1, 0, 1, 'final', 0)

# taken from django-registration

def get_version():
    "Returns a PEP 386-compliant version number from VERSION."
    assert len(VERSION) == 5
    assert VERSION[3] in ('alpha', 'beta', 'rc', 'final')

    # Now build the two parts of the version number:
    # main = X.Y[.Z]
    # sub = .devN - for pre-alpha releases
    #     | {a|b|c}N - for alpha, beta and rc releases

    parts = 2 if VERSION[2] == 0 else 3
    main = '.'.join(str(x) for x in VERSION[:parts])

    sub = ''
    if VERSION[3] != 'final':
        mapping = {'alpha': 'a', 'beta': 'b', 'rc': 'c'}
        sub = mapping[VERSION[3]] + str(VERSION[4])

    return str(main + sub)
<commit_msg>Dev version update<commit_after># -*- coding: utf-8 -*-
VERSION = (1, 1, 0, 'alpha', 0)

# taken from django-registration

def get_version():
    "Returns a PEP 386-compliant version number from VERSION."
    assert len(VERSION) == 5
    assert VERSION[3] in ('alpha', 'beta', 'rc', 'final')

    # Now build the two parts of the version number:
    # main = X.Y[.Z]
    # sub = .devN - for pre-alpha releases
    #     | {a|b|c}N - for alpha, beta and rc releases

    parts = 2 if VERSION[2] == 0 else 3
    main = '.'.join(str(x) for x in VERSION[:parts])

    sub = ''
    if VERSION[3] != 'final':
        mapping = {'alpha': 'a', 'beta': 'b', 'rc': 'c'}
        sub = mapping[VERSION[3]] + str(VERSION[4])

    return str(main + sub)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import print_function, division, absolute_import, unicode_literals
from .rendering import NodeVisitor

class GrammarVisitor(NodeVisitor):
    def visit_grammar(self, g, vrules):
        return g

    def visit_rule(self, r, vexp):
        return r

    def visit_rule_ref(self, rr):
        return rr

    def visit_special(self, s):
        return s

    def visit_override(self, o, vexp):
        return o

    def visit_named(self, n, vexp):
        return n

    def visit_cut(self, c):
        return c

    def visit_optional(self, o, vexp):
        return o

    def visit_repeat(self, r, vexp):
        return r

    def visit_repeat_plus(self, r, vexp):
        return r

    def visit_group(self, g, vexp):
        return g

    def visit_choice(self, c, vopt):
        return c

    def visit_sequence(self, s, vseq):
        return s

    def visit_lookahead(self, l, vexp):
        return l

    def visit_lookahead_not(self, l, vexp):
        return l

    def visit_pattern(self, p):
        return p

    def visit_token(self, t):
        return t

    def visit_void(self, v):
        return v

    def visit_fail(self, f):
        return f

    def visit_eof(self, e):
        return e

    def visit_Grammar(self, g):
        return self.visit_grammar(g, [r.accept(self) for r in g.rules])

    def visit_Rule(self, r):
        return self.visit_rule(r, r.exp.accept(self))

    def visit_RuleRef(self, r):
        return self.visit_rule_ref(r)

    def visit_Special(self, s):
        return self.visit_special(s)

    def visit_Override(self, o):
        return self.visit_override(o, o.exp.accept(self))

    def visit_Named(self, n):
        return self.visit_named(n, n.exp.accept(self))

    def visit_Cut(self, c):
        return self.visit_cut(c)

    def visit_Optional(self, o):
        return self.visit_optional(o, o.exp.accept(self))

    def visit_Repeat(self, r):
        return self.visit_repeat(r, r.exp.accept(self))

    def visit_RepeatPlus(self, r):
        return self.visit_repeat_plus(r, r.exp.accept(self))

    def visit_Group(self, g):
        return self.visit_group(g, g.exp.accept(self))

    def visit_Choice(self, c):
        return self.visit_choice(c, [o.accept(self) for o in c.options])

    def visit_Sequence(self, s):
        return self.visit_sequence(s, [e.accept(self) for e in s.sequence])

    def visit_Lookahead(self, l):
        return self.visit_lookahead(l, l.exp.accept(self))

    def visit_LookaheadNot(self, l):
        return self.visit_lookahead_not(l, l.exp.accept(self))

    def visit_Pattern(self, p):
        return self.visit_pattern(p)

    def visit_Token(self, t):
        return self.visit_token(t)

    def visit_Void(self, v):
        return self.visit_void(v)

    def visit_Fail(self, f):
        return self.visit_fail(f)

    def visit_EOF(self, e):
        return self.visit_eof(e)
<commit_msg>Removed the obsolete visitor module (again).<commit_after><|endoftext|>


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

import os

import numpy as np
import tables as tb

from km3pipe import Pipeline
from km3pipe.io import _insert_tabname_into_colnames, _read_group
from km3pipe.io.hdf5 import HDF5Pump, HDF5Sink
from km3pipe.testing import TestCase

class TestMultiTable(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ], dtype=[('a', '<f8'), ('b', '<f8'), ('c', '<f8'), ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ], dtype=[('aa', '<f8'), ('bb', '<f8'), ('cc', '<f8'), ])
        self.tabs = {'foo': self.foo, 'bar': self.bar}
        self.where = '/lala'
        self.h5name = './test.h5'
        self.h5file = tb.open_file(
            # create the file in memory only
            self.h5name, 'w', driver="H5FD_CORE", driver_core_backing_store=0)
        for name, tab in self.tabs.items():
            self.h5file.create_table(self.where, name=name, obj=tab,
                                     createparents=True)

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

    def test_name_insert(self):
        exp_foo = ('foo_a', 'foo_b', 'foo_c')
        exp_bar = ('bar_aa', 'bar_bb', 'bar_cc')
        pref_foo = _insert_tabname_into_colnames(self.tabs['foo'], 'foo')
        pref_bar = _insert_tabname_into_colnames(self.tabs['bar'], 'bar')
        self.assertEqual(exp_foo, pref_foo.dtype.names)
        self.assertEqual(exp_bar, pref_bar.dtype.names)

    def test_group_read(self):
        tabs = _read_group(self.h5file, self.where)
        exp_cols = (
            'bar_aa', 'bar_bb', 'bar_cc',
            'foo_a', 'foo_b', 'foo_c',
        )
        exp_shape = (2, 6)
        res_shape = tabs.shape
        res_cols = tuple(tabs.columns)
        print(exp_cols)
        print(res_cols)
        self.assertEqual(exp_shape, res_shape)
        self.assertEqual(exp_cols, res_cols)

class TestWriteRead(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ], dtype=[('a', '<f8'), ('b', '<f8'), ('c', '<f8'), ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ], dtype=[('aa', '<f8'), ('bb', '<f8'), ('cc', '<f8'), ])
        self.tabs = {'foo': self.foo, 'bar': self.bar}
        self.where = '/lala'
        self.h5name = './test.h5'
        self.h5file = tb.open_file(
            # create the file in memory only
            self.h5name, 'w', driver="H5FD_CORE", driver_core_backing_store=0)
        for name, tab in self.tabs.items():
            self.h5file.create_table(self.where, name=name, obj=tab,
                                     createparents=True)

    def tearDown(self):
        self.h5file.close()
<commit_msg>add test env<commit_after>#!/usr/bin/env python

import os

import numpy as np
import tables as tb

from km3pipe.io import _insert_prefix_to_dtype, _read_group
from km3pipe.io.hdf5 import H5Chain
from km3pipe.testing import TestCase

class TestMultiTable(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ], dtype=[('a', '<f8'), ('b', '<f8'), ('c', '<f8'), ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ], dtype=[('aa', '<f8'), ('bb', '<f8'), ('cc', '<f8'), ])
        self.tabs = {'foo': self.foo, 'bar': self.bar}
        self.where = '/lala'
        self.h5name = './test.h5'
        self.h5file = tb.open_file(
            # create the file in memory only
            self.h5name, 'w', driver="H5FD_CORE", driver_core_backing_store=0)
        for name, tab in self.tabs.items():
            self.h5file.create_table(self.where, name=name, obj=tab,
                                     createparents=True)

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

    def test_name_insert(self):
        exp_foo = ('foo_a', 'foo_b', 'foo_c')
        exp_bar = ('bar_aa', 'bar_bb', 'bar_cc')
        pref_foo = _insert_prefix_to_dtype(self.tabs['foo'], 'foo')
        pref_bar = _insert_prefix_to_dtype(self.tabs['bar'], 'bar')
        self.assertEqual(exp_foo, pref_foo.dtype.names)
        self.assertEqual(exp_bar, pref_bar.dtype.names)

    def test_group_read(self):
        tabs = _read_group(self.h5file, self.where)
        exp_cols = (
            'bar_aa', 'bar_bb', 'bar_cc',
            'foo_a', 'foo_b', 'foo_c',
        )
        exp_shape = (2, 6)
        res_shape = tabs.shape
        res_cols = tuple(tabs.columns)
        print(exp_cols)
        print(res_cols)
        self.assertEqual(exp_shape, res_shape)
        self.assertEqual(exp_cols, res_cols)

class TestH5Chain(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ], dtype=[('a', '<f8'), ('b', '<f8'), ('c', '<f8'), ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ], dtype=[('aa', '<f8'), ('bb', '<f8'), ('cc', '<f8'), ])
        self.yay = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ], dtype=[('aaa', '<f8'), ('bbb', '<f8'), ('ccc', '<f8'), ])
        self.tabs = {'foo': self.foo, 'bar': self.bar, 'yay': self.yay}
        self.where = {'foo': '/', 'bar': '/lala', 'yay': '/lala'}
        self.h5name = './test.h5'
        self.h5file = tb.open_file(self.h5name, 'a')
        for name, tab in self.tabs.items():
            self.h5file.create_table(self.where, name=name, obj=tab,
                                     createparents=True)
        self.h5file.close()

    # def tearDown(self):
    #     os.remove(self.h5name)

    def test_single_file(self):
        files = {self.h5name: 1}
        c = H5Chain(files)
        print(c['foo'])
        self.assertTrue(False)
<|endoftext|>


<commit_before>import time
from .block import Block
from .instance import BlockchainInstance
from bitsharesbase.operationids import getOperationNameForId

class Blockchain(object):
    """ This class allows to access the blockchain and read data
        from it

        :param bitshares.bitshares.BitShares blockchain_instance: BitShares
                 instance
        :param str mode: (default) Irreversible block (``irreversible``) or
                 actual head block (``head``)
        :param int max_block_wait_repetition: (default) 3 maximum wait time for
            next block ismax_block_wait_repetition * block_interval

        This class let's you deal with blockchain related data and methods.
    """
    def __init__(
        self,
        mode="irreversible",
        max_block_wait_repetition=3,
        *args,
        **kwargs
    ):
        BlockchainInstance.__init__(self, *args, **kwargs)
        assert mode in ["irreversible", "head"], "invalid value for 'mode'!"
        self.mode = mode
        self.max_block_wait_repetition = max_block_wait_repetition

    def is_irreversible_mode(self):
        return self.mode == 'last_irreversible_block_num'

    def info(self):
        """ This call returns the *dynamic global properties*
        """
        return self.blockchain.rpc.get_dynamic_global_properties()

    def chainParameters(self):
        """ The blockchain parameters, such as fees, and committee-controlled
            parameters are returned here
        """
        return self.config()["parameters"]

    def get_network(self):
        """ Identify the network

            :returns: Network parameters
            :rtype: dict
        """
        return self.blockchain.rpc.get_network()

    def get_chain_properties(self):
        """ Return chain properties
        """
        return self.blockchain.rpc.get_chain_properties()

    def config(self):
        """ Returns object 2.0.0
        """
        return self.blockchain.rpc.get_object("2.0.0")

    def get_current_block_num(self):
        """ This call returns the current block

            .. note:: The block number returned depends on the ``mode`` used
                      when instanciating from this class.
        """
        return self.info().get(self.mode)

    def get_current_block(self):
        """ This call returns the current block

            .. note:: The block number returned depends on the ``mode`` used
                      when instanciating from this class.
        """
        return Block(
            self.get_current_block_num(),
            blockchain_instance=self.blockchain
        )

    def block_time(self, block_num):
        """ Returns a datetime of the block with the given block
            number.

            :param int block_num: Block number
        """
        return Block(
            block_num,
            blockchain_instance=self.blockchain
        ).time()

    def block_timestamp(self, block_num):
        """ Returns the timestamp of the block with the given block
            number.

            :param int block_num: Block number
        """
        return int(Block(
            block_num,
            blockchain_instance=self.blockchain
        ).time().timestamp())

    def blocks(self, start=None, stop=None):
        """ Yields blocks starting from ``start``.

            :param int start: Starting block
            :param int stop: Stop at this block
            :param str mode: We here have the choice between
             "head" (the last block) and "irreversible" (the block that is
             confirmed by 2/3 of all block producers and is thus irreversible)
        """
        # Let's find out how often blocks are generated!
        self.block_interval = self.chainParameters().get("block_interval")

        if not start:
            start = self.get_current_block_num()

        # We are going to loop indefinitely
        while True:

            # Get chain properies to identify the
            if stop:
                head_block = stop
            else:
                head_block = self.get_current_block_num()

            # Blocks from start until head block
            for blocknum in range(start, head_block + 1):
                # Get full block
                block = self.wait_for_and_get_block(blocknum)
                block.update({"block_num": blocknum})
                yield block
            # Set new start
            start = head_block + 1

            if stop and start > stop:
                # raise StopIteration
                return

            # Sleep for one block
            time.sleep(self.block_interval)

    def wait_for_and_get_block(self, block_number, blocks_waiting_for=None):
        """ Get the desired block from the chain, if the current head block is
            smaller (for both head and irreversible) then we wait, but a
            maxmimum of blocks_waiting_for * max_block_wait_repetition time
            before failure.

            :param int block_number: desired block number
            :param int blocks_waiting_for: (default) difference between
                block_number and current head how many blocks we are willing to
                wait, positive int
        """
        if not blocks_waiting_for:
            blocks_waiting_for = max(
                1, block_number - self.get_current_block_num())

        repetition = 0
        # can't return the block before the chain has reached it (support
        # future block_num)
        while self.get_current_block_num() < block_number:
            repetition += 1
            time.sleep(self.block_interval)
            if (
                repetition >
                blocks_waiting_for * self.max_block_wait_repetition
            ):
                raise Exception("Wait time for new block exceeded, aborting")
        # block has to be returned properly
        block = self.blockchain.rpc.get_block(block_number)
        repetition = 0
        while not block:
            repetition += 1
            time.sleep(self.block_interval)
            if repetition > self.max_block_wait_repetition:
                raise Exception("Wait time for new block exceeded, aborting")
            block = self.blockchain.rpc.get_block(block_number)
        return block

    def ops(self, start=None, stop=None, **kwargs):
        """ Yields all operations (including virtual operations) starting from
            ``start``.

            :param int start: Starting block
            :param int stop: Stop at this block
            :param str mode: We here have the choice between
             "head" (the last block) and "irreversible" (the block that is
             confirmed by 2/3 of all block producers and is thus irreversible)
            :param bool only_virtual_ops: Only yield virtual operations

            This call returns a list that only carries one operation and
            its type!
        """

        for block in self.blocks(start=start, stop=stop, **kwargs):
            for tx in block["transactions"]:
                for op in tx["operations"]:
                    # Replace opid by op name
                    op[0] = getOperationNameForId(op[0])
                    yield {
                        "block_num": block["block_num"],
                        "op": op,
                        "timestamp": block["timestamp"]
                    }

    def stream(self, opNames=[], *args, **kwargs):
        """ Yield specific operations (e.g. comments) only

            :param array opNames: List of operations to filter for
            :param int start: Start at this block
            :param int stop: Stop at this block
            :param str mode: We here have the choice between
                 * "head": the last block
                 * "irreversible": the block that is confirmed by 2/3 of all
                    block producers and is thus irreversible!

            The dict output is formated such that ``type`` caries the
            operation type, timestamp and block_num are taken from the
            block the operation was stored in and the other key depend
            on the actualy operation.
        """
        for op in self.ops(**kwargs):
            if not opNames or op["op"][0] in opNames:
                r = {
                    "type": op["op"][0],
                    "timestamp": op.get("timestamp"),
                    "block_num": op.get("block_num"),
                }
                r.update(op["op"][1])
                yield r

    def awaitTxConfirmation(self, transaction, limit=10):
        """ Returns the transaction as seen by the blockchain after being
            included into a block

            .. note:: If you want instant confirmation, you need to instantiate
                      class:`bitshares.blockchain.Blockchain` with
                      ``mode="head"``, otherwise, the call will wait until
                      confirmed in an irreversible block.

            .. note:: This method returns once the blockchain has included a
                      transaction with the **same signature**. Even though the
                      signature is not usually used to identify a transaction,
                      it still cannot be forfeited and is derived from the
                      transaction contented and thus identifies a transaction
                      uniquely.
        """
        counter = 10
        for block in self.blocks():
            counter += 1
            for tx in block["transactions"]:
                if (
                    sorted(tx["signatures"]) ==
                    sorted(transaction["signatures"])
                ):
                    return tx
            if counter > limit:
                raise Exception(
                    "The operation has not been added after 10 blocks!")

    def get_all_accounts(self, start='', stop='', steps=1e3, **kwargs):
        """ Yields account names between start and stop.

            :param str start: Start at this account name
            :param str stop: Stop at this account name
            :param int steps: Obtain ``steps`` ret with a single call from RPC
        """
        lastname = start
        while True:
            ret = self.blockchain.rpc.lookup_accounts(lastname, steps)
            for account in ret:
                yield account[0]
                if account[0] == stop:
                    raise StopIteration
            if lastname == ret[-1][0]:
                raise StopIteration
            lastname = ret[-1][0]
            if len(ret) < steps:
                raise StopIteration
<commit_msg>revert changes to blockchain.py<commit_after>import time
from .block import Block
from .instance import BlockchainInstance
from bitsharesbase.operationids import getOperationNameForId

class Blockchain(object):
    """ This class allows to access the blockchain and read data
        from it

        :param bitshares.bitshares.BitShares blockchain_instance: BitShares
                 instance
        :param str mode: (default) Irreversible block (``irreversible``) or
                 actual head block (``head``)
        :param int max_block_wait_repetition: (default) 3 maximum wait time for
            next block ismax_block_wait_repetition * block_interval

        This class let's you deal with blockchain related data and methods.
    """
    def __init__(
        self,
        blockchain_instance=None,
        mode="irreversible",
        max_block_wait_repetition=None,
        *args,
        **kwargs
    ):
        BlockchainInstance.__init__(self, *args, **kwargs)

        if mode == "irreversible":
            self.mode = 'last_irreversible_block_num'
        elif mode == "head":
            self.mode = "head_block_number"
        else:
            raise ValueError("invalid value for 'mode'!")

        if max_block_wait_repetition:
            self.max_block_wait_repetition = max_block_wait_repetition
        else:
            self.max_block_wait_repetition = 3

    def is_irreversible_mode(self):
        return self.mode == 'last_irreversible_block_num'

    def info(self):
        """ This call returns the *dynamic global properties*
        """
        return self.blockchain.rpc.get_dynamic_global_properties()

    def chainParameters(self):
        """ The blockchain parameters, such as fees, and committee-controlled
            parameters are returned here
        """
        return self.config()["parameters"]

    def get_network(self):
        """ Identify the network

            :returns: Network parameters
            :rtype: dict
        """
        return self.blockchain.rpc.get_network()

    def get_chain_properties(self):
        """ Return chain properties
        """
        return self.blockchain.rpc.get_chain_properties()

    def config(self):
        """ Returns object 2.0.0
        """
        return self.blockchain.rpc.get_object("2.0.0")

    def get_current_block_num(self):
        """ This call returns the current block

            .. note:: The block number returned depends on the ``mode`` used
                      when instanciating from this class.
        """
        return self.info().get(self.mode)

    def get_current_block(self):
        """ This call returns the current block

            .. note:: The block number returned depends on the ``mode`` used
                      when instanciating from this class.
        """
        return Block(
            self.get_current_block_num(),
            blockchain_instance=self.blockchain
        )

    def block_time(self, block_num):
        """ Returns a datetime of the block with the given block
            number.

            :param int block_num: Block number
        """
        return Block(
            block_num,
            blockchain_instance=self.blockchain
        ).time()

    def block_timestamp(self, block_num):
        """ Returns the timestamp of the block with the given block
            number.

            :param int block_num: Block number
        """
        return int(Block(
            block_num,
            blockchain_instance=self.blockchain
        ).time().timestamp())

    def blocks(self, start=None, stop=None):
        """ Yields blocks starting from ``start``.

            :param int start: Starting block
            :param int stop: Stop at this block
            :param str mode: We here have the choice between
             "head" (the last block) and "irreversible" (the block that is
             confirmed by 2/3 of all block producers and is thus irreversible)
        """
        # Let's find out how often blocks are generated!
        self.block_interval = self.chainParameters().get("block_interval")

        if not start:
            start = self.get_current_block_num()

        # We are going to loop indefinitely
        while True:

            # Get chain properies to identify the
            if stop:
                head_block = stop
            else:
                head_block = self.get_current_block_num()

            # Blocks from start until head block
            for blocknum in range(start, head_block + 1):
                # Get full block
                block = self.wait_for_and_get_block(blocknum)
                block.update({"block_num": blocknum})
                yield block
            # Set new start
            start = head_block + 1

            if stop and start > stop:
                # raise StopIteration
                return

            # Sleep for one block
            time.sleep(self.block_interval)

    def wait_for_and_get_block(self, block_number, blocks_waiting_for=None):
        """ Get the desired block from the chain, if the current head block is
            smaller (for both head and irreversible) then we wait, but a
            maxmimum of blocks_waiting_for * max_block_wait_repetition time
            before failure.

            :param int block_number: desired block number
            :param int blocks_waiting_for: (default) difference between
                block_number and current head how many blocks we are willing to
                wait, positive int
        """
        if not blocks_waiting_for:
            blocks_waiting_for = max(
                1, block_number - self.get_current_block_num())

        repetition = 0
        # can't return the block before the chain has reached it (support
        # future block_num)
        while self.get_current_block_num() < block_number:
            repetition += 1
            time.sleep(self.block_interval)
            if (
                repetition >
                blocks_waiting_for * self.max_block_wait_repetition
            ):
                raise Exception("Wait time for new block exceeded, aborting")
        # block has to be returned properly
        block = self.blockchain.rpc.get_block(block_number)
        repetition = 0
        while not block:
            repetition += 1
            time.sleep(self.block_interval)
            if repetition > self.max_block_wait_repetition:
                raise Exception("Wait time for new block exceeded, aborting")
            block = self.blockchain.rpc.get_block(block_number)
        return block

    def ops(self, start=None, stop=None, **kwargs):
        """ Yields all operations (including virtual operations) starting from
            ``start``.

            :param int start: Starting block
            :param int stop: Stop at this block
            :param str mode: We here have the choice between
             "head" (the last block) and "irreversible" (the block that is
             confirmed by 2/3 of all block producers and is thus irreversible)
            :param bool only_virtual_ops: Only yield virtual operations

            This call returns a list that only carries one operation and
            its type!
        """

        for block in self.blocks(start=start, stop=stop, **kwargs):
            for tx in block["transactions"]:
                for op in tx["operations"]:
                    # Replace opid by op name
                    op[0] = getOperationNameForId(op[0])
                    yield {
                        "block_num": block["block_num"],
                        "op": op,
                        "timestamp": block["timestamp"]
                    }

    def stream(self, opNames=[], *args, **kwargs):
        """ Yield specific operations (e.g. comments) only

            :param array opNames: List of operations to filter for
            :param int start: Start at this block
            :param int stop: Stop at this block
            :param str mode: We here have the choice between
                 * "head": the last block
                 * "irreversible": the block that is confirmed by 2/3 of all
                    block producers and is thus irreversible!

            The dict output is formated such that ``type`` caries the
            operation type, timestamp and block_num are taken from the
            block the operation was stored in and the other key depend
            on the actualy operation.
        """
        for op in self.ops(**kwargs):
            if not opNames or op["op"][0] in opNames:
                r = {
                    "type": op["op"][0],
                    "timestamp": op.get("timestamp"),
                    "block_num": op.get("block_num"),
                }
                r.update(op["op"][1])
                yield r

    def awaitTxConfirmation(self, transaction, limit=10):
        """ Returns the transaction as seen by the blockchain after being
            included into a block

            .. note:: If you want instant confirmation, you need to instantiate
                      class:`bitshares.blockchain.Blockchain` with
                      ``mode="head"``, otherwise, the call will wait until
                      confirmed in an irreversible block.

            .. note:: This method returns once the blockchain has included a
                      transaction with the **same signature**. Even though the
                      signature is not usually used to identify a transaction,
                      it still cannot be forfeited and is derived from the
                      transaction contented and thus identifies a transaction
                      uniquely.
        """
        counter = 10
        for block in self.blocks():
            counter += 1
            for tx in block["transactions"]:
                if (
                    sorted(tx["signatures"]) ==
                    sorted(transaction["signatures"])
                ):
                    return tx
            if counter > limit:
                raise Exception(
                    "The operation has not been added after 10 blocks!")

    def get_all_accounts(self, start='', stop='', steps=1e3, **kwargs):
        """ Yields account names between start and stop.

            :param str start: Start at this account name
            :param str stop: Stop at this account name
            :param int steps: Obtain ``steps`` ret with a single call from RPC
        """
        lastname = start
        while True:
            ret = self.blockchain.rpc.lookup_accounts(lastname, steps)
            for account in ret:
                yield account[0]
                if account[0] == stop:
                    raise StopIteration
            if lastname == ret[-1][0]:
                raise StopIteration
            lastname = ret[-1][0]
            if len(ret) < steps:
                raise StopIteration
<|endoftext|>


<commit_before>import csv
from functools import update_wrapper
from itertools import count

from django.conf.urls import url
from django.contrib import admin
from django.contrib import messages
from django.core.exceptions import ValidationError, PermissionDenied
from django.core.urlresolvers import reverse
from django.db import transaction
from django.http import HttpResponse
from django import forms
from django.utils import six
from django.utils.translation import ugettext_lazy as _
from django.views.generic.edit import FormView

class CSVImportError(ValidationError):
    def __init__(self, *args, **kwargs):
        self.rownumber = kwargs.pop('rownumber')
        super(CSVImportError, self).__init__(*args, **kwargs)

class ImportCSVForm(forms.Form):
    csv_file = forms.FileField(required=True, label=_('CSV File'))
    has_headers = forms.BooleanField(
        label=_('Has headers'),
        help_text=_('Check this whether or not your CSV file '
                    'has a row with columns headers.'),
        initial=True,
        required=False,
    )

class ImportCSVAdminView(FormView):
    form_class = ImportCSVForm
    model_admin = None

    def _get_meta(self):
        opts = self.model_admin.model._meta
        app_label = opts.app_label
        object_name = opts.object_name.lower()
        return (app_label, object_name)

    def get_template_names(self):
        importcsv_template = self.model_admin.importcsv_template
        if importcsv_template is not None:
            return importcsv_template
        else:
            app_label, object_name = self._get_meta()
            return [
                'admin/%s/%s/csv_import.html' % (app_label, object_name),
                'admin/%s/csv_import.html' % app_label,
                'admin/csv_import.html',
            ]

    def get_success_url(self):
        app_label, object_name = self._get_meta()
        return reverse('admin:%s_%s_changelist' % (app_label, object_name))

    def form_valid(self, form):
        try:
            self.import_csv(form.cleaned_data['csv_file'], form.cleaned_data['has_headers'])
        except CSVImportError as e:
            importer = self.model_admin.importer_class()
            for field, errors in six.iteritems(e.message_dict):
                label = importer[field].label
                for error in errors:
                    messages.error(self.request,
                        _("Couldn't import row number #%d: %s - %s") % (e.rownumber,
                                                                        label,
                                                                        error)
                    )
            return self.form_invalid(form)
        return super(ImportCSVAdminView, self).form_valid(form)

    def get_context_data(self, **kwargs):
        context = super(ImportCSVAdminView, self).get_context_data(**kwargs)
        context['opts'] = self.model_admin.model._meta
        return context

    @transaction.commit_on_success
    def import_csv(self, file_, has_headers):
        reader = csv.DictReader(
            file_,
            fieldnames=self.model_admin.importer_class._meta.fields,
            dialect=self.model_admin.dialect,
        )

        reader_iter = iter(six.moves.zip(count(1), reader))
        if has_headers:
            six.advance_iterator(reader_iter)

        for i, row in reader_iter:
            self.process_row(i, row)

    def process_row(self, i, row):
        importer = self.model_admin.importer_class(data=row)
        if not importer.is_valid():
            raise CSVImportError(importer.errors, rownumber=i)
        return importer.save()

class ImportCSVModelAdmin(admin.ModelAdmin):
    importcsv_view_class = ImportCSVAdminView
    importcsv_template = None

    dialect = csv.excel
    skip_firstline = True

    @property
    def change_list_template(self):
        opts = self.model._meta
        return [
            'admin/%s/%s/change_list_csv.html' % (opts.app_label, opts.object_name.lower()),
            'admin/%s/change_list_csv.html' % opts.app_label,
            'admin/change_list_csv.html',
        ]

    def get_urls(self):
        def wrap(view):
            def wrapper(request, *args, **kwargs):
                if not self.has_add_permission(request):
                    raise PermissionDenied
                return self.admin_site.admin_view(view)(request, *args, **kwargs)
            return update_wrapper(wrapper, view)

        info = self.model._meta.app_label, self.model._meta.module_name

        urlpatterns = super(ImportCSVModelAdmin, self).get_urls()

        extra_urls = [
            url(r'^import-csv/$',
                wrap(self.importcsv_view),
                name='%s_%s_importcsv' % info),
            url(r'^import-csv/template.csv$',
                wrap(self.download_csv_template),
                name='%s_%s_csvtemplate' % info),
        ]
        return extra_urls + urlpatterns

    @property
    def importcsv_view(self):
        return self.importcsv_view_class.as_view(model_admin=self)

    def download_csv_template(self, request):
        def get_label(form, fname):
            field = form[fname]
            label = field.label
            if field.field.required:
                label = '%s*' % label
            return label

        importer = self.importer_class()

        response = HttpResponse(content_type='text/csv')
        response['Content-Disposition'] = 'attachment; filename="template.csv"'

        writer = csv.writer(response, dialect=self.dialect)

        fields = importer._meta.fields
        labels = [get_label(importer, fname) for fname in fields]
        writer.writerow(labels)

        return response
<commit_msg>FIX: Use built-in enumerate() function<commit_after>import csv
from functools import update_wrapper
from itertools import count

from django.conf.urls import url
from django.contrib import admin
from django.contrib import messages
from django.core.exceptions import ValidationError, PermissionDenied
from django.core.urlresolvers import reverse
from django.db import transaction
from django.http import HttpResponse
from django import forms
from django.utils import six
from django.utils.translation import ugettext_lazy as _
from django.views.generic.edit import FormView

class CSVImportError(ValidationError):
    def __init__(self, *args, **kwargs):
        self.rownumber = kwargs.pop('rownumber')
        super(CSVImportError, self).__init__(*args, **kwargs)

class ImportCSVForm(forms.Form):
    csv_file = forms.FileField(required=True, label=_('CSV File'))
    has_headers = forms.BooleanField(
        label=_('Has headers'),
        help_text=_('Check this whether or not your CSV file '
                    'has a row with columns headers.'),
        initial=True,
        required=False,
    )

class ImportCSVAdminView(FormView):
    form_class = ImportCSVForm
    model_admin = None

    def _get_meta(self):
        opts = self.model_admin.model._meta
        app_label = opts.app_label
        object_name = opts.object_name.lower()
        return (app_label, object_name)

    def get_template_names(self):
        importcsv_template = self.model_admin.importcsv_template
        if importcsv_template is not None:
            return importcsv_template
        else:
            app_label, object_name = self._get_meta()
            return [
                'admin/%s/%s/csv_import.html' % (app_label, object_name),
                'admin/%s/csv_import.html' % app_label,
                'admin/csv_import.html',
            ]

    def get_success_url(self):
        app_label, object_name = self._get_meta()
        return reverse('admin:%s_%s_changelist' % (app_label, object_name))

    def form_valid(self, form):
        try:
            self.import_csv(form.cleaned_data['csv_file'], form.cleaned_data['has_headers'])
        except CSVImportError as e:
            importer = self.model_admin.importer_class()
            for field, errors in six.iteritems(e.message_dict):
                label = importer[field].label
                for error in errors:
                    messages.error(self.request,
                        _("Couldn't import row number #%d: %s - %s") % (e.rownumber,
                                                                        label,
                                                                        error)
                    )
            return self.form_invalid(form)
        return super(ImportCSVAdminView, self).form_valid(form)

    def get_context_data(self, **kwargs):
        context = super(ImportCSVAdminView, self).get_context_data(**kwargs)
        context['opts'] = self.model_admin.model._meta
        return context

    @transaction.commit_on_success
    def import_csv(self, file_, has_headers):
        reader = csv.DictReader(
            file_,
            fieldnames=self.model_admin.importer_class._meta.fields,
            dialect=self.model_admin.dialect,
        )

        reader_iter = enumerate(reader, 1)
        if has_headers:
            six.advance_iterator(reader_iter)

        for i, row in reader_iter:
            self.process_row(i, row)

    def process_row(self, i, row):
        importer = self.model_admin.importer_class(data=row)
        if not importer.is_valid():
            raise CSVImportError(importer.errors, rownumber=i)
        return importer.save()

class ImportCSVModelAdmin(admin.ModelAdmin):
    importcsv_view_class = ImportCSVAdminView
    importcsv_template = None

    dialect = csv.excel
    skip_firstline = True

    @property
    def change_list_template(self):
        opts = self.model._meta
        return [
            'admin/%s/%s/change_list_csv.html' % (opts.app_label, opts.object_name.lower()),
            'admin/%s/change_list_csv.html' % opts.app_label,
            'admin/change_list_csv.html',
        ]

    def get_urls(self):
        def wrap(view):
            def wrapper(request, *args, **kwargs):
                if not self.has_add_permission(request):
                    raise PermissionDenied
                return self.admin_site.admin_view(view)(request, *args, **kwargs)
            return update_wrapper(wrapper, view)

        info = self.model._meta.app_label, self.model._meta.module_name

        urlpatterns = super(ImportCSVModelAdmin, self).get_urls()

        extra_urls = [
            url(r'^import-csv/$',
                wrap(self.importcsv_view),
                name='%s_%s_importcsv' % info),
            url(r'^import-csv/template.csv$',
                wrap(self.download_csv_template),
                name='%s_%s_csvtemplate' % info),
        ]
        return extra_urls + urlpatterns

    @property
    def importcsv_view(self):
        return self.importcsv_view_class.as_view(model_admin=self)

    def download_csv_template(self, request):
        def get_label(form, fname):
            field = form[fname]
            label = field.label
            if field.field.required:
                label = '%s*' % label
            return label

        importer = self.importer_class()

        response = HttpResponse(content_type='text/csv')
        response['Content-Disposition'] = 'attachment; filename="template.csv"'

        writer = csv.writer(response, dialect=self.dialect)

        fields = importer._meta.fields
        labels = [get_label(importer, fname) for fname in fields]
        writer.writerow(labels)

        return response
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding=utf-8
# Filename: triggermap.py
# vim: ts=4 sw=4 et
"""
This script creates histogram which shows the trigger contribution for events.

Usage:
    trigger_map.py [-d DU -p PLOT_FILENAME] FILENAME
    trigger_map.py --version

Option:
    FILENAME          Name of the input file.
    -u DU             Only plot for the given DU.
    -d DET_ID         Detector ID [default: 29].
    -p PLOT_FILENAME  The filename of the plot [default: trigger_map.png].
    -h --help         Show this screen.

"""
from __future__ import division

__author__ = "Tamas Gal"
__copyright__ = "Copyright 2016, Tamas Gal and the KM3NeT collaboration."
__credits__ = []
__license__ = "MIT"
__maintainer__ = "Tamas Gal"
__email__ = "tgal@km3net.de"
__status__ = "Development"

from docopt import docopt
import matplotlib
# Force matplotlib to not use any Xwindows backend.
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as md
from matplotlib.colors import LogNorm
from matplotlib import pylab
import numpy as np

import km3pipe as kp
import km3pipe.style
km3pipe.style.use('km3pipe')

class TriggerMap(kp.Module):
    """Creates a plot to show the number of triggered hits for each DOM."""
    def configure(self):
        self.det = self.require('detector')
        self.plot_filename = self.require('plot_filename')
        self.subtitle = self.get('subtitle', default='')
        self.du = self.get('du')
        if self.du is not None:
            self.n_dus = 1
            self.n_doms = 18
        else:
            self.n_dus = self.det.n_dus
            self.n_doms = int(self.det.n_doms / self.n_dus)
        self.hit_counts = []

    def process(self, blob):
        hits = blob['Hits'].triggered_hits
        dom_ids = np.unique(hits.dom_id)
        hit_counts = np.zeros(self.n_dus * self.n_doms)
        for dom_id in dom_ids:
            n_hits = np.sum(hits.dom_id == dom_id)
            du, floor, _ = self.det.doms[dom_id]
            if self.du is not None and du != self.du:
                continue
            if self.du:
                du = 1
            hit_counts[(du - 1) * self.n_doms + floor - 1] += n_hits
        self.hit_counts.append(hit_counts)

        return blob

    def finish(self):
        self.create_plot()

    def create_plot(self):
        title = "Trigger Map\n{}".format(self.subtitle)
        fig, ax = plt.subplots(figsize=(16, 8))
        ax.grid(True)
        ax.set_axisbelow(True)
        hit_mat = np.array([np.array(x) for x in self.hit_counts]).transpose()
        im = ax.matshow(hit_mat,
                        interpolation='nearest', filternorm=None, cmap='plasma',
                        aspect='auto', origin='lower', zorder=3,
                        norm=LogNorm(vmin=1, vmax=np.amax(hit_mat)))
        yticks = np.arange(self.n_doms * self.n_dus)
        ytick_label_templ = "DU{0:.0f}-DOM{1:02d}" if self.du else "DOM{1:02d}"
        ytick_labels = [ytick_label_templ
                        .format(np.ceil((y+1)/self.n_doms), y % (self.n_doms) + 1)  \
                        for y in yticks]
        ax.set_yticks(yticks)
        ax.set_yticklabels(ytick_labels)
        ax.tick_params(labelbottom=True)
        ax.tick_params(labeltop=True)
        ax.set_xlabel("event id")
        ax.set_title(title)
        cb = fig.colorbar(im, pad=0.05)
        cb.set_label("number of triggered hits")

        fig.tight_layout()

        plt.savefig(self.plot_filename, dpi=120, bbox_inches="tight")

def main():
    args = docopt(__doc__, version='1.0')
    du = int(args['-u']) if args['-d'] else None
    det_id = int(args['-d'])
    det = kp.hardware.Detector(det_id=det_id)
    pipe = kp.Pipeline()
    pipe.attach(kp.io.jpp.EventPump, filename=args['FILENAME'])
    pipe.attach(TriggerMap,
                detector=det,
                du=du,
                plot_filename=args['-p'],
                subtitle=args['FILENAME'])
    pipe.drain()

if __name__ == '__main__':
    main()
<commit_msg>Fix signature in doc<commit_after>#!/usr/bin/env python
# coding=utf-8
# Filename: triggermap.py
# vim: ts=4 sw=4 et
"""
This script creates histogram which shows the trigger contribution for events.

Usage:
    triggermap [-d DET_ID -p PLOT_FILENAME -u DU] FILENAME
    triggermap --version

Option:
    FILENAME          Name of the input file.
    -u DU             Only plot for the given DU.
    -d DET_ID         Detector ID [default: 29].
    -p PLOT_FILENAME  The filename of the plot [default: trigger_map.png].
    -h --help         Show this screen.

"""
from __future__ import division

__author__ = "Tamas Gal"
__copyright__ = "Copyright 2016, Tamas Gal and the KM3NeT collaboration."
__credits__ = []
__license__ = "MIT"
__maintainer__ = "Tamas Gal"
__email__ = "tgal@km3net.de"
__status__ = "Development"

from docopt import docopt
import matplotlib
# Force matplotlib to not use any Xwindows backend.
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as md
from matplotlib.colors import LogNorm
from matplotlib import pylab
import numpy as np

import km3pipe as kp
import km3pipe.style
km3pipe.style.use('km3pipe')

class TriggerMap(kp.Module):
    """Creates a plot to show the number of triggered hits for each DOM."""
    def configure(self):
        self.det = self.require('detector')
        self.plot_filename = self.require('plot_filename')
        self.subtitle = self.get('subtitle', default='')
        self.du = self.get('du')
        if self.du is not None:
            self.n_dus = 1
            self.n_doms = 18
        else:
            self.n_dus = self.det.n_dus
            self.n_doms = int(self.det.n_doms / self.n_dus)
        self.hit_counts = []

    def process(self, blob):
        hits = blob['Hits'].triggered_hits
        dom_ids = np.unique(hits.dom_id)
        hit_counts = np.zeros(self.n_dus * self.n_doms)
        for dom_id in dom_ids:
            n_hits = np.sum(hits.dom_id == dom_id)
            du, floor, _ = self.det.doms[dom_id]
            if self.du is not None and du != self.du:
                continue
            if self.du:
                du = 1
            hit_counts[(du - 1) * self.n_doms + floor - 1] += n_hits
        self.hit_counts.append(hit_counts)

        return blob

    def finish(self):
        self.create_plot()

    def create_plot(self):
        title = "Trigger Map\n{}".format(self.subtitle)
        fig, ax = plt.subplots(figsize=(16, 8))
        ax.grid(True)
        ax.set_axisbelow(True)
        hit_mat = np.array([np.array(x) for x in self.hit_counts]).transpose()
        im = ax.matshow(hit_mat,
                        interpolation='nearest', filternorm=None, cmap='plasma',
                        aspect='auto', origin='lower', zorder=3,
                        norm=LogNorm(vmin=1, vmax=np.amax(hit_mat)))
        yticks = np.arange(self.n_doms * self.n_dus)
        ytick_label_templ = "DU{0:.0f}-DOM{1:02d}" if self.du else "DOM{1:02d}"
        ytick_labels = [ytick_label_templ
                        .format(np.ceil((y+1)/self.n_doms), y % (self.n_doms) + 1)  \
                        for y in yticks]
        ax.set_yticks(yticks)
        ax.set_yticklabels(ytick_labels)
        ax.tick_params(labelbottom=True)
        ax.tick_params(labeltop=True)
        ax.set_xlabel("event id")
        ax.set_title(title)
        cb = fig.colorbar(im, pad=0.05)
        cb.set_label("number of triggered hits")

        fig.tight_layout()

        plt.savefig(self.plot_filename, dpi=120, bbox_inches="tight")

def main():
    args = docopt(__doc__, version='1.0')
    du = int(args['-u']) if args['-d'] else None
    det_id = int(args['-d'])
    det = kp.hardware.Detector(det_id=det_id)
    pipe = kp.Pipeline()
    pipe.attach(kp.io.jpp.EventPump, filename=args['FILENAME'])
    pipe.attach(TriggerMap,
                detector=det,
                du=du,
                plot_filename=args['-p'],
                subtitle=args['FILENAME'])
    pipe.drain()

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


<commit_before># (c) 2012-2014, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

# Make coding more python3-ish
from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

from collections import defaultdict
from six import iteritems

from ansible.compat.tests import unittest
from ansible.compat.tests.mock import patch, MagicMock
from ansible.inventory import Inventory
from ansible.playbook.play import Play
from ansible.vars import VariableManager

from units.mock.loader import DictDataLoader
from units.mock.path import mock_unfrackpath_noop

class TestVariableManager(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_basic_manager(self):
        fake_loader = DictDataLoader({})

        v = VariableManager()
        vars = v.get_vars(loader=fake_loader, use_cache=False)
        if 'omit' in vars:
            del vars['omit']
        if 'vars' in vars:
            del vars['vars']
        if 'ansible_version' in vars:
            del vars['ansible_version']
        if 'ansible_check_mode' in vars:
            del vars['ansible_check_mode']

        self.assertEqual(vars, dict(playbook_dir='.'))

    def test_variable_manager_extra_vars(self):
        fake_loader = DictDataLoader({})

        extra_vars = dict(a=1, b=2, c=3)
        v = VariableManager()
        v.extra_vars = extra_vars

        vars = v.get_vars(loader=fake_loader, use_cache=False)

        for (key, val) in iteritems(extra_vars):
            self.assertEqual(vars.get(key), val)

        self.assertIsNot(v.extra_vars, extra_vars)

    def test_variable_manager_host_vars_file(self):
        fake_loader = DictDataLoader({
            "host_vars/hostname1.yml": """
               foo: bar
            """,
            "other_path/host_vars/hostname1.yml": """
               foo: bam
               baa: bat
            """,
            "host_vars/host.name.yml": """
               host_with_dots: true
            """,
        })

        v = VariableManager()
        v.add_host_vars_file("host_vars/hostname1.yml", loader=fake_loader)
        v.add_host_vars_file("other_path/host_vars/hostname1.yml", loader=fake_loader)
        self.assertIn("hostname1", v._host_vars_files)
        self.assertEqual(v._host_vars_files["hostname1"], [dict(foo="bar"), dict(foo="bam", baa="bat")])

        mock_host = MagicMock()
        mock_host.get_name.return_value = "hostname1"
        mock_host.get_vars.return_value = dict()
        mock_host.get_groups.return_value = ()
        mock_host.get_group_vars.return_value = dict()

        self.assertEqual(v.get_vars(loader=fake_loader, host=mock_host, use_cache=False).get("foo"), "bam")
        self.assertEqual(v.get_vars(loader=fake_loader, host=mock_host, use_cache=False).get("baa"), "bat")

        v.add_host_vars_file("host_vars/host.name", loader=fake_loader)
        self.assertEqual(v._host_vars_files["host.name"], [dict(host_with_dots=True)])

    def test_variable_manager_group_vars_file(self):
        fake_loader = DictDataLoader({
            "group_vars/all.yml": """
               foo: bar
            """,
            "group_vars/somegroup.yml": """
               bam: baz
            """,
            "other_path/group_vars/somegroup.yml": """
               baa: bat
            """,
            "group_vars/some.group.yml": """
               group_with_dots: true
            """,
        })

        v = VariableManager()
        v.add_group_vars_file("group_vars/all.yml", loader=fake_loader)
        v.add_group_vars_file("group_vars/somegroup.yml", loader=fake_loader)
        v.add_group_vars_file("other_path/group_vars/somegroup.yml", loader=fake_loader)
        self.assertIn("somegroup", v._group_vars_files)
        self.assertEqual(v._group_vars_files["all"], [dict(foo="bar")])
        self.assertEqual(v._group_vars_files["somegroup"], [dict(bam="baz"), dict(baa="bat")])

        mock_group = MagicMock()
        mock_group.name = "somegroup"
        mock_group.get_ancestors.return_value = ()
        mock_group.get_vars.return_value = dict()

        mock_host = MagicMock()
        mock_host.get_name.return_value = "hostname1"
        mock_host.get_vars.return_value = dict()
        mock_host.get_groups.return_value = (mock_group,)
        mock_host.get_group_vars.return_value = dict()

        vars = v.get_vars(loader=fake_loader, host=mock_host, use_cache=False)
        self.assertEqual(vars.get("foo"), "bar")
        self.assertEqual(vars.get("baa"), "bat")

        v.add_group_vars_file("group_vars/some.group", loader=fake_loader)
        self.assertEqual(v._group_vars_files["some.group"], [dict(group_with_dots=True)])

    def test_variable_manager_play_vars(self):
        fake_loader = DictDataLoader({})

        mock_play = MagicMock()
        mock_play.get_vars.return_value = dict(foo="bar")
        mock_play.get_roles.return_value = []
        mock_play.get_vars_files.return_value = []

        v = VariableManager()
        self.assertEqual(v.get_vars(loader=fake_loader, play=mock_play, use_cache=False).get("foo"), "bar")

    def test_variable_manager_play_vars_files(self):
        fake_loader = DictDataLoader({
            "/path/to/somefile.yml": """
               foo: bar
            """
        })

        mock_play = MagicMock()
        mock_play.get_vars.return_value = dict()
        mock_play.get_roles.return_value = []
        mock_play.get_vars_files.return_value = ['/path/to/somefile.yml']

        v = VariableManager()
        self.assertEqual(v.get_vars(loader=fake_loader, play=mock_play, use_cache=False).get("foo"), "bar")

    def test_variable_manager_task_vars(self):
        fake_loader = DictDataLoader({})

        mock_task = MagicMock()
        mock_task._role = None
        mock_task.loop = None
        mock_task.get_vars.return_value = dict(foo="bar")
        mock_task.get_include_params.return_value = dict()

        v = VariableManager()
        self.assertEqual(v.get_vars(loader=fake_loader, task=mock_task, use_cache=False).get("foo"), "bar")

    @patch.object(Inventory, 'basedir')
    @patch('ansible.playbook.role.definition.unfrackpath', mock_unfrackpath_noop)
    def test_variable_manager_precedence(self, mock_basedir):
        '''
        Tests complex variations and combinations of get_vars() with different
        objects to modify the context under which variables are merged.
        '''

        v = VariableManager()
        v._fact_cache = defaultdict(dict)

        fake_loader = DictDataLoader({
            # inventory1
            '/etc/ansible/inventory1': """
            [group2:children]
            group1

            [group1]
            host1 host_var=host_var_from_inventory_host1

            [group1:vars]
            group_var = group_var_from_inventory_group1

            [group2:vars]
            group_var = group_var_from_inventory_group2
            """,

            # role defaults_only1
            '/etc/ansible/roles/defaults_only1/defaults/main.yml': """
            default_var: "default_var_from_defaults_only1"
            host_var: "host_var_from_defaults_only1"
            group_var: "group_var_from_defaults_only1"
            group_var_all: "group_var_all_from_defaults_only1"
            extra_var: "extra_var_from_defaults_only1"
            """,
            '/etc/ansible/roles/defaults_only1/tasks/main.yml': """
            - debug: msg="here i am"
            """,

            # role defaults_only2
            '/etc/ansible/roles/defaults_only2/defaults/main.yml': """
            default_var: "default_var_from_defaults_only2"
            host_var: "host_var_from_defaults_only2"
            group_var: "group_var_from_defaults_only2"
            group_var_all: "group_var_all_from_defaults_only2"
            extra_var: "extra_var_from_defaults_only2"
            """,
        })

        mock_basedir.return_value = './'
        inv1 = Inventory(loader=fake_loader, variable_manager=v, host_list='/etc/ansible/inventory1')
        inv1.set_playbook_basedir('./')

        play1 = Play.load(dict(
           hosts=['all'],
           roles=['defaults_only1', 'defaults_only2'],
        ), loader=fake_loader, variable_manager=v)

        # first we assert that the defaults as viewed as a whole are the merged results
        # of the defaults from each role, with the last role defined "winning" when
        # there is a variable naming conflict
        res = v.get_vars(loader=fake_loader, play=play1)
        self.assertEqual(res['default_var'], 'default_var_from_defaults_only2')

        # next, we assert that when vars are viewed from the context of a task within a
        # role, that task will see its own role defaults before any other role's
        blocks = play1.compile()
        task = blocks[1].block[0]
        res = v.get_vars(loader=fake_loader, play=play1, task=task)
        self.assertEqual(res['default_var'], 'default_var_from_defaults_only1')

        # next we assert the precendence of inventory variables
        v.set_inventory(inv1)
        h1 = inv1.get_host('host1')

        res = v.get_vars(loader=fake_loader, play=play1, host=h1)
        self.assertEqual(res['group_var'], 'group_var_from_inventory_group1')
        self.assertEqual(res['host_var'], 'host_var_from_inventory_host1')

        # next we test with group_vars/ files loaded
        fake_loader.push("/etc/ansible/group_vars/all", """
        group_var_all: group_var_all_from_group_vars_all
        """)
        fake_loader.push("/etc/ansible/group_vars/group1", """
        group_var: group_var_from_group_vars_group1
        """)
        fake_loader.push("/etc/ansible/group_vars/group3", """
        # this is a dummy, which should not be used anywhere
        group_var: group_var_from_group_vars_group3
        """)
        fake_loader.push("/etc/ansible/host_vars/host1", """
        host_var: host_var_from_host_vars_host1
        """)

        v.add_group_vars_file("/etc/ansible/group_vars/all", loader=fake_loader)
        v.add_group_vars_file("/etc/ansible/group_vars/group1", loader=fake_loader)
        v.add_group_vars_file("/etc/ansible/group_vars/group2", loader=fake_loader)
        v.add_host_vars_file("/etc/ansible/host_vars/host1", loader=fake_loader)

        res = v.get_vars(loader=fake_loader, play=play1, host=h1)
        self.assertEqual(res['group_var'], 'group_var_from_group_vars_group1')
        self.assertEqual(res['group_var_all'], 'group_var_all_from_group_vars_all')
        self.assertEqual(res['host_var'], 'host_var_from_host_vars_host1')

        # add in the fact cache
        v._fact_cache['host1'] = dict(fact_cache_var="fact_cache_var_from_fact_cache")

        res = v.get_vars(loader=fake_loader, play=play1, host=h1)
        self.assertEqual(res['fact_cache_var'], 'fact_cache_var_from_fact_cache')
<commit_msg>Adding unit tests for playbook group/host vars<commit_after># (c) 2012-2014, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

# Make coding more python3-ish
from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

from collections import defaultdict
from six import iteritems

from ansible.compat.tests import unittest
from ansible.compat.tests.mock import patch, MagicMock
from ansible.inventory import Inventory
from ansible.playbook.play import Play
from ansible.vars import VariableManager

from units.mock.loader import DictDataLoader
from units.mock.path import mock_unfrackpath_noop

class TestVariableManager(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_basic_manager(self):
        fake_loader = DictDataLoader({})

        v = VariableManager()
        vars = v.get_vars(loader=fake_loader, use_cache=False)
        if 'omit' in vars:
            del vars['omit']
        if 'vars' in vars:
            del vars['vars']
        if 'ansible_version' in vars:
            del vars['ansible_version']
        if 'ansible_check_mode' in vars:
            del vars['ansible_check_mode']

        self.assertEqual(vars, dict(playbook_dir='.'))

    def test_variable_manager_extra_vars(self):
        fake_loader = DictDataLoader({})

        extra_vars = dict(a=1, b=2, c=3)
        v = VariableManager()
        v.extra_vars = extra_vars

        vars = v.get_vars(loader=fake_loader, use_cache=False)

        for (key, val) in iteritems(extra_vars):
            self.assertEqual(vars.get(key), val)

        self.assertIsNot(v.extra_vars, extra_vars)

    def test_variable_manager_host_vars_file(self):
        fake_loader = DictDataLoader({
            "host_vars/hostname1.yml": """
               foo: bar
            """,
            "other_path/host_vars/hostname1.yml": """
               foo: bam
               baa: bat
            """,
            "host_vars/host.name.yml": """
               host_with_dots: true
            """,
        })

        v = VariableManager()
        v.add_host_vars_file("host_vars/hostname1.yml", loader=fake_loader)
        v.add_host_vars_file("other_path/host_vars/hostname1.yml", loader=fake_loader)
        self.assertIn("hostname1", v._host_vars_files)
        self.assertEqual(v._host_vars_files["hostname1"], [dict(foo="bar"), dict(foo="bam", baa="bat")])

        mock_host = MagicMock()
        mock_host.get_name.return_value = "hostname1"
        mock_host.get_vars.return_value = dict()
        mock_host.get_groups.return_value = ()
        mock_host.get_group_vars.return_value = dict()

        self.assertEqual(v.get_vars(loader=fake_loader, host=mock_host, use_cache=False).get("foo"), "bam")
        self.assertEqual(v.get_vars(loader=fake_loader, host=mock_host, use_cache=False).get("baa"), "bat")

        v.add_host_vars_file("host_vars/host.name", loader=fake_loader)
        self.assertEqual(v._host_vars_files["host.name"], [dict(host_with_dots=True)])

    def test_variable_manager_group_vars_file(self):
        fake_loader = DictDataLoader({
            "group_vars/all.yml": """
               foo: bar
            """,
            "group_vars/somegroup.yml": """
               bam: baz
            """,
            "other_path/group_vars/somegroup.yml": """
               baa: bat
            """,
            "group_vars/some.group.yml": """
               group_with_dots: true
            """,
        })

        v = VariableManager()
        v.add_group_vars_file("group_vars/all.yml", loader=fake_loader)
        v.add_group_vars_file("group_vars/somegroup.yml", loader=fake_loader)
        v.add_group_vars_file("other_path/group_vars/somegroup.yml", loader=fake_loader)
        self.assertIn("somegroup", v._group_vars_files)
        self.assertEqual(v._group_vars_files["all"], [dict(foo="bar")])
        self.assertEqual(v._group_vars_files["somegroup"], [dict(bam="baz"), dict(baa="bat")])

        mock_group = MagicMock()
        mock_group.name = "somegroup"
        mock_group.get_ancestors.return_value = ()
        mock_group.get_vars.return_value = dict()

        mock_host = MagicMock()
        mock_host.get_name.return_value = "hostname1"
        mock_host.get_vars.return_value = dict()
        mock_host.get_groups.return_value = (mock_group,)
        mock_host.get_group_vars.return_value = dict()

        vars = v.get_vars(loader=fake_loader, host=mock_host, use_cache=False)
        self.assertEqual(vars.get("foo"), "bar")
        self.assertEqual(vars.get("baa"), "bat")

        v.add_group_vars_file("group_vars/some.group", loader=fake_loader)
        self.assertEqual(v._group_vars_files["some.group"], [dict(group_with_dots=True)])

    def test_variable_manager_play_vars(self):
        fake_loader = DictDataLoader({})

        mock_play = MagicMock()
        mock_play.get_vars.return_value = dict(foo="bar")
        mock_play.get_roles.return_value = []
        mock_play.get_vars_files.return_value = []

        v = VariableManager()
        self.assertEqual(v.get_vars(loader=fake_loader, play=mock_play, use_cache=False).get("foo"), "bar")

    def test_variable_manager_play_vars_files(self):
        fake_loader = DictDataLoader({
            "/path/to/somefile.yml": """
               foo: bar
            """
        })

        mock_play = MagicMock()
        mock_play.get_vars.return_value = dict()
        mock_play.get_roles.return_value = []
        mock_play.get_vars_files.return_value = ['/path/to/somefile.yml']

        v = VariableManager()
        self.assertEqual(v.get_vars(loader=fake_loader, play=mock_play, use_cache=False).get("foo"), "bar")

    def test_variable_manager_task_vars(self):
        fake_loader = DictDataLoader({})

        mock_task = MagicMock()
        mock_task._role = None
        mock_task.loop = None
        mock_task.get_vars.return_value = dict(foo="bar")
        mock_task.get_include_params.return_value = dict()

        v = VariableManager()
        self.assertEqual(v.get_vars(loader=fake_loader, task=mock_task, use_cache=False).get("foo"), "bar")

    @patch.object(Inventory, 'basedir')
    @patch('ansible.playbook.role.definition.unfrackpath', mock_unfrackpath_noop)
    def test_variable_manager_precedence(self, mock_basedir):
        '''
        Tests complex variations and combinations of get_vars() with different
        objects to modify the context under which variables are merged.
        '''

        v = VariableManager()
        v._fact_cache = defaultdict(dict)

        fake_loader = DictDataLoader({
            # inventory1
            '/etc/ansible/inventory1': """
            [group2:children]
            group1

            [group1]
            host1 host_var=host_var_from_inventory_host1

            [group1:vars]
            group_var = group_var_from_inventory_group1

            [group2:vars]
            group_var = group_var_from_inventory_group2
            """,

            # role defaults_only1
            '/etc/ansible/roles/defaults_only1/defaults/main.yml': """
            default_var: "default_var_from_defaults_only1"
            host_var: "host_var_from_defaults_only1"
            group_var: "group_var_from_defaults_only1"
            group_var_all: "group_var_all_from_defaults_only1"
            extra_var: "extra_var_from_defaults_only1"
            """,
            '/etc/ansible/roles/defaults_only1/tasks/main.yml': """
            - debug: msg="here i am"
            """,

            # role defaults_only2
            '/etc/ansible/roles/defaults_only2/defaults/main.yml': """
            default_var: "default_var_from_defaults_only2"
            host_var: "host_var_from_defaults_only2"
            group_var: "group_var_from_defaults_only2"
            group_var_all: "group_var_all_from_defaults_only2"
            extra_var: "extra_var_from_defaults_only2"
            """,
        })

        mock_basedir.return_value = './'
        inv1 = Inventory(loader=fake_loader, variable_manager=v, host_list='/etc/ansible/inventory1')
        inv1.set_playbook_basedir('./')

        play1 = Play.load(dict(
           hosts=['all'],
           roles=['defaults_only1', 'defaults_only2'],
        ), loader=fake_loader, variable_manager=v)

        # first we assert that the defaults as viewed as a whole are the merged results
        # of the defaults from each role, with the last role defined "winning" when
        # there is a variable naming conflict
        res = v.get_vars(loader=fake_loader, play=play1)
        self.assertEqual(res['default_var'], 'default_var_from_defaults_only2')

        # next, we assert that when vars are viewed from the context of a task within a
        # role, that task will see its own role defaults before any other role's
        blocks = play1.compile()
        task = blocks[1].block[0]
        res = v.get_vars(loader=fake_loader, play=play1, task=task)
        self.assertEqual(res['default_var'], 'default_var_from_defaults_only1')

        # next we assert the precendence of inventory variables
        v.set_inventory(inv1)
        h1 = inv1.get_host('host1')

        res = v.get_vars(loader=fake_loader, play=play1, host=h1)
        self.assertEqual(res['group_var'], 'group_var_from_inventory_group1')
        self.assertEqual(res['host_var'], 'host_var_from_inventory_host1')

        # next we test with group_vars/ files loaded
        fake_loader.push("/etc/ansible/group_vars/all", """
        group_var_all: group_var_all_from_group_vars_all
        """)
        fake_loader.push("/etc/ansible/group_vars/group1", """
        group_var: group_var_from_group_vars_group1
        """)
        fake_loader.push("/etc/ansible/group_vars/group3", """
        # this is a dummy, which should not be used anywhere
        group_var: group_var_from_group_vars_group3
        """)
        fake_loader.push("/etc/ansible/host_vars/host1", """
        host_var: host_var_from_host_vars_host1
        """)
        fake_loader.push("group_vars/group1", """
        playbook_group_var: playbook_group_var
        """)
        fake_loader.push("host_vars/host1", """
        playbook_host_var: playbook_host_var
        """)

        v.add_group_vars_file("/etc/ansible/group_vars/all", loader=fake_loader)
        v.add_group_vars_file("/etc/ansible/group_vars/group1", loader=fake_loader)
        v.add_group_vars_file("/etc/ansible/group_vars/group2", loader=fake_loader)
        v.add_group_vars_file("group_vars/group1", loader=fake_loader)
        v.add_host_vars_file("/etc/ansible/host_vars/host1", loader=fake_loader)
        v.add_host_vars_file("host_vars/host1", loader=fake_loader)

        res = v.get_vars(loader=fake_loader, play=play1, host=h1)
        self.assertEqual(res['group_var'], 'group_var_from_group_vars_group1')
        self.assertEqual(res['group_var_all'], 'group_var_all_from_group_vars_all')
        self.assertEqual(res['playbook_group_var'], 'playbook_group_var')
        self.assertEqual(res['host_var'], 'host_var_from_host_vars_host1')
        self.assertEqual(res['playbook_host_var'], 'playbook_host_var')

        # add in the fact cache
        v._fact_cache['host1'] = dict(fact_cache_var="fact_cache_var_from_fact_cache")

        res = v.get_vars(loader=fake_loader, play=play1, host=h1)
        self.assertEqual(res['fact_cache_var'], 'fact_cache_var_from_fact_cache')
<|endoftext|>