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

from robot import utils
from robot.common import BaseTestSuite, BaseTestCase
from robot.parsing import TestSuiteData
from robot.errors import ExecutionFailed
from robot.variables import GLOBAL_VARIABLES
from robot.output import LOGGER

from fixture import Setup, Teardown
from timeouts import TestTimeout
from keywords import KeywordFactory
from namespace import Namespace
from userkeyword import UserLibrary

def TestSuite(datasources, settings):
    suitedata = TestSuiteData(datasources, settings['SuiteNames'])
    suite = RunnableTestSuite(suitedata)
    suite.set_options(settings)
    return suite

class RunnableTestSuite(BaseTestSuite):

    _exit_on_failure_error = ('Critical failure occurred and ExitOnFailure '
                              'option is in use')

    def __init__(self, data, testdefaults=None, parent=None):
        BaseTestSuite.__init__(self, data.name, data.source, parent)
        self.variables = GLOBAL_VARIABLES.copy()
        self.variables.set_from_variable_table(data.variables)
        self.source = data.source
        self.doc = data.doc or ''
        self.metadata = data.metadata
        self.imports = data.imports
        self.user_keywords = UserLibrary(data.user_keywords)
        self.setup = data.suite_setup or []
        self.teardown = data.suite_teardown or []
        if not testdefaults:
            testdefaults = _TestCaseDefaults()
        testdefaults.add_defaults(data)
        for suite in data.suites:
            RunnableTestSuite(suite, testdefaults.copy(), parent=self)
        for test in data.tests:
            RunnableTestCase(test, testdefaults, parent=self)
        self._exit_on_failure = False

    def run(self, output, parent=None, error=None):
        self.status = 'RUNNING'
        self.starttime = utils.get_timestamp()
        self.namespace = Namespace(self, parent)
        self.namespace.variables['${SUITE_NAME}'] = self.longname
        init_err = self._init_suite(self.namespace.variables)
        output.start_suite(self)
        setup_err = self._run_fixture(self.setup, output, error, init_err)
        child_err = self._get_child_error(error, init_err, setup_err)
        for suite in self.suites:
            suite.run(output, self, child_err)
            if self._exit_on_failure and not child_err and \
                   suite.critical_stats.failed:
                child_err = self._exit_on_failure_error
        executed_tests = []
        for test in self.tests:
            normname = utils.normalize(test.name)
            if normname in executed_tests:
                LOGGER.warn("Multiple test cases with name '%s' executed in "
                            "test suite '%s'"% (test.name, self.longname))
            executed_tests.append(normname)
            test.run(output, self.namespace, child_err)
            if self._exit_on_failure and not child_err and \
                    test.status == 'FAIL' and test.critical == 'yes':
                child_err = self._exit_on_failure_error
            self._set_prev_test_variables(self.namespace.variables, test)
        self.set_status()
        self.message = self._get_my_error(error, init_err, setup_err)
        self.namespace.variables['${SUITE_STATUS}'] = self.status
        self.namespace.variables['${SUITE_MESSAGE}'] = self.get_full_message()
        td_err = self._run_fixture(self.teardown, output, error, init_err)
        if td_err:
            self.suite_teardown_failed('Suite teardown failed:\n%s' % td_err)
        self.endtime = utils.get_timestamp()
        self.elapsedtime = utils.get_elapsed_time(self.starttime, self.endtime)
        self._set_prev_test_variables(GLOBAL_VARIABLES,
                                      varz=self.namespace.variables)
        output.end_suite(self)
        self.namespace.end_suite()

    def _init_suite(self, varz):
        errors = []
        self.doc = varz.replace_meta('Documentation', self.doc, errors)
        self.setup = Setup(varz.replace_meta('Setup', self.setup, errors))
        self.teardown = Teardown(
            varz.replace_meta('Teardown', self.teardown, errors))
        for name, value in self.metadata.items():
            self.metadata[name] = varz.replace_meta(name, value, errors)
        if errors:
            return 'Suite initialization failed:\n%s' % '\n'.join(errors)

    def _get_child_error(self, error, init_error, setup_error):
        if error:
            return error
        if init_error:
            return 'Initialization of the parent suite failed.'
        if setup_error:
            return 'Setup of the parent suite failed.'
        return None

    def _get_my_error(self, error, init_error, setup_error):
        if error:
            return error
        if init_error:
            return init_error
        if setup_error:
            return 'Suite setup failed:\n%s' % setup_error
        return ''

    def _run_fixture(self, fixture, output, error, init_error):
        if fixture and not (error or init_error):
            try:
                fixture.run(output, self.namespace)
            except ExecutionFailed:
                return utils.get_error_message()

    def _set_prev_test_variables(self, destination, test=None, varz=None):
        if test is not None:
            name, status, message = test.name, test.status, test.message
        else:
            name, status, message = varz['${PREV_TEST_NAME}'], \
                    varz['${PREV_TEST_STATUS}'], varz['${PREV_TEST_MESSAGE}']
        destination['${PREV_TEST_NAME}'] = name
        destination['${PREV_TEST_STATUS}'] = status
        destination['${PREV_TEST_MESSAGE}'] = message

class RunnableTestCase(BaseTestCase):

    def __init__(self, data, defaults, parent):
        BaseTestCase.__init__(self, data.name, parent)
        self.doc = data.doc or ''
        self.setup = utils.get_not_none(data.setup, defaults.test_setup)
        self.teardown = utils.get_not_none(data.teardown,
                                           defaults.test_teardown)
        self.tags = defaults.force_tags \
                    + utils.get_not_none(data.tags, defaults.default_tags)
        self.timeout = utils.get_not_none(data.timeout, defaults.test_timeout)
        self.keywords = [ KeywordFactory(kw) for kw in data.keywords ]

    def run(self, output, namespace, error=None):
        self.starttime = utils.get_timestamp()
        init_err = self._init_test(namespace.variables)
        error = error or init_err
        namespace.start_test(self)
        output.start_test(self)
        if not error:
            self._run(output, namespace)
        else:
            self.status = 'FAIL'
            self.message = error
        self.endtime = utils.get_timestamp()
        self.elapsedtime = utils.get_elapsed_time(self.starttime, self.endtime)
        output.end_test(self)
        namespace.end_test()

    def _init_test(self, varz):
        errors = []
        self.doc = varz.replace_meta('Documentation', self.doc, errors)
        self.setup = Setup(varz.replace_meta('Setup', self.setup, errors))
        self.teardown = Teardown(
            varz.replace_meta('Teardown', self.teardown, errors))
        self.tags = utils.normalize_tags(
            varz.replace_meta('Tags', self.tags, errors))
        self.timeout = TestTimeout(
            *varz.replace_meta('Timeout', self.timeout, errors))
        if errors:
            return 'Test case initialization failed:\n%s' % '\n'.join(errors)
        if not self.keywords:
            return 'Test case contains no keywords'

    def _run(self, output, namespace):
        namespace.variables['${TEST_NAME}'] = self.name
        namespace.variables['@{TEST_TAGS}'] = self.tags
        self.status = 'RUNNING'
        self.timeout.start()
        setup_err = self._run_fixture(self.setup, output, namespace)
        kw_err = self._run_keywords(output, namespace, setup_err)
        self.message = self._get_message(setup_err, kw_err)
        self.status = self.message == '' and 'PASS' or 'FAIL'
        namespace.variables['${TEST_STATUS}'] = self.status
        namespace.variables['${TEST_MESSAGE}'] = self.message
        td_err = self._run_fixture(self.teardown, output, namespace)
        if td_err:
            self.message = self._get_teardown_error(self.message, td_err)
            self.status = 'FAIL'
        if self.status == 'PASS' and self.timeout.timed_out():
            self.status = 'FAIL'
            self.message = self.timeout.get_message()

    def _run_fixture(self, fixture, output, namespace):
        if fixture:
            try:
                fixture.run(output, namespace)
            except ExecutionFailed:
                return utils.get_error_message()

    def _run_keywords(self, output, namespace, setup_err):
        if not setup_err:
            for kw in self.keywords:
                try:
                    kw.run(output, namespace)
                except ExecutionFailed:
                    return utils.get_error_message()

    def _get_message(self, setup_err, kw_err):
        if setup_err:
            return 'Setup failed:\n%s' % setup_err
        return kw_err or ''

    def _get_teardown_error(self, message, teardown_err):
        if message == '':
            return 'Teardown failed:\n%s' % teardown_err
        return '%s\n\nAlso teardown failed:\n%s' % (message, teardown_err)

class _TestCaseDefaults:

    def __init__(self):
        self.force_tags = []
        self.default_tags = []
        self.test_setup = []
        self.test_teardown = []
        self.test_timeout = []

    def add_defaults(self, suite):
        if suite.force_tags:
            self.force_tags.extend(suite.force_tags)
        if suite.default_tags:
            self.default_tags = suite.default_tags
        if suite.test_setup:
            self.test_setup = suite.test_setup
        if suite.test_teardown:
            self.test_teardown = suite.test_teardown
        if suite.test_timeout:
            self.test_timeout = suite.test_timeout

    def copy(self):
        copy = _TestCaseDefaults()
        copy.force_tags = self.force_tags[:]
        copy.default_tags = self.default_tags[:]
        copy.test_setup = self.test_setup[:]
        copy.test_teardown = self.test_teardown[:]
        copy.test_timeout = self.test_timeout[:]
        return copy
<commit_msg>set test status to RUNNING before calling start_test (issue 328)<commit_after>#  Copyright 2008-2009 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from robot import utils
from robot.common import BaseTestSuite, BaseTestCase
from robot.parsing import TestSuiteData
from robot.errors import ExecutionFailed
from robot.variables import GLOBAL_VARIABLES
from robot.output import LOGGER

from fixture import Setup, Teardown
from timeouts import TestTimeout
from keywords import KeywordFactory
from namespace import Namespace
from userkeyword import UserLibrary

def TestSuite(datasources, settings):
    suitedata = TestSuiteData(datasources, settings['SuiteNames'])
    suite = RunnableTestSuite(suitedata)
    suite.set_options(settings)
    return suite

class RunnableTestSuite(BaseTestSuite):

    _exit_on_failure_error = ('Critical failure occurred and ExitOnFailure '
                              'option is in use')

    def __init__(self, data, testdefaults=None, parent=None):
        BaseTestSuite.__init__(self, data.name, data.source, parent)
        self.variables = GLOBAL_VARIABLES.copy()
        self.variables.set_from_variable_table(data.variables)
        self.source = data.source
        self.doc = data.doc or ''
        self.metadata = data.metadata
        self.imports = data.imports
        self.user_keywords = UserLibrary(data.user_keywords)
        self.setup = data.suite_setup or []
        self.teardown = data.suite_teardown or []
        if not testdefaults:
            testdefaults = _TestCaseDefaults()
        testdefaults.add_defaults(data)
        for suite in data.suites:
            RunnableTestSuite(suite, testdefaults.copy(), parent=self)
        for test in data.tests:
            RunnableTestCase(test, testdefaults, parent=self)
        self._exit_on_failure = False

    def run(self, output, parent=None, error=None):
        self.status = 'RUNNING'
        self.starttime = utils.get_timestamp()
        self.namespace = Namespace(self, parent)
        self.namespace.variables['${SUITE_NAME}'] = self.longname
        init_err = self._init_suite(self.namespace.variables)
        output.start_suite(self)
        setup_err = self._run_fixture(self.setup, output, error, init_err)
        child_err = self._get_child_error(error, init_err, setup_err)
        for suite in self.suites:
            suite.run(output, self, child_err)
            if self._exit_on_failure and not child_err and \
                   suite.critical_stats.failed:
                child_err = self._exit_on_failure_error
        executed_tests = []
        for test in self.tests:
            normname = utils.normalize(test.name)
            if normname in executed_tests:
                LOGGER.warn("Multiple test cases with name '%s' executed in "
                            "test suite '%s'"% (test.name, self.longname))
            executed_tests.append(normname)
            test.run(output, self.namespace, child_err)
            if self._exit_on_failure and not child_err and \
                    test.status == 'FAIL' and test.critical == 'yes':
                child_err = self._exit_on_failure_error
            self._set_prev_test_variables(self.namespace.variables, test)
        self.set_status()
        self.message = self._get_my_error(error, init_err, setup_err)
        self.namespace.variables['${SUITE_STATUS}'] = self.status
        self.namespace.variables['${SUITE_MESSAGE}'] = self.get_full_message()
        td_err = self._run_fixture(self.teardown, output, error, init_err)
        if td_err:
            self.suite_teardown_failed('Suite teardown failed:\n%s' % td_err)
        self.endtime = utils.get_timestamp()
        self.elapsedtime = utils.get_elapsed_time(self.starttime, self.endtime)
        self._set_prev_test_variables(GLOBAL_VARIABLES,
                                      varz=self.namespace.variables)
        output.end_suite(self)
        self.namespace.end_suite()

    def _init_suite(self, varz):
        errors = []
        self.doc = varz.replace_meta('Documentation', self.doc, errors)
        self.setup = Setup(varz.replace_meta('Setup', self.setup, errors))
        self.teardown = Teardown(
            varz.replace_meta('Teardown', self.teardown, errors))
        for name, value in self.metadata.items():
            self.metadata[name] = varz.replace_meta(name, value, errors)
        if errors:
            return 'Suite initialization failed:\n%s' % '\n'.join(errors)

    def _get_child_error(self, error, init_error, setup_error):
        if error:
            return error
        if init_error:
            return 'Initialization of the parent suite failed.'
        if setup_error:
            return 'Setup of the parent suite failed.'
        return None

    def _get_my_error(self, error, init_error, setup_error):
        if error:
            return error
        if init_error:
            return init_error
        if setup_error:
            return 'Suite setup failed:\n%s' % setup_error
        return ''

    def _run_fixture(self, fixture, output, error, init_error):
        if fixture and not (error or init_error):
            try:
                fixture.run(output, self.namespace)
            except ExecutionFailed:
                return utils.get_error_message()

    def _set_prev_test_variables(self, destination, test=None, varz=None):
        if test is not None:
            name, status, message = test.name, test.status, test.message
        else:
            name, status, message = varz['${PREV_TEST_NAME}'], \
                    varz['${PREV_TEST_STATUS}'], varz['${PREV_TEST_MESSAGE}']
        destination['${PREV_TEST_NAME}'] = name
        destination['${PREV_TEST_STATUS}'] = status
        destination['${PREV_TEST_MESSAGE}'] = message

class RunnableTestCase(BaseTestCase):

    def __init__(self, data, defaults, parent):
        BaseTestCase.__init__(self, data.name, parent)
        self.doc = data.doc or ''
        self.setup = utils.get_not_none(data.setup, defaults.test_setup)
        self.teardown = utils.get_not_none(data.teardown,
                                           defaults.test_teardown)
        self.tags = defaults.force_tags \
                    + utils.get_not_none(data.tags, defaults.default_tags)
        self.timeout = utils.get_not_none(data.timeout, defaults.test_timeout)
        self.keywords = [ KeywordFactory(kw) for kw in data.keywords ]

    def run(self, output, namespace, error=None):
        self.status = 'RUNNING'
        self.starttime = utils.get_timestamp()
        init_err = self._init_test(namespace.variables)
        error = error or init_err
        namespace.start_test(self)
        output.start_test(self)
        if not error:
            self._run(output, namespace)
        else:
            self.status = 'FAIL'
            self.message = error
        self.endtime = utils.get_timestamp()
        self.elapsedtime = utils.get_elapsed_time(self.starttime, self.endtime)
        output.end_test(self)
        namespace.end_test()

    def _init_test(self, varz):
        errors = []
        self.doc = varz.replace_meta('Documentation', self.doc, errors)
        self.setup = Setup(varz.replace_meta('Setup', self.setup, errors))
        self.teardown = Teardown(
            varz.replace_meta('Teardown', self.teardown, errors))
        self.tags = utils.normalize_tags(
            varz.replace_meta('Tags', self.tags, errors))
        self.timeout = TestTimeout(
            *varz.replace_meta('Timeout', self.timeout, errors))
        if errors:
            return 'Test case initialization failed:\n%s' % '\n'.join(errors)
        if not self.keywords:
            return 'Test case contains no keywords'

    def _run(self, output, namespace):
        namespace.variables['${TEST_NAME}'] = self.name
        namespace.variables['@{TEST_TAGS}'] = self.tags
        self.timeout.start()
        setup_err = self._run_fixture(self.setup, output, namespace)
        kw_err = self._run_keywords(output, namespace, setup_err)
        self.message = self._get_message(setup_err, kw_err)
        self.status = self.message == '' and 'PASS' or 'FAIL'
        namespace.variables['${TEST_STATUS}'] = self.status
        namespace.variables['${TEST_MESSAGE}'] = self.message
        td_err = self._run_fixture(self.teardown, output, namespace)
        if td_err:
            self.message = self._get_teardown_error(self.message, td_err)
            self.status = 'FAIL'
        if self.status == 'PASS' and self.timeout.timed_out():
            self.status = 'FAIL'
            self.message = self.timeout.get_message()

    def _run_fixture(self, fixture, output, namespace):
        if fixture:
            try:
                fixture.run(output, namespace)
            except ExecutionFailed:
                return utils.get_error_message()

    def _run_keywords(self, output, namespace, setup_err):
        if not setup_err:
            for kw in self.keywords:
                try:
                    kw.run(output, namespace)
                except ExecutionFailed:
                    return utils.get_error_message()

    def _get_message(self, setup_err, kw_err):
        if setup_err:
            return 'Setup failed:\n%s' % setup_err
        return kw_err or ''

    def _get_teardown_error(self, message, teardown_err):
        if message == '':
            return 'Teardown failed:\n%s' % teardown_err
        return '%s\n\nAlso teardown failed:\n%s' % (message, teardown_err)

class _TestCaseDefaults:

    def __init__(self):
        self.force_tags = []
        self.default_tags = []
        self.test_setup = []
        self.test_teardown = []
        self.test_timeout = []

    def add_defaults(self, suite):
        if suite.force_tags:
            self.force_tags.extend(suite.force_tags)
        if suite.default_tags:
            self.default_tags = suite.default_tags
        if suite.test_setup:
            self.test_setup = suite.test_setup
        if suite.test_teardown:
            self.test_teardown = suite.test_teardown
        if suite.test_timeout:
            self.test_timeout = suite.test_timeout

    def copy(self):
        copy = _TestCaseDefaults()
        copy.force_tags = self.force_tags[:]
        copy.default_tags = self.default_tags[:]
        copy.test_setup = self.test_setup[:]
        copy.test_teardown = self.test_teardown[:]
        copy.test_timeout = self.test_timeout[:]
        return copy
<|endoftext|>


<commit_before>
import os
import re
import json

def get_data(filename, cache_data=True, **kwargs):
    ''' A function to load tiff data files and its faster loading caches.

    In case the file has been loaded previously and caching was enabled
    (default=True) then a way faster loading cache file has been saved. This
    function will load this one automatically.

    :Args:
        - filename : specify either an absolute or a relative path to a file
          you would like to load.

    :Kwargs:
        - cache_data=True : if no cache is available, create a faster loading
          cache file.
     '''
    from ..io import data_handler
    return data_handler.DataHandler().get_data(filename, cache_data=cache_data, **kwargs)

def get_meta_file(filename, meta_file_name='data_meta.json', dirretracts=3):
    ''' Search for data_meta.json file in the parent directories '''
    trunk = filename
    for _ in range(dirretracts):
        trunk, tail = os.path.split(trunk)
        meta_file = os.path.join(trunk, meta_file_name)
        print(meta_file)
        if os.path.exists(meta_file):
            return meta_file

def get_meta_info(filename, dirrectracts=3, idselector=r'.*/(?P<id>\w*)_.*.\w*$'):
    ''' Search for data_meta.json and return its content + recid '''
    recidmatch = re.match(idselector, filename)
    if recidmatch is None:
        raise ValueError('Could not identify id-pattern.')
    recid = recidmatch.groupdict()['id']
    metafile = get_meta_file(filename, dirretracts=dirrectracts)
    if metafile is None:
        return
    with open(metafile, 'r') as jfi:
        metadict = json.load(jfi)
    return metadict, recid
<commit_msg>remove print<commit_after>
import os
import re
import json

def get_data(filename, cache_data=True, **kwargs):
    ''' A function to load tiff data files and its faster loading caches.

    In case the file has been loaded previously and caching was enabled
    (default=True) then a way faster loading cache file has been saved. This
    function will load this one automatically.

    :Args:
        - filename : specify either an absolute or a relative path to a file
          you would like to load.

    :Kwargs:
        - cache_data=True : if no cache is available, create a faster loading
          cache file.
     '''
    from ..io import data_handler
    return data_handler.DataHandler().get_data(filename, cache_data=cache_data, **kwargs)

def get_meta_file(filename, meta_file_name='data_meta.json', dirretracts=3):
    ''' Search for data_meta.json file in the parent directories '''
    trunk = filename
    for _ in range(dirretracts):
        trunk, tail = os.path.split(trunk)
        meta_file = os.path.join(trunk, meta_file_name)
        if os.path.exists(meta_file):
            return meta_file

def get_meta_info(filename, dirrectracts=3, idselector=r'.*/(?P<id>\w*)_.*.\w*$'):
    ''' Search for data_meta.json and return its content + recid '''
    recidmatch = re.match(idselector, filename)
    if recidmatch is None:
        raise ValueError('Could not identify id-pattern.')
    recid = recidmatch.groupdict()['id']
    metafile = get_meta_file(filename, dirretracts=dirrectracts)
    if metafile is None:
        return
    with open(metafile, 'r') as jfi:
        metadict = json.load(jfi)
    return metadict, recid
<|endoftext|>


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

import abc

from oslo_log import log as logging
from osprofiler import profiler

from mistral.db.v2 import api as db_api
from mistral import exceptions as exc
from mistral.lang import parser as spec_parser
from mistral import utils as u
from mistral.workflow import commands
from mistral.workflow import data_flow
from mistral.workflow import states

LOG = logging.getLogger(__name__)

@profiler.trace('wf-controller-get-controller', hide_args=True)
def get_controller(wf_ex, wf_spec=None):
    """Gets a workflow controller instance by given workflow execution object.

    :param wf_ex: Workflow execution object.
    :param wf_spec: Workflow specification object. If passed, the method works
        faster.
    :returns: Workflow controller class.
    """

    if not wf_spec:
        wf_spec = spec_parser.get_workflow_spec_by_execution_id(wf_ex.id)

    wf_type = wf_spec.get_type()

    ctrl_cls = None

    for cls in u.iter_subclasses(WorkflowController):
        if cls.__workflow_type__ == wf_type:
            ctrl_cls = cls
            break

    if not ctrl_cls:
        raise exc.MistralError(
            'Failed to find a workflow controller [type=%s]' % wf_type
        )

    return ctrl_cls(wf_ex, wf_spec)

class TaskLogicalState(object):
    """Task logical state.

    This data structure describes what state a task should have according
    to the logic of the workflow type and state of other tasks.
    """

    def __init__(self, state, state_info=None, cardinality=0,
                 triggered_by=None):
        self.state = state
        self.state_info = state_info
        self.cardinality = cardinality
        self.triggered_by = triggered_by or []

    def get_state(self):
        return self.state

    def get_state_info(self):
        return self.state_info

    def get_cardinality(self):
        return self.cardinality

    def get_triggered_by(self):
        return self.get_triggered_by

class WorkflowController(object):
    """Workflow Controller base class.

    Different workflow controllers implement different workflow algorithms.
    In practice it may actually mean that there may be multiple ways of
    describing workflow models (and even languages) that will be supported
    by Mistral.
    """

    def __init__(self, wf_ex, wf_spec=None):
        """Creates a new workflow controller.

        :param wf_ex: Workflow execution.

        :param wf_spec: Workflow specification.
        """
        self.wf_ex = wf_ex

        if wf_spec is None:
            wf_spec = spec_parser.get_workflow_spec_by_execution_id(wf_ex.id)

        self.wf_spec = wf_spec

    @profiler.trace('workflow-controller-continue-workflow', hide_args=True)
    def continue_workflow(self, task_ex=None):
        """Calculates a list of commands to continue the workflow.

        Given a workflow specification this method makes required analysis
        according to this workflow type rules and identifies a list of
        commands needed to continue the workflow.

        :param task_ex: Task execution that caused workflow continuation.
            Optional. If not specified, it means that no certain task caused
            this operation (e.g. workflow has been just started or resumed
            manually).
        :return: List of workflow commands (instances of
            mistral.workflow.commands.WorkflowCommand).
        """

        if self._is_paused_or_completed():
            return []

        return self._find_next_commands(task_ex)

    def rerun_tasks(self, task_execs, reset=True):
        """Gets commands to rerun existing task executions.

        :param task_execs: List of task executions.
        :param reset: If true, then purge action executions for the tasks.
        :return: List of workflow commands.
        """
        if self._is_paused_or_completed():
            return []

        cmds = [
            commands.RunExistingTask(self.wf_ex, self.wf_spec, t_e, reset,
                                     rerun=True)
            for t_e in task_execs
        ]

        LOG.debug("Commands to rerun workflow tasks: %s", cmds)

        return cmds

    @abc.abstractmethod
    def get_logical_task_state(self, task_ex):
        """Determines a logical state of the given task.

        :param task_ex: Task execution.
        :return: Tuple (state, state_info, cardinality) where 'state' and
            'state_info' are the corresponding values which the given
             task should have according to workflow rules and current
            states of other tasks. 'cardinality' gives the estimation on
            the number of preconditions that are not yet met in case if
            state is WAITING. This number can be used to estimate how
            frequently we can refresh the state of this task.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def find_indirectly_affected_task_executions(self, task_name):
        """Get a set of task executions indirectly affected by the given.

        :param task_name: Task name.
        :return: Task executions that can be indirectly affected by a task
            identified by the given name.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def is_error_handled_for(self, task_ex):
        """Determines if error is handled for specific task.

        :param task_ex: Task execution.
        :return: True if either there is no error at all or
            error is considered handled.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def all_errors_handled(self):
        """Determines if all errors (if any) are handled.

        :return: True if either there aren't errors at all or all
            errors are considered handled.
        """
        raise NotImplementedError

    def any_cancels(self):
        """Determines if there are any task cancellations.

        :return: True if there is one or more tasks in cancelled state.
        """
        return len(self._get_task_executions(state=states.CANCELLED)) > 0

    @abc.abstractmethod
    def evaluate_workflow_final_context(self):
        """Evaluates final workflow context assuming that workflow has finished.

        :return: Final workflow context.
        """
        raise NotImplementedError

    def get_task_inbound_context(self, task_spec):
        # TODO(rakhmerov): This method should also be able to work with task_ex
        # to cover 'split' (aka 'merge') use case.
        upstream_task_execs = self._get_upstream_task_executions(task_spec)

        return data_flow.evaluate_upstream_context(upstream_task_execs)

    @abc.abstractmethod
    def _get_upstream_task_executions(self, task_spec):
        """Gets workflow upstream tasks for the given task.

        :param task_spec: Task specification.
        :return: List of upstream task executions for the given task spec.
        """
        raise NotImplementedError

    def may_complete_workflow(self, task_ex):
        """Determines if the task execution may lead to workflow completion."""

        return states.is_completed(task_ex.state)

    @abc.abstractmethod
    def _find_next_commands(self, task_ex):
        """Finds commands that should run next.

        A concrete algorithm of finding such tasks depends on a concrete
        workflow controller.

        :return: List of workflow commands.
        """

        # If task execution was passed then we should make all calculations
        # only based on it.
        if task_ex:
            return []

        # Add all tasks in IDLE state.
        return [
            commands.RunExistingTask(self.wf_ex, self.wf_spec, t)
            for t in self._get_task_executions(state=states.IDLE)
        ]

    def _is_paused_or_completed(self):
        return states.is_paused_or_completed(self.wf_ex.state)

    def _get_task_executions(self, **kwargs):
        return db_api.get_task_executions(
            workflow_execution_id=self.wf_ex.id,
            sort_keys=[],  # disable sorting
            **kwargs
        )
<commit_msg>Use get_task_executions_count for any_cancels method<commit_after># Copyright 2014 - Mirantis, Inc.
# Copyright 2015 - StackStorm, Inc.
# Copyright 2015 - Huawei Technologies Co. Ltd
# Copyright 2016 - Brocade Communications Systems, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import abc

from oslo_log import log as logging
from osprofiler import profiler

from mistral.db.v2 import api as db_api
from mistral import exceptions as exc
from mistral.lang import parser as spec_parser
from mistral import utils as u
from mistral.workflow import commands
from mistral.workflow import data_flow
from mistral.workflow import states

LOG = logging.getLogger(__name__)

@profiler.trace('wf-controller-get-controller', hide_args=True)
def get_controller(wf_ex, wf_spec=None):
    """Gets a workflow controller instance by given workflow execution object.

    :param wf_ex: Workflow execution object.
    :param wf_spec: Workflow specification object. If passed, the method works
        faster.
    :returns: Workflow controller class.
    """

    if not wf_spec:
        wf_spec = spec_parser.get_workflow_spec_by_execution_id(wf_ex.id)

    wf_type = wf_spec.get_type()

    ctrl_cls = None

    for cls in u.iter_subclasses(WorkflowController):
        if cls.__workflow_type__ == wf_type:
            ctrl_cls = cls
            break

    if not ctrl_cls:
        raise exc.MistralError(
            'Failed to find a workflow controller [type=%s]' % wf_type
        )

    return ctrl_cls(wf_ex, wf_spec)

class TaskLogicalState(object):
    """Task logical state.

    This data structure describes what state a task should have according
    to the logic of the workflow type and state of other tasks.
    """

    def __init__(self, state, state_info=None, cardinality=0,
                 triggered_by=None):
        self.state = state
        self.state_info = state_info
        self.cardinality = cardinality
        self.triggered_by = triggered_by or []

    def get_state(self):
        return self.state

    def get_state_info(self):
        return self.state_info

    def get_cardinality(self):
        return self.cardinality

    def get_triggered_by(self):
        return self.get_triggered_by

class WorkflowController(object):
    """Workflow Controller base class.

    Different workflow controllers implement different workflow algorithms.
    In practice it may actually mean that there may be multiple ways of
    describing workflow models (and even languages) that will be supported
    by Mistral.
    """

    def __init__(self, wf_ex, wf_spec=None):
        """Creates a new workflow controller.

        :param wf_ex: Workflow execution.

        :param wf_spec: Workflow specification.
        """
        self.wf_ex = wf_ex

        if wf_spec is None:
            wf_spec = spec_parser.get_workflow_spec_by_execution_id(wf_ex.id)

        self.wf_spec = wf_spec

    @profiler.trace('workflow-controller-continue-workflow', hide_args=True)
    def continue_workflow(self, task_ex=None):
        """Calculates a list of commands to continue the workflow.

        Given a workflow specification this method makes required analysis
        according to this workflow type rules and identifies a list of
        commands needed to continue the workflow.

        :param task_ex: Task execution that caused workflow continuation.
            Optional. If not specified, it means that no certain task caused
            this operation (e.g. workflow has been just started or resumed
            manually).
        :return: List of workflow commands (instances of
            mistral.workflow.commands.WorkflowCommand).
        """

        if self._is_paused_or_completed():
            return []

        return self._find_next_commands(task_ex)

    def rerun_tasks(self, task_execs, reset=True):
        """Gets commands to rerun existing task executions.

        :param task_execs: List of task executions.
        :param reset: If true, then purge action executions for the tasks.
        :return: List of workflow commands.
        """
        if self._is_paused_or_completed():
            return []

        cmds = [
            commands.RunExistingTask(self.wf_ex, self.wf_spec, t_e, reset,
                                     rerun=True)
            for t_e in task_execs
        ]

        LOG.debug("Commands to rerun workflow tasks: %s", cmds)

        return cmds

    @abc.abstractmethod
    def get_logical_task_state(self, task_ex):
        """Determines a logical state of the given task.

        :param task_ex: Task execution.
        :return: Tuple (state, state_info, cardinality) where 'state' and
            'state_info' are the corresponding values which the given
             task should have according to workflow rules and current
            states of other tasks. 'cardinality' gives the estimation on
            the number of preconditions that are not yet met in case if
            state is WAITING. This number can be used to estimate how
            frequently we can refresh the state of this task.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def find_indirectly_affected_task_executions(self, task_name):
        """Get a set of task executions indirectly affected by the given.

        :param task_name: Task name.
        :return: Task executions that can be indirectly affected by a task
            identified by the given name.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def is_error_handled_for(self, task_ex):
        """Determines if error is handled for specific task.

        :param task_ex: Task execution.
        :return: True if either there is no error at all or
            error is considered handled.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def all_errors_handled(self):
        """Determines if all errors (if any) are handled.

        :return: True if either there aren't errors at all or all
            errors are considered handled.
        """
        raise NotImplementedError

    def any_cancels(self):
        """Determines if there are any task cancellations.

        :return: True if there is one or more tasks in cancelled state.
        """
        return db_api.get_task_executions_count(
            workflow_execution_id=self.wf_ex.id,
            state=states.CANCELLED
        ) > 0

    @abc.abstractmethod
    def evaluate_workflow_final_context(self):
        """Evaluates final workflow context assuming that workflow has finished.

        :return: Final workflow context.
        """
        raise NotImplementedError

    def get_task_inbound_context(self, task_spec):
        # TODO(rakhmerov): This method should also be able to work with task_ex
        # to cover 'split' (aka 'merge') use case.
        upstream_task_execs = self._get_upstream_task_executions(task_spec)

        return data_flow.evaluate_upstream_context(upstream_task_execs)

    @abc.abstractmethod
    def _get_upstream_task_executions(self, task_spec):
        """Gets workflow upstream tasks for the given task.

        :param task_spec: Task specification.
        :return: List of upstream task executions for the given task spec.
        """
        raise NotImplementedError

    def may_complete_workflow(self, task_ex):
        """Determines if the task execution may lead to workflow completion."""

        return states.is_completed(task_ex.state)

    @abc.abstractmethod
    def _find_next_commands(self, task_ex):
        """Finds commands that should run next.

        A concrete algorithm of finding such tasks depends on a concrete
        workflow controller.

        :return: List of workflow commands.
        """

        # If task execution was passed then we should make all calculations
        # only based on it.
        if task_ex:
            return []

        # Add all tasks in IDLE state.
        return [
            commands.RunExistingTask(self.wf_ex, self.wf_spec, t)
            for t in self._get_task_executions(state=states.IDLE)
        ]

    def _is_paused_or_completed(self):
        return states.is_paused_or_completed(self.wf_ex.state)

    def _get_task_executions(self, **kwargs):
        return db_api.get_task_executions(
            workflow_execution_id=self.wf_ex.id,
            sort_keys=[],  # disable sorting
            **kwargs
        )
<|endoftext|>


<commit_before>#!/speechlab/software/EPD7/epd-7.1-1-rh5-x86_64/bin/python

import os
import sys
import glob
import argparse
from scai_utils import *

help_doc = "Convert raw DICOM images to nii format. (scai)"

# === Heuristics === #
heur = [["functionalsparse", "bold"],
        ["DIFFUSIONHIGHRES10Min", "diffusion"],
        ["AAHeadScout32", "aascout"],
        ["tflmghmultiecho1mmiso", "T1"],
        ["localizer", "localizer"]]

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=help_doc)
    ap.add_argument("inputDir", type=str, \
                    help="Input raw DICOM data directory")
    ap.add_argument("outputDir", type=str, \
                    help="Output directory (for saving the .nii.gz files)")
    ap.add_argument("sID", type=str, \
                    help="Subject ID (e.g., ANS_M01)")

    if len(sys.argv) == 1:
        ap.print_help()
        sys.exit(0)

    args = ap.parse_args()

    check_dir(args.inputDir)
    check_dir(args.outputDir, bCreate=True)

    # Create temporary directory #
    tmpDir = os.path.join(args.outputDir, "tmp_convert_dcm")
    check_dir(args.outputDir, bCreate=True)
    delete_file_if_exists(tmpDir, recursive=True)
    check_dir(tmpDir, bCreate=True)

    dcm2nii_cmd = "dcm2nii -a n -o %s %s" \
                  % (os.path.abspath(tmpDir), \
                     os.path.join(os.path.abspath(args.inputDir), "*"))
    saydo(dcm2nii_cmd)

    # Apply heuristics to move and rename files #
    niis = glob.glob(os.path.join(tmpDir, "*.nii.gz"))
    niis.sort()
    bvec_bval_done = False
    for (i0, nii) in enumerate(niis):
        [nfp, nfn] = os.path.split(nii)

        imgStr = ""
        imgType = ""

        for (i1, t_item) in enumerate(heur):
            if (t_item[0] in nfn):
                imgStr = t_item[0]
                imgType = t_item[1]
                break

        if imgType == "":
            raise Exception, \
                  "Unrecognized image type for image file name %s" % nii

        # == Parse file name == #
        assert(nfn.endswith(".nii.gz"))
        assert(nfn.count("_") == 1)
        ius = nfn.index("_")
        assert(ius >= 8)
        timeStr = nfn[ius - 8 : ius + 7]
        assert(nfn.count(timeStr) == 1)

        prefix = nfn[:ius - 8]
        nfn1 = nfn.replace(timeStr, "")
        nfn2 = nfn1[len(prefix) :]
        suffix = nfn2.replace(imgStr, "").replace(".nii.gz", "")

        subDir = os.path.join(args.outputDir, imgType)
        check_dir(subDir, bCreate=True)

        if prefix == "":
            imgFN = os.path.join(subDir, \
                                  "%s_%s_%s.nii.gz" \
                                  % (args.sID, imgType, suffix))
        else:
            imgFN = os.path.join(subDir, \
                                  "%s_%s_%s_%s.nii.gz" \
                                  % (args.sID, imgType, suffix, prefix))

        saydo("cp %s %s" % (nii, imgFN))
        check_file(imgFN)

        if imgType == "diffusion" and not bvec_bval_done:
            bvec = nii.replace(".nii.gz", ".bvec")
            bval = nii.replace(".nii.gz", ".bval")
            check_file(bvec)
            check_file(bval)

            bvec_new = imgFN.replace(".nii.gz", ".bvec")
            bval_new = imgFN.replace(".nii.gz", ".bval")

            saydo("cp %s %s" % (bvec, bvec_new))
            saydo("cp %s %s" % (bval, bval_new))

            bvec_bval_done = True

    # === Remove temporary directory === #
    saydo("rm -rf %s" % tmpDir)
<commit_msg>Using /usr/bin/env python<commit_after>#!/usr/bin/env python

import os
import sys
import glob
import argparse
from scai_utils import *

help_doc = "Convert raw DICOM images to nii format. (scai)"

# === Heuristics === #
heur = [["functionalsparse", "bold"],
        ["DIFFUSIONHIGHRES10Min", "diffusion"],
        ["AAHeadScout32", "aascout"],
        ["tflmghmultiecho1mmiso", "T1"],
        ["localizer", "localizer"]]

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=help_doc)
    ap.add_argument("inputDir", type=str, \
                    help="Input raw DICOM data directory")
    ap.add_argument("outputDir", type=str, \
                    help="Output directory (for saving the .nii.gz files)")
    ap.add_argument("sID", type=str, \
                    help="Subject ID (e.g., ANS_M01)")

    if len(sys.argv) == 1:
        ap.print_help()
        sys.exit(0)

    args = ap.parse_args()

    check_dir(args.inputDir)
    check_dir(args.outputDir, bCreate=True)

    # Create temporary directory #
    tmpDir = os.path.join(args.outputDir, "tmp_convert_dcm")
    check_dir(args.outputDir, bCreate=True)
    delete_file_if_exists(tmpDir, recursive=True)
    check_dir(tmpDir, bCreate=True)

    dcm2nii_cmd = "dcm2nii -a n -o %s %s" \
                  % (os.path.abspath(tmpDir), \
                     os.path.join(os.path.abspath(args.inputDir), "*"))
    saydo(dcm2nii_cmd)

    # Apply heuristics to move and rename files #
    niis = glob.glob(os.path.join(tmpDir, "*.nii.gz"))
    niis.sort()
    bvec_bval_done = False
    for (i0, nii) in enumerate(niis):
        [nfp, nfn] = os.path.split(nii)

        imgStr = ""
        imgType = ""

        for (i1, t_item) in enumerate(heur):
            if (t_item[0] in nfn):
                imgStr = t_item[0]
                imgType = t_item[1]
                break

        if imgType == "":
            raise Exception, \
                  "Unrecognized image type for image file name %s" % nii

        # == Parse file name == #
        assert(nfn.endswith(".nii.gz"))
        assert(nfn.count("_") == 1)
        ius = nfn.index("_")
        assert(ius >= 8)
        timeStr = nfn[ius - 8 : ius + 7]
        assert(nfn.count(timeStr) == 1)

        prefix = nfn[:ius - 8]
        nfn1 = nfn.replace(timeStr, "")
        nfn2 = nfn1[len(prefix) :]
        suffix = nfn2.replace(imgStr, "").replace(".nii.gz", "")

        subDir = os.path.join(args.outputDir, imgType)
        check_dir(subDir, bCreate=True)

        if prefix == "":
            imgFN = os.path.join(subDir, \
                                  "%s_%s_%s.nii.gz" \
                                  % (args.sID, imgType, suffix))
        else:
            imgFN = os.path.join(subDir, \
                                  "%s_%s_%s_%s.nii.gz" \
                                  % (args.sID, imgType, suffix, prefix))

        saydo("cp %s %s" % (nii, imgFN))
        check_file(imgFN)

        if imgType == "diffusion" and not bvec_bval_done:
            bvec = nii.replace(".nii.gz", ".bvec")
            bval = nii.replace(".nii.gz", ".bval")
            check_file(bvec)
            check_file(bval)

            bvec_new = imgFN.replace(".nii.gz", ".bvec")
            bval_new = imgFN.replace(".nii.gz", ".bval")

            saydo("cp %s %s" % (bvec, bvec_new))
            saydo("cp %s %s" % (bval, bval_new))

            bvec_bval_done = True

    # === Remove temporary directory === #
    saydo("rm -rf %s" % tmpDir)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
# © 2015 Grupo ESOC Ingeniería de Servicios, S.L.U.

from datetime import datetime, timedelta
from pytz import timezone
from openerp import _, api, fields, models
from openerp.exceptions import ValidationError

class Generator(models.TransientModel):
    _name = "event.track.generator"

    event_id = fields.Many2one(
        "event.event",
        string="Event",
        default=lambda self: self.env.context["active_id"],
        required=True)
    event_date_begin = fields.Datetime(
        related="event_id.date_begin",
        help="Change it in the event form. "
             "No tracks before this date will be generated.")
    event_date_end = fields.Datetime(
        related="event_id.date_end",
        help="Change it in the event form. "
             "No tracks after this date will be generated.")
    event_date_tz = fields.Selection(
        related="event_id.date_tz",
        help="Change it in the event form. Timezone of the generated tracks.")
    name = fields.Char(
        "Track title",
        required=True,
        help="Title that will be assigned to all created tracks.")
    location_id = fields.Many2one("event.track.location", "Location")
    speaker_ids = fields.Many2many(
        "res.partner",
        string="Speakers",
        relation="event_track_generate_generator_speaker_ids")
    tag_ids = fields.Many2many(
        "event.track.tag",
        string="Tags",
        relation="event_track_generate_generator_tag_ids")
    mondays = fields.Boolean(help="Create tracks on Mondays.")
    tuesdays = fields.Boolean(help="Create tracks on Tuesdays.")
    wednesdays = fields.Boolean(help="Create tracks on Wednesdays.")
    thursdays = fields.Boolean(help="Create tracks on Thursdays.")
    fridays = fields.Boolean(help="Create tracks on Fridays.")
    saturdays = fields.Boolean(help="Create tracks on Saturdays.")
    sundays = fields.Boolean(help="Create tracks on Sundays.")
    start_time = fields.Float(
        required=True,
        help="Each track will start at this time (in the event's timezone).")
    duration = fields.Float(
        required=True,
        help="Each track will have this duration.")
    end_time = fields.Float(
        compute="_compute_end_time",
        help="Each track will end at this time.")
    delete_existing_tracks = fields.Boolean()
    publish_tracks_in_website = fields.Boolean()
    adjust_start_time = fields.Boolean(
        default=True,
        help="Make event's start time match the start of the first track.")
    adjust_end_time = fields.Boolean(
        default=True,
        help="Make event's end time match the end of the last track.")

    @api.one
    @api.depends("start_time", "duration")
    def _compute_end_time(self):
        self.end_time = self.start_time + self.duration

    @api.one
    def action_generate(self):
        """Generate event tracks according to received data.

        This is the main method of this class, triggered by the UI.
        """
        # You need at least one weekday
        weekdays = self.weekdays()[0]
        if not any(weekdays):
            raise ValidationError(_("You must select at least one weekday."))

        # Delete existing
        if self.delete_existing_tracks:
            self.event_id.track_ids.unlink()

        # Create new
        self.generate_tracks()

        # Adjust event's dates
        self.adjust_dates()

    @api.one
    def adjust_dates(self):
        """Adjust event dates if asked to do so."""
        # Cheek if the user wanted to adjust dates
        if self.event_id.track_ids.exists() and (self.adjust_start_time or
                                                 self.adjust_end_time):
            sorted_ = self.event_id.track_ids.sorted(lambda r: r.date)

            # Start date
            if self.adjust_start_time:
                self.event_id.date_begin = sorted_[0].date

            # End date
            if self.adjust_end_time:
                self.event_id.date_end = fields.Datetime.to_string(
                    fields.Datetime.from_string(sorted_[-1].date) +
                    timedelta(hours=sorted_[-1].duration))

    @api.one
    def create_track(self, **values):
        """Create a new track record with the provided values."""
        data = {
            "name": self.name,
            "event_id": self.event_id.id,
            "duration": self.duration,
            "location_id": self.location_id.id,
            "speaker_ids": [(6, False, self.speaker_ids.ids)],
            "tag_ids": [(6, False, self.tag_ids.ids)],
            "user_id": self.event_id.user_id.id,
            "website_published": self.publish_tracks_in_website}
        data.update(values)
        return self.env["event.track"].create(data)

    @api.one
    def datetime_fields(self):
        """Fields converted to Python's Datetime-based objects."""
        result = {
            "event_start": fields.Datetime.from_string(self.event_date_begin),
            "event_end": fields.Datetime.from_string(self.event_date_end),
            "duration_delta": timedelta(hours=self.duration),
            "day_delta": timedelta(days=1),
            "start_time":
                datetime.min + timedelta(hours=self.start_time),
        }

        # Needed to manually fix timezone offset, for start_time
        result["tzdiff"] = (timezone(self.event_id.date_tz or
                                     self.env.context["tz"] or
                                     self.env.user.tz or
                                     "UTC")
                            .utcoffset(result["event_start"]))
        return result

    @api.one
    def existing_tracks(self, date):
        """Return existing tracks that match some criteria."""
        return self.env["event.track"].search(
            (("event_id", "=", self.event_id.id),
             ("date", "=", date),
             ("duration", "=", self.duration)))

    @api.one
    def generate_tracks(self):
        """Know which tracks must be generated and do it."""
        counter = 0
        dt = self.datetime_fields()[0]
        weekdays = self.weekdays()[0]

        # Check that tracks fit between event start and end dates
        current = dt["event_start"]
        while current <= dt["event_end"]:
            # Get start date and time with fixed timezone offset
            current_start = datetime.combine(
                current.date(),
                dt["start_time"].time()) - dt["tzdiff"]
            if (current_start >= dt["event_start"] and
                    weekdays[current.weekday()]):
                current_end = current_start + dt["duration_delta"]
                if current_end <= dt["event_end"]:
                    # Need string for the ORM
                    current_start = fields.Datetime.to_string(current_start)

                    # Check that no track exists with this data
                    if not self.existing_tracks(current_start)[0]:
                        self.create_track(date=current_start)
                        counter += 1

            # Next day
            current += dt["day_delta"]

    @api.one
    def weekdays(self):
        """Sorted weekdays user selection."""
        return (self.mondays,
                self.tuesdays,
                self.wednesdays,
                self.thursdays,
                self.fridays,
                self.saturdays,
                self.sundays)
<commit_msg>Drop deprecated `@api.one`.<commit_after># -*- coding: utf-8 -*-
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
# © 2015 Grupo ESOC Ingeniería de Servicios, S.L.U.

from datetime import datetime, timedelta
from pytz import timezone
from openerp import _, api, fields, models
from openerp.exceptions import ValidationError

class Generator(models.TransientModel):
    _name = "event.track.generator"

    event_id = fields.Many2one(
        "event.event",
        string="Event",
        default=lambda self: self.env.context["active_id"],
        required=True)
    event_date_begin = fields.Datetime(
        related="event_id.date_begin",
        help="Change it in the event form. "
             "No tracks before this date will be generated.")
    event_date_end = fields.Datetime(
        related="event_id.date_end",
        help="Change it in the event form. "
             "No tracks after this date will be generated.")
    event_date_tz = fields.Selection(
        related="event_id.date_tz",
        help="Change it in the event form. Timezone of the generated tracks.")
    name = fields.Char(
        "Track title",
        required=True,
        help="Title that will be assigned to all created tracks.")
    location_id = fields.Many2one("event.track.location", "Location")
    speaker_ids = fields.Many2many(
        "res.partner",
        string="Speakers",
        relation="event_track_generate_generator_speaker_ids")
    tag_ids = fields.Many2many(
        "event.track.tag",
        string="Tags",
        relation="event_track_generate_generator_tag_ids")
    mondays = fields.Boolean(help="Create tracks on Mondays.")
    tuesdays = fields.Boolean(help="Create tracks on Tuesdays.")
    wednesdays = fields.Boolean(help="Create tracks on Wednesdays.")
    thursdays = fields.Boolean(help="Create tracks on Thursdays.")
    fridays = fields.Boolean(help="Create tracks on Fridays.")
    saturdays = fields.Boolean(help="Create tracks on Saturdays.")
    sundays = fields.Boolean(help="Create tracks on Sundays.")
    start_time = fields.Float(
        required=True,
        help="Each track will start at this time (in the event's timezone).")
    duration = fields.Float(
        required=True,
        help="Each track will have this duration.")
    end_time = fields.Float(
        compute="_compute_end_time",
        help="Each track will end at this time.")
    delete_existing_tracks = fields.Boolean()
    publish_tracks_in_website = fields.Boolean()
    adjust_start_time = fields.Boolean(
        default=True,
        help="Make event's start time match the start of the first track.")
    adjust_end_time = fields.Boolean(
        default=True,
        help="Make event's end time match the end of the last track.")

    @api.multi
    @api.depends("start_time", "duration")
    def _compute_end_time(self):
        self.end_time = self.start_time + self.duration

    @api.multi
    def action_generate(self):
        """Generate event tracks according to received data.

        This is the main method of this class, triggered by the UI.
        """
        # You need at least one weekday
        weekdays = self.weekdays()
        if not any(weekdays):
            raise ValidationError(_("You must select at least one weekday."))

        # Delete existing
        if self.delete_existing_tracks:
            self.event_id.track_ids.unlink()

        # Create new
        self.generate_tracks()

        # Adjust event's dates
        self.adjust_dates()

    @api.multi
    def adjust_dates(self):
        """Adjust event dates if asked to do so."""
        # Cheek if the user wanted to adjust dates
        if self.event_id.track_ids.exists() and (self.adjust_start_time or
                                                 self.adjust_end_time):
            sorted_ = self.event_id.track_ids.sorted(lambda r: r.date)

            # Start date
            if self.adjust_start_time:
                self.event_id.date_begin = sorted_[0].date

            # End date
            if self.adjust_end_time:
                self.event_id.date_end = fields.Datetime.to_string(
                    fields.Datetime.from_string(sorted_[-1].date) +
                    timedelta(hours=sorted_[-1].duration))

    @api.multi
    def create_track(self, **values):
        """Create a new track record with the provided values."""
        data = {
            "name": self.name,
            "event_id": self.event_id.id,
            "duration": self.duration,
            "location_id": self.location_id.id,
            "speaker_ids": [(6, False, self.speaker_ids.ids)],
            "tag_ids": [(6, False, self.tag_ids.ids)],
            "user_id": self.event_id.user_id.id,
            "website_published": self.publish_tracks_in_website}
        data.update(values)
        return self.env["event.track"].create(data)

    @api.multi
    def datetime_fields(self):
        """Fields converted to Python's Datetime-based objects."""
        result = {
            "event_start": fields.Datetime.from_string(self.event_date_begin),
            "event_end": fields.Datetime.from_string(self.event_date_end),
            "duration_delta": timedelta(hours=self.duration),
            "day_delta": timedelta(days=1),
            "start_time":
                datetime.min + timedelta(hours=self.start_time),
        }

        # Needed to manually fix timezone offset, for start_time
        result["tzdiff"] = (timezone(self.event_id.date_tz or
                                     self.env.context["tz"] or
                                     self.env.user.tz or
                                     "UTC")
                            .utcoffset(result["event_start"]))
        return result

    @api.multi
    def existing_tracks(self, date):
        """Return existing tracks that match some criteria."""
        return self.env["event.track"].search(
            (("event_id", "=", self.event_id.id),
             ("date", "=", date),
             ("duration", "=", self.duration)))

    @api.multi
    def generate_tracks(self):
        """Know which tracks must be generated and do it."""
        counter = 0
        dt = self.datetime_fields()
        weekdays = self.weekdays()

        # Check that tracks fit between event start and end dates
        current = dt["event_start"]
        while current <= dt["event_end"]:
            # Get start date and time with fixed timezone offset
            current_start = datetime.combine(
                current.date(),
                dt["start_time"].time()) - dt["tzdiff"]
            if (current_start >= dt["event_start"] and
                    weekdays[current.weekday()]):
                current_end = current_start + dt["duration_delta"]
                if current_end <= dt["event_end"]:
                    # Need string for the ORM
                    current_start = fields.Datetime.to_string(current_start)

                    # Check that no track exists with this data
                    if not self.existing_tracks(current_start):
                        self.create_track(date=current_start)
                        counter += 1

            # Next day
            current += dt["day_delta"]

    @api.multi
    def weekdays(self):
        """Sorted weekdays user selection."""
        return (self.mondays,
                self.tuesdays,
                self.wednesdays,
                self.thursdays,
                self.fridays,
                self.saturdays,
                self.sundays)
<|endoftext|>


<commit_before>import unittest
import numpy as np

import tempfile
import os
import shutil

from neurosynth.tests.utils import get_test_dataset, get_test_data_path
from neurosynth.analysis import meta
from neurosynth.base.dataset import Dataset, ImageTable
from neurosynth.base import imageutils
import json

class TestBase(unittest.TestCase):

    def setUp(self):
        """ Create a new Dataset and add features. """
        self.dataset = get_test_dataset()
        self.skipped_tests = []

    def test_dataset_save_and_load(self):
        # smoke test of saving and loading
        t = tempfile.mktemp()
        self.dataset.save(t, keep_mappables=True)
        self.assertTrue(os.path.exists(t))
        dataset = Dataset.load(t)
        self.assertIsNotNone(dataset)
        self.assertIsNotNone(dataset.mappables)
        self.assertEqual(len(dataset.mappables), 5)
        # Now with the mappables deleted
        dataset.save(t)
        self.assertTrue(os.path.exists(t))
        dataset = Dataset.load(t)
        self.assertEqual(len(dataset.mappables), 0)
        os.unlink(t)

    def test_export_to_json(self):
        js = self.dataset.to_json()
        self.assertTrue(isinstance(js, basestring))
        data = json.loads(js)
        self.assertEqual(len(data['mappables']), 5)

    def test_dataset_initializes(self):
        """ Test whether dataset initializes properly. """
        self.assertIsNotNone(self.dataset.volume)
        self.assertIsNotNone(self.dataset.image_table)
        self.assertEqual(len(self.dataset.mappables), 5)
        self.assertIsNotNone(self.dataset.volume)
        self.assertIsNotNone(self.dataset.r)
        self.assertIsNotNone(self.dataset.mappables[0].data['extra_field'])

    def test_get_mappables(self):
        mappables = self.dataset.get_mappables(['study2', 'study5'])
        self.assertEqual(len(mappables), 2)
        mappables = self.dataset.get_mappables(['study3', 'study4', 'study5'], get_image_data=True)
        self.assertEqual(mappables.shape, (228453, 3))

    def test_image_table_loads(self):
        """ Test ImageTable initialization. """
        self.assertIsNotNone(self.dataset.image_table)
        it = self.dataset.image_table
        self.assertEqual(len(it.ids), 5)
        self.assertIsNotNone(it.volume)
        self.assertIsNotNone(it.r)
        self.assertEqual(it.data.shape, (228453, 5))
        # Add tests for values in table

    def test_feature_table_loads(self):
        """ Test FeatureTable initialization. """
        tt = self.dataset.feature_table
        self.assertIsNotNone(tt)
        self.assertEqual(len(self.dataset.get_feature_names()), 5)
        self.assertEqual(tt.data.shape, (5, 5))
        self.assertEqual(tt.data.columns[3], 'f4')
        self.assertEqual(tt.data.to_dense().iloc[0,0], 0.0003)

    def test_feature_search(self):
        """ Test feature-based Mappable search. Tests both the FeatureTable method
        and the Dataset wrapper. """
        tt = self.dataset.feature_table
        features = tt.search_features(['f*'])
        self.assertEqual(len(features), 4)
        d = self.dataset
        ids = d.get_ids_by_features(['f*'], threshold=0.001)
        self.assertEqual(len(ids), 4)
        img_data = d.get_ids_by_features(
            ['f1', 'f3', 'g1'], 0.001, func=np.max, get_image_data=True)
        self.assertEqual(img_data.shape, (228453, 5))

        # And some smoke-tests:
        # run a meta-analysis
        ma = meta.MetaAnalysis(d, ids)
        # save the results
        tempdir = tempfile.mkdtemp()
        ma.save_results(tempdir + os.path.sep)
        shutil.rmtree(tempdir)

    def test_selection_by_expression(self):
        """ Tests the expression-based search using the lexer/parser.
        Note that this functionality is optional, so only run the test if
        ply is available. """
        try:
            import ply.lex as lex
            ids = self.dataset.get_ids_by_expression("f*", func=np.mean, threshold=0.003)
            self.assertEqual(list(ids), ['study1', 'study3', 'study4'])
            os.unlink('lextab.py')
            os.unlink('parser.out')
            os.unlink('parsetab.py')
        except:
            pass

    def test_selection_by_mask(self):
        """ Test mask-based Mappable selection.
        Only one peak in the test dataset (in study5) should be within the sgACC. """
        ids = self.dataset.get_ids_by_mask(
            get_test_data_path() + 'sgacc_mask.nii.gz')
        self.assertEquals(len(ids), 1)
        self.assertEquals('study5', ids[0])

    def test_unmask(self):
        """ Test unmasking on 1d and 2d vectors (going back to 3d and 4d)

        TODO: test directly on Mask class and its functions, and on
        some smaller example data.  But then it should get into a separate
        TestCase to not 'reload' the same Dataset.
        So for now let's just reuse loaded Dataset and provide
        rudimentary testing
        """
        dataset = self.dataset
        ids = dataset.get_ids_by_mask(
            get_test_data_path() + 'sgacc_mask.nii.gz')
        nvoxels = dataset.volume.in_mask[0].shape[0]

        nvols = 2
        data2d = np.arange(nvoxels * nvols).reshape((nvoxels, -1))

        data2d_unmasked_separately = [
            dataset.volume.unmask(data2d[:, i]) for i in xrange(nvols)]
        data2d_unmasked = dataset.volume.unmask(data2d)
        self.assertEqual(data2d_unmasked.shape,
                         data2d_unmasked_separately[0].shape + (nvols,))
        # and check corresponding volumes
        for i in xrange(nvols):
            self.assertTrue(np.all(data2d_unmasked[
                            ..., i] == data2d_unmasked_separately[i]))

    def test_selection_by_peaks(self):
        """ Test peak-based Mappable selection. """
        ids = self.dataset.get_ids_by_peaks(np.array(
            [[3, 30, -9]]))  # Test with numpy array
        ids2 = self.dataset.get_ids_by_peaks(
            [[3, 30, -9]])  # Test with list of lists
        self.assertEquals(ids, ids2)
        self.assertEquals(len(ids), 1)
        self.assertEquals('study5', ids[0])
        # Repeat with list of lists

    def test_get_feature_counts(self):
        # If we set threshold too high -- nothing should get through and
        # all should be 0s
        feature_counts = self.dataset.get_feature_counts(threshold=1.)
        self.assertEqual(feature_counts,
                         dict((f, 0) for f in self.dataset.get_feature_names()))

        feature_counts = self.dataset.get_feature_counts()
        # all should have some non-0 loading with default threshold,
        # otherwise what is the point of having them?
        for f, c in feature_counts.items():
            self.assertGreater(c, 0, "feature %s has no hits" % f)
        # and should be equal to the ones computed directly (we do not do
        # any fancy queries atm), assumes default threshold of 0.001
        feature_counts_ = dict([(feature, np.sum(self.dataset.feature_table.data.to_dense().ix[:, col] > 0.001))
                               for col, feature in enumerate(self.dataset.feature_table.feature_names)])
        self.assertEqual(feature_counts, feature_counts_)

    def test_get_feature_data(self):
        """ Test retrieval of Mappable x feature data. """
        feature_data = self.dataset.get_feature_data(ids=['study3', 'study1'])
        self.assertEqual(feature_data.shape, (2,5))
        feature_data = self.dataset.get_feature_data(features=['f1', 'f4', 'g1'], dense=True)
        self.assertEqual(feature_data.shape, (5,3))

    def test_get_image_data(self):
        """ Test retrieval of voxel x Mappable data. """
        image_data = self.dataset.get_image_data(voxels=range(2000,3000))
        self.assertEquals(image_data.shape, (1000,5))
        image_data = self.dataset.get_image_data(ids=['study1','study2','study5'], dense=True)
        self.assertEquals(image_data.shape, (228453,3))
        # Or directly through the image table
        image_data = self.dataset.image_table.get_image_data(ids=['study1','study2','study5'], dense=False)
        self.assertEquals(image_data.shape, (228453,3))

    def test_trim_image_table(self):
        self.dataset.image_table.trim(['study1', 'study2', 'study3'])
        self.assertEqual(self.dataset.get_image_data().shape, (228453, 3))

    def test_either_dataset_or_manual_image_vars(self):
        with self.assertRaises(AssertionError):
            image_table = ImageTable()

    def test_img_to_json(self):
        path = get_test_data_path() + 'sgacc_mask.nii.gz'
        json_data = imageutils.img_to_json(path)
        data = json.loads(json_data)
        self.assertEqual(data['max'], 1)
        self.assertEqual(data['dims'], [91, 109, 91])
        self.assertEqual(data['values'], [1.0])
        self.assertEqual(len(data['indices'][0]), 1142)

    def test_get_features_by_ids(self):
        features = self.dataset.feature_table.get_features_by_ids(['study1','study3'], threshold=0.01)
        self.assertEquals(len(features), 3)
        features = self.dataset.feature_table.get_features_by_ids(['study2','study5'], func=np.sum,
            threshold=0.0, get_weights=True)
        self.assertEquals(len(features), 5)
        self.assertEqual(features['f3'], 0.01)

    def test_get_feature_names(self):
        features = self.dataset.get_feature_names()
        self.assertEquals(len(features), 5)
        self.assertEquals(features[2], 'f3')
        features = self.dataset.get_feature_names(['g1', 'f4', 'f2'])
        self.assertEquals(features, ['f2', 'f4', 'g1'])

    def test_grid_creation(self):
        mask = self.dataset.volume.volume
        # Test with mask
        grid = imageutils.create_grid(image=mask, scale=4, mask=mask)
        self.assertEquals(grid.shape, (91, 109, 91))
        self.assertEquals(len(np.unique(grid.get_data())), 4359)
        # Test without mask
        grid = imageutils.create_grid(image=mask, scale=4)
        self.assertGreater(len(np.unique(grid.get_data())), 4359)

suite = unittest.TestLoader().loadTestsFromTestCase(TestBase)

if __name__ == '__main__':
    unittest.main()
<commit_msg>improved lexer test<commit_after>import unittest
import numpy as np

import tempfile
import os
import shutil

from neurosynth.tests.utils import get_test_dataset, get_test_data_path
from neurosynth.analysis import meta
from neurosynth.base.dataset import Dataset, ImageTable
from neurosynth.base import imageutils
import json

class TestBase(unittest.TestCase):

    def setUp(self):
        """ Create a new Dataset and add features. """
        self.dataset = get_test_dataset()
        self.skipped_tests = []

    def test_dataset_save_and_load(self):
        # smoke test of saving and loading
        t = tempfile.mktemp()
        self.dataset.save(t, keep_mappables=True)
        self.assertTrue(os.path.exists(t))
        dataset = Dataset.load(t)
        self.assertIsNotNone(dataset)
        self.assertIsNotNone(dataset.mappables)
        self.assertEqual(len(dataset.mappables), 5)
        # Now with the mappables deleted
        dataset.save(t)
        self.assertTrue(os.path.exists(t))
        dataset = Dataset.load(t)
        self.assertEqual(len(dataset.mappables), 0)
        os.unlink(t)

    def test_export_to_json(self):
        js = self.dataset.to_json()
        self.assertTrue(isinstance(js, basestring))
        data = json.loads(js)
        self.assertEqual(len(data['mappables']), 5)

    def test_dataset_initializes(self):
        """ Test whether dataset initializes properly. """
        self.assertIsNotNone(self.dataset.volume)
        self.assertIsNotNone(self.dataset.image_table)
        self.assertEqual(len(self.dataset.mappables), 5)
        self.assertIsNotNone(self.dataset.volume)
        self.assertIsNotNone(self.dataset.r)
        self.assertIsNotNone(self.dataset.mappables[0].data['extra_field'])

    def test_get_mappables(self):
        mappables = self.dataset.get_mappables(['study2', 'study5'])
        self.assertEqual(len(mappables), 2)
        mappables = self.dataset.get_mappables(['study3', 'study4', 'study5'], get_image_data=True)
        self.assertEqual(mappables.shape, (228453, 3))

    def test_image_table_loads(self):
        """ Test ImageTable initialization. """
        self.assertIsNotNone(self.dataset.image_table)
        it = self.dataset.image_table
        self.assertEqual(len(it.ids), 5)
        self.assertIsNotNone(it.volume)
        self.assertIsNotNone(it.r)
        self.assertEqual(it.data.shape, (228453, 5))
        # Add tests for values in table

    def test_feature_table_loads(self):
        """ Test FeatureTable initialization. """
        tt = self.dataset.feature_table
        self.assertIsNotNone(tt)
        self.assertEqual(len(self.dataset.get_feature_names()), 5)
        self.assertEqual(tt.data.shape, (5, 5))
        self.assertEqual(tt.data.columns[3], 'f4')
        self.assertEqual(tt.data.to_dense().iloc[0,0], 0.0003)

    def test_feature_search(self):
        """ Test feature-based Mappable search. Tests both the FeatureTable method
        and the Dataset wrapper. """
        tt = self.dataset.feature_table
        features = tt.search_features(['f*'])
        self.assertEqual(len(features), 4)
        d = self.dataset
        ids = d.get_ids_by_features(['f*'], threshold=0.001)
        self.assertEqual(len(ids), 4)
        img_data = d.get_ids_by_features(
            ['f1', 'f3', 'g1'], 0.001, func=np.max, get_image_data=True)
        self.assertEqual(img_data.shape, (228453, 5))

        # And some smoke-tests:
        # run a meta-analysis
        ma = meta.MetaAnalysis(d, ids)
        # save the results
        tempdir = tempfile.mkdtemp()
        ma.save_results(tempdir + os.path.sep)
        shutil.rmtree(tempdir)

    def test_selection_by_expression(self):
        """ Tests the expression-based search using the lexer/parser.
        Note that this functionality is optional, so only run the test if
        ply is available. """
        try:
            import ply.lex as lex
            have_ply = True
        except:
            have_ply = False

        if have_ply:
            ids = self.dataset.get_ids_by_expression("* &~ (g*)", func=np.sum, threshold=0.003)
            self.assertEqual(list(ids), ['study3', 'study5'])
            ids = self.dataset.get_ids_by_expression("f* > 0.002", func=np.mean, threshold=0.003)
            self.assertEqual(list(ids), ['study1', 'study3', 'study4'])
            ids = self.dataset.get_ids_by_expression("f* | g*", func=np.mean, threshold=0.003)
            self.assertEqual(list(ids), ['study1', 'study2', 'study3', 'study4'])
            try:
                os.unlink('lextab.py')
                os.unlink('parser.out')
                os.unlink('parsetab.py')
            except: pass

    def test_selection_by_mask(self):
        """ Test mask-based Mappable selection.
        Only one peak in the test dataset (in study5) should be within the sgACC. """
        ids = self.dataset.get_ids_by_mask(
            get_test_data_path() + 'sgacc_mask.nii.gz')
        self.assertEquals(len(ids), 1)
        self.assertEquals('study5', ids[0])

    def test_unmask(self):
        """ Test unmasking on 1d and 2d vectors (going back to 3d and 4d)

        TODO: test directly on Mask class and its functions, and on
        some smaller example data.  But then it should get into a separate
        TestCase to not 'reload' the same Dataset.
        So for now let's just reuse loaded Dataset and provide
        rudimentary testing
        """
        dataset = self.dataset
        ids = dataset.get_ids_by_mask(
            get_test_data_path() + 'sgacc_mask.nii.gz')
        nvoxels = dataset.volume.in_mask[0].shape[0]

        nvols = 2
        data2d = np.arange(nvoxels * nvols).reshape((nvoxels, -1))

        data2d_unmasked_separately = [
            dataset.volume.unmask(data2d[:, i]) for i in xrange(nvols)]
        data2d_unmasked = dataset.volume.unmask(data2d)
        self.assertEqual(data2d_unmasked.shape,
                         data2d_unmasked_separately[0].shape + (nvols,))
        # and check corresponding volumes
        for i in xrange(nvols):
            self.assertTrue(np.all(data2d_unmasked[
                            ..., i] == data2d_unmasked_separately[i]))

    def test_selection_by_peaks(self):
        """ Test peak-based Mappable selection. """
        ids = self.dataset.get_ids_by_peaks(np.array(
            [[3, 30, -9]]))  # Test with numpy array
        ids2 = self.dataset.get_ids_by_peaks(
            [[3, 30, -9]])  # Test with list of lists
        self.assertEquals(ids, ids2)
        self.assertEquals(len(ids), 1)
        self.assertEquals('study5', ids[0])
        # Repeat with list of lists

    def test_get_feature_counts(self):
        # If we set threshold too high -- nothing should get through and
        # all should be 0s
        feature_counts = self.dataset.get_feature_counts(threshold=1.)
        self.assertEqual(feature_counts,
                         dict((f, 0) for f in self.dataset.get_feature_names()))

        feature_counts = self.dataset.get_feature_counts()
        # all should have some non-0 loading with default threshold,
        # otherwise what is the point of having them?
        for f, c in feature_counts.items():
            self.assertGreater(c, 0, "feature %s has no hits" % f)
        # and should be equal to the ones computed directly (we do not do
        # any fancy queries atm), assumes default threshold of 0.001
        feature_counts_ = dict([(feature, np.sum(self.dataset.feature_table.data.to_dense().ix[:, col] > 0.001))
                               for col, feature in enumerate(self.dataset.feature_table.feature_names)])
        self.assertEqual(feature_counts, feature_counts_)

    def test_get_feature_data(self):
        """ Test retrieval of Mappable x feature data. """
        feature_data = self.dataset.get_feature_data(ids=['study3', 'study1'])
        self.assertEqual(feature_data.shape, (2,5))
        feature_data = self.dataset.get_feature_data(features=['f1', 'f4', 'g1'], dense=True)
        self.assertEqual(feature_data.shape, (5,3))

    def test_get_image_data(self):
        """ Test retrieval of voxel x Mappable data. """
        image_data = self.dataset.get_image_data(voxels=range(2000,3000))
        self.assertEquals(image_data.shape, (1000,5))
        image_data = self.dataset.get_image_data(ids=['study1','study2','study5'], dense=True)
        self.assertEquals(image_data.shape, (228453,3))
        # Or directly through the image table
        image_data = self.dataset.image_table.get_image_data(ids=['study1','study2','study5'], dense=False)
        self.assertEquals(image_data.shape, (228453,3))

    def test_trim_image_table(self):
        self.dataset.image_table.trim(['study1', 'study2', 'study3'])
        self.assertEqual(self.dataset.get_image_data().shape, (228453, 3))

    def test_either_dataset_or_manual_image_vars(self):
        with self.assertRaises(AssertionError):
            image_table = ImageTable()

    def test_img_to_json(self):
        path = get_test_data_path() + 'sgacc_mask.nii.gz'
        json_data = imageutils.img_to_json(path)
        data = json.loads(json_data)
        self.assertEqual(data['max'], 1)
        self.assertEqual(data['dims'], [91, 109, 91])
        self.assertEqual(data['values'], [1.0])
        self.assertEqual(len(data['indices'][0]), 1142)

    def test_get_features_by_ids(self):
        features = self.dataset.feature_table.get_features_by_ids(['study1','study3'], threshold=0.01)
        self.assertEquals(len(features), 3)
        features = self.dataset.feature_table.get_features_by_ids(['study2','study5'], func=np.sum,
            threshold=0.0, get_weights=True)
        self.assertEquals(len(features), 5)
        self.assertEqual(features['f3'], 0.01)

    def test_get_feature_names(self):
        features = self.dataset.get_feature_names()
        self.assertEquals(len(features), 5)
        self.assertEquals(features[2], 'f3')
        features = self.dataset.get_feature_names(['g1', 'f4', 'f2'])
        self.assertEquals(features, ['f2', 'f4', 'g1'])

    def test_grid_creation(self):
        mask = self.dataset.volume.volume
        # Test with mask
        grid = imageutils.create_grid(image=mask, scale=4, mask=mask)
        self.assertEquals(grid.shape, (91, 109, 91))
        self.assertEquals(len(np.unique(grid.get_data())), 4359)
        # Test without mask
        grid = imageutils.create_grid(image=mask, scale=4)
        self.assertGreater(len(np.unique(grid.get_data())), 4359)

suite = unittest.TestLoader().loadTestsFromTestCase(TestBase)

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