<commit_before>import contextlib
import uuid
from datetime import datetime
from unittest.mock import MagicMock, Mock, patch

from django.test import TestCase

from corehq.apps.formplayer_api.smsforms.tests.test_formplayer_interface import (
    MockFormplayerRequest,
)
from corehq.apps.sms.api import _process_incoming
from corehq.apps.sms.models import (
    INCOMING,
    SMS,
    MobileBackendInvitation,
    PhoneNumber,
    SQLMobileBackendMapping,
)
from corehq.apps.smsforms.models import (
    SMSChannel,
    SQLXFormsSession,
    XFormsSessionSynchronization,
    get_channel_for_contact,
)
from corehq.messaging.smsbackends.test.models import SQLTestSMSBackend
from corehq.util.test_utils import flag_enabled

QUESTION_RESPONSE = {
    "event": {
        "datatype": "text",
        "choices": None,
        "caption": "Name",
        "type": "question",
        "answer": None,
        "required": 0,
        "ix": "1",
        "help": None
    }
}

ANSWER_RESPONSE = {
    "event": {
        "datatype": None,
        "choices": None,
        "caption": None,
        "type": "form-complete",
        "answer": None,
        "required": 0,
        "ix": "1",
        "help": None,
        "output": "<xml>dummy form xml</xml>"
    }
}

@flag_enabled("ONE_PHONE_NUMBER_MULTIPLE_CONTACTS")
@flag_enabled("INBOUND_SMS_LENIENCY")
@patch('corehq.apps.smsforms.util.critical_section_for_smsforms_sessions',
       new=lambda contact_id: contextlib.nullcontext())
class FormSessionMultipleContactsTestCase(TestCase):

    @classmethod
    def setUpTestData(cls):
        cls.domain_name1 = "domain1"
        cls.domain_name2 = "domain2"

        cls.backend = SQLTestSMSBackend.objects.create(
            name='BACKEND',
            domain=cls.domain_name1,
            is_global=False,
            hq_api_id=SQLTestSMSBackend.get_api_id()
        )

        # share the backend with the 2nd domain
        MobileBackendInvitation.objects.create(domain=cls.domain_name2, backend=cls.backend, accepted=True)

        # set up the backend mapping
        SQLMobileBackendMapping.set_default_domain_backend(cls.domain_name1, cls.backend)
        SQLMobileBackendMapping.set_default_domain_backend(cls.domain_name2, cls.backend)

        cls.phone_number = '01112223333'
        cls.number1 = PhoneNumber.objects.create(
            domain=cls.domain_name1,
            owner_doc_type='CommCareCase',
            owner_id='fake-owner-id1',
            phone_number=cls.phone_number,
            backend_id=None,
            ivr_backend_id=None,
            verified=True,
            is_two_way=True,
            pending_verification=False,
            contact_last_modified=datetime.utcnow()
        )

        cls.number2 = PhoneNumber.objects.create(
            domain=cls.domain_name2,
            owner_doc_type='CommCareCase',
            owner_id='fake-owner-id2',
            phone_number=cls.phone_number,
            backend_id=None,
            ivr_backend_id=None,
            verified=True,
            is_two_way=False,
            pending_verification=False,
            contact_last_modified=datetime.utcnow()
        )

    def setUp(self):
        get_channel_for_contact.clear(self.number1.owner_id, self.number1.phone_number)
        get_channel_for_contact.clear(self.number2.owner_id, self.number2.phone_number)

    @patch('corehq.apps.sms.api.get_location_id_by_verified_number', MagicMock(return_value=None))
    @patch('corehq.apps.sms.api._domain_accepts_inbound', MagicMock(return_value=True))
    @patch('corehq.apps.sms.api.domain_has_privilege', MagicMock(return_value=True))
    @patch('corehq.apps.sms.api._allow_load_handlers', MagicMock(return_value=True))
    @patch('corehq.apps.smsforms.util.submit_form_locally', MagicMock(
        return_value=Mock(xform=Mock(form_id="123"))))
    def test_sms_form_session_across_multiple_domains_with_same_backend_and_number(self):
        self._test(self.number1)  # test session from domain1
        self._test(self.number2)  # test session from domain2 immediately after

    def _test(self, number):
        session = self._make_session(number)
        self._claim_channel(session)

        channel = session.get_channel()

        mocker = MockFormplayerRequest("current", QUESTION_RESPONSE)
        mocker.add_action("answer", ANSWER_RESPONSE)
        msg = self._get_sms("+" + number.phone_number)
        with mocker:
            # with the mocks configured this should result in the session being completed
            _process_incoming(msg)

        # message gets attached to the correct session
        self.assertEqual(msg.xforms_session_couch_id, session.couch_id)

        # confirm session closed
        session.refresh_from_db()
        self.assertFalse(session.session_is_open)
        self.assertTrue(session.completed)
        self.assertEqual(session.submission_id, "123")

        # sticky session is released
        session_info = XFormsSessionSynchronization.get_running_session_info_for_channel(channel)

        self.assertIsNone(session_info.session_id)
        self.assertEqual(session_info.contact_id, session.connection_id)

    def _claim_channel(self, session):
        self.assertEqual(session.get_channel(), SMSChannel(
            backend_id=self.backend.couch_id, phone_number=self.phone_number))

        self.assertTrue(XFormsSessionSynchronization.channel_is_available_for_session(session))
        XFormsSessionSynchronization.claim_channel_for_session(session)

    def _make_session(self, number):
        session = SQLXFormsSession.create_session_object(
            number.domain,
            Mock(get_id=number.owner_id),
            number.phone_number,
            Mock(get_id='app_id'),
            Mock(xmlns='xmlns'),
            expire_after=24 * 60,
        )
        session.session_id = uuid.uuid4().hex
        session.save()
        return session

    def _get_sms(self, phone_number):
        msg = SMS(
            phone_number=phone_number,
            direction=INCOMING,
            date=datetime.utcnow(),
            text="test message",
            domain_scope=self.backend.domain,
            backend_api=self.backend.get_api_id(),
            backend_id=self.backend.couch_id,
            backend_message_id=None,
        )
        return msg
<commit_msg>add option to add the '+' prefix to the SMS number<commit_after>import contextlib
import uuid
from datetime import datetime
from unittest.mock import MagicMock, Mock, patch

from django.test import TestCase

from corehq.apps.formplayer_api.smsforms.tests.test_formplayer_interface import (
    MockFormplayerRequest,
)
from corehq.apps.sms.api import _process_incoming
from corehq.apps.sms.models import (
    INCOMING,
    SMS,
    MobileBackendInvitation,
    PhoneNumber,
    SQLMobileBackendMapping,
)
from corehq.apps.smsforms.models import (
    SMSChannel,
    SQLXFormsSession,
    XFormsSessionSynchronization,
    get_channel_for_contact,
)
from corehq.messaging.smsbackends.test.models import SQLTestSMSBackend
from corehq.util.test_utils import flag_enabled

QUESTION_RESPONSE = {
    "event": {
        "datatype": "text",
        "choices": None,
        "caption": "Name",
        "type": "question",
        "answer": None,
        "required": 0,
        "ix": "1",
        "help": None
    }
}

ANSWER_RESPONSE = {
    "event": {
        "datatype": None,
        "choices": None,
        "caption": None,
        "type": "form-complete",
        "answer": None,
        "required": 0,
        "ix": "1",
        "help": None,
        "output": "<xml>dummy form xml</xml>"
    }
}

@flag_enabled("ONE_PHONE_NUMBER_MULTIPLE_CONTACTS")
@flag_enabled("INBOUND_SMS_LENIENCY")
@patch('corehq.apps.smsforms.util.critical_section_for_smsforms_sessions',
       new=lambda contact_id: contextlib.nullcontext())
class FormSessionMultipleContactsTestCase(TestCase):

    @classmethod
    def setUpTestData(cls):
        cls.domain_name1 = "domain1"
        cls.domain_name2 = "domain2"

        cls.backend = SQLTestSMSBackend.objects.create(
            name='BACKEND',
            domain=cls.domain_name1,
            is_global=False,
            hq_api_id=SQLTestSMSBackend.get_api_id()
        )

        # share the backend with the 2nd domain
        MobileBackendInvitation.objects.create(domain=cls.domain_name2, backend=cls.backend, accepted=True)

        # set up the backend mapping
        SQLMobileBackendMapping.set_default_domain_backend(cls.domain_name1, cls.backend)
        SQLMobileBackendMapping.set_default_domain_backend(cls.domain_name2, cls.backend)

        cls.phone_number = '01112223333'
        cls.number1 = PhoneNumber.objects.create(
            domain=cls.domain_name1,
            owner_doc_type='CommCareCase',
            owner_id='fake-owner-id1',
            phone_number=cls.phone_number,
            backend_id=None,
            ivr_backend_id=None,
            verified=True,
            is_two_way=True,
            pending_verification=False,
            contact_last_modified=datetime.utcnow()
        )

        cls.number2 = PhoneNumber.objects.create(
            domain=cls.domain_name2,
            owner_doc_type='CommCareCase',
            owner_id='fake-owner-id2',
            phone_number=cls.phone_number,
            backend_id=None,
            ivr_backend_id=None,
            verified=True,
            is_two_way=False,
            pending_verification=False,
            contact_last_modified=datetime.utcnow()
        )

    def setUp(self):
        get_channel_for_contact.clear(self.number1.owner_id, self.number1.phone_number)
        get_channel_for_contact.clear(self.number2.owner_id, self.number2.phone_number)

    @patch('corehq.apps.sms.api.get_location_id_by_verified_number', MagicMock(return_value=None))
    @patch('corehq.apps.sms.api._domain_accepts_inbound', MagicMock(return_value=True))
    @patch('corehq.apps.sms.api.domain_has_privilege', MagicMock(return_value=True))
    @patch('corehq.apps.sms.api._allow_load_handlers', MagicMock(return_value=True))
    @patch('corehq.apps.smsforms.util.submit_form_locally', MagicMock(
        return_value=Mock(xform=Mock(form_id="123"))))
    def test_sms_form_session_in_primary_domain_with_plus_prefix(self):
        self._test(self.number1, with_prefix=True)

    def _test(self, number, with_prefix):
        session = self._make_session(number)
        self._claim_channel(session)

        channel = session.get_channel()

        mocker = MockFormplayerRequest("current", QUESTION_RESPONSE)
        mocker.add_action("answer", ANSWER_RESPONSE)
        sms_number = number.phone_number
        if with_prefix:
            sms_number = "+" + sms_number
        msg = self._get_sms(sms_number)
        with mocker:
            # with the mocks configured this should result in the session being completed
            _process_incoming(msg)

        # message gets attached to the correct session
        self.assertEqual(msg.xforms_session_couch_id, session.couch_id)

        # confirm session closed
        session.refresh_from_db()
        self.assertFalse(session.session_is_open)
        self.assertTrue(session.completed)
        self.assertEqual(session.submission_id, "123")

        # sticky session is released
        session_info = XFormsSessionSynchronization.get_running_session_info_for_channel(channel)

        self.assertIsNone(session_info.session_id)
        self.assertEqual(session_info.contact_id, session.connection_id)

    def _claim_channel(self, session):
        self.assertEqual(session.get_channel(), SMSChannel(
            backend_id=self.backend.couch_id, phone_number=self.phone_number))

        self.assertTrue(XFormsSessionSynchronization.channel_is_available_for_session(session))
        XFormsSessionSynchronization.claim_channel_for_session(session)

    def _make_session(self, number):
        session = SQLXFormsSession.create_session_object(
            number.domain,
            Mock(get_id=number.owner_id),
            number.phone_number,
            Mock(get_id='app_id'),
            Mock(xmlns='xmlns'),
            expire_after=24 * 60,
        )
        session.session_id = uuid.uuid4().hex
        session.save()
        return session

    def _get_sms(self, phone_number):
        msg = SMS(
            phone_number=phone_number,
            direction=INCOMING,
            date=datetime.utcnow(),
            text="test message",
            domain_scope=self.backend.domain,
            backend_api=self.backend.get_api_id(),
            backend_id=self.backend.couch_id,
            backend_message_id=None,
        )
        return msg
<|endoftext|>


<commit_before>########
# Copyright (c) 2014 GigaSpaces Technologies Ltd. All rights reserved
#
# 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 time

from cinderclient import exceptions as cinder_exc

import cosmo_tester.framework.util as framework_util
from cosmo_tester.framework.testenv import TestCase

class CinderVolumeTestBase(TestCase):

    VOLUME_SIZE = 1
    DEVICE_NAME = '/dev/vdc'

    def setUp(self):
        super(CinderVolumeTestBase, self).setUp()
        _, _, self.cinderclient = self.env.handler.openstack_clients()

class CinderVolumeTest(CinderVolumeTestBase):

    def test_volume_use_existing(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = \
            blueprint_path / 'use-existing-volume-blueprint.yaml'

        volume_name = 'volume-system-test'

        volume = self.cinderclient.volumes.create(size=self.VOLUME_SIZE,
                                                  display_name=volume_name)
        self.addCleanup(self.cinderclient.volumes.delete, volume.id)

        self._wait_for_volume_available(volume)

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'existing_volume_id': volume.id,
            'device_name': self.DEVICE_NAME
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _wait_for_volume_available(self, volume, num_tries=10):
        for _ in range(num_tries):
            v = self.cinderclient.volumes.get(volume.id)
            if v.status == 'available':
                return True
            else:
                time.sleep(2)
        return False

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 2)

        for key, value in nodes_state.items():
            if 'volume' in key:
                self.assertTrue('external_name' in value['runtime_properties'])
                self.assertTrue('external_id' in value['runtime_properties'])
                self.assertTrue('external_type'
                                in value['runtime_properties'])
                self.assertEqual('volume',
                                 value['runtime_properties']['external_type'])
                self.assertEqual(value['state'], 'started')

                volume_id = value['runtime_properties']['external_id']
                volume_created = True
                try:
                    volume = self.cinderclient.volumes.get(volume_id)
                except cinder_exc.NotFound:
                    volume_created = False
                self.assertTrue(volume_created)
                self.assertEqual(self.VOLUME_SIZE, volume.size)
                self.assertEqual(1, len(volume.attachments))
                self.assertEqual(self.DEVICE_NAME,
                                 volume.attachments[0]['device'])

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 2)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 2)
        volume_node_verified = False
        for node in nodes:
            if node.id == 'test_volume':
                self.assertEqual(len(node.relationships), 1)
                self.assertTrue('device_name' in node.properties)
                volume_node_verified = True
        self.assertTrue(volume_node_verified)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
        self._check_executions(deployment)

    def _check_executions(self, deployment):
        executions = self.client.executions.list(deployment_id=deployment.id)

        self.assertEqual(len(executions), 2)

        execution_from_list = executions[0]
        execution_by_id = self.client.executions.get(execution_from_list.id)

        self.assertEqual(execution_from_list.id, execution_by_id.id)
        self.assertEqual(execution_from_list.workflow_id,
                         execution_by_id.workflow_id)
        self.assertEqual(execution_from_list['blueprint_id'],
                         execution_by_id['blueprint_id'])

        events, _ = self.client.events.get(execution_by_id.id)

        self.assertGreater(len(events), 0)

class CinderVolumeFSTest(CinderVolumeTestBase):

    FS_TYPE = 'ext4'
    FS_MOUNT_PATH = '/test-mount'

    def test_volume_file_system(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = blueprint_path / 'file-system-blueprint.yaml'

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'volume_size': self.VOLUME_SIZE,
            'fs_type': self.FS_TYPE,
            'fs_mount_path': self.FS_MOUNT_PATH
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 4)

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 4)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 4)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
<commit_msg>CFY-1850 flake8<commit_after>########
# Copyright (c) 2014 GigaSpaces Technologies Ltd. All rights reserved
#
# 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 time

from cinderclient import exceptions as cinder_exc

from cosmo_tester.framework.testenv import TestCase

class CinderVolumeTestBase(TestCase):

    VOLUME_SIZE = 1
    DEVICE_NAME = '/dev/vdc'

    def setUp(self):
        super(CinderVolumeTestBase, self).setUp()
        _, _, self.cinderclient = self.env.handler.openstack_clients()

class CinderVolumeTest(CinderVolumeTestBase):

    def test_volume_use_existing(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = \
            blueprint_path / 'use-existing-volume-blueprint.yaml'

        volume_name = 'volume-system-test'

        volume = self.cinderclient.volumes.create(size=self.VOLUME_SIZE,
                                                  display_name=volume_name)
        self.addCleanup(self.cinderclient.volumes.delete, volume.id)

        self._wait_for_volume_available(volume)

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'existing_volume_id': volume.id,
            'device_name': self.DEVICE_NAME
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _wait_for_volume_available(self, volume, num_tries=10):
        for _ in range(num_tries):
            v = self.cinderclient.volumes.get(volume.id)
            if v.status == 'available':
                return True
            else:
                time.sleep(2)
        return False

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 2)

        for key, value in nodes_state.items():
            if 'volume' in key:
                self.assertTrue('external_name' in value['runtime_properties'])
                self.assertTrue('external_id' in value['runtime_properties'])
                self.assertTrue('external_type'
                                in value['runtime_properties'])
                self.assertEqual('volume',
                                 value['runtime_properties']['external_type'])
                self.assertEqual(value['state'], 'started')

                volume_id = value['runtime_properties']['external_id']
                volume_created = True
                try:
                    volume = self.cinderclient.volumes.get(volume_id)
                except cinder_exc.NotFound:
                    volume_created = False
                self.assertTrue(volume_created)
                self.assertEqual(self.VOLUME_SIZE, volume.size)
                self.assertEqual(1, len(volume.attachments))
                self.assertEqual(self.DEVICE_NAME,
                                 volume.attachments[0]['device'])

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 2)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 2)
        volume_node_verified = False
        for node in nodes:
            if node.id == 'test_volume':
                self.assertEqual(len(node.relationships), 1)
                self.assertTrue('device_name' in node.properties)
                volume_node_verified = True
        self.assertTrue(volume_node_verified)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
        self._check_executions(deployment)

    def _check_executions(self, deployment):
        executions = self.client.executions.list(deployment_id=deployment.id)

        self.assertEqual(len(executions), 2)

        execution_from_list = executions[0]
        execution_by_id = self.client.executions.get(execution_from_list.id)

        self.assertEqual(execution_from_list.id, execution_by_id.id)
        self.assertEqual(execution_from_list.workflow_id,
                         execution_by_id.workflow_id)
        self.assertEqual(execution_from_list['blueprint_id'],
                         execution_by_id['blueprint_id'])

        events, _ = self.client.events.get(execution_by_id.id)

        self.assertGreater(len(events), 0)

class CinderVolumeFSTest(CinderVolumeTestBase):

    FS_TYPE = 'ext4'
    FS_MOUNT_PATH = '/test-mount'

    def test_volume_file_system(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = blueprint_path / 'file-system-blueprint.yaml'

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'volume_size': self.VOLUME_SIZE,
            'fs_type': self.FS_TYPE,
            'fs_mount_path': self.FS_MOUNT_PATH
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 4)

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 4)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 4)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
<|endoftext|>


<commit_before>########
# Copyright (c) 2014 GigaSpaces Technologies Ltd. All rights reserved
#
# 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 time

from cinderclient import exceptions as cinder_exc

from cosmo_tester.framework.testenv import TestCase

class CinderVolumeTestBase(TestCase):

    VOLUME_SIZE = 1
    DEVICE_NAME = '/dev/vdc'

    def setUp(self):
        super(CinderVolumeTestBase, self).setUp()
        _, _, self.cinderclient = self.env.handler.openstack_clients()

class CinderVolumeTest(CinderVolumeTestBase):

    def test_volume_create_new(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = blueprint_path / 'create-new-volume-blueprint.yaml'

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'volume_size': self.VOLUME_SIZE,
            'device_name': self.DEVICE_NAME
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def test_volume_use_existing(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = blueprint_path / 'use-existing-volume-blueprint.yaml'

        volume_name = 'volume-system-test'

        volume = self.cinderclient.volumes.create(size=self.VOLUME_SIZE,
                                                  display_name=volume_name)
        self.addCleanup(self.cinderclient.volumes.delete, volume.id)

        self._wait_for_volume_available(volume)

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'existing_volume_id': volume.id,
            'device_name': self.DEVICE_NAME
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _wait_for_volume_available(self, volume, num_tries=10):
        for _ in range(num_tries):
            v = self.cinderclient.volumes.get(volume.id)
            if v.status == 'available':
                return True
            else:
                time.sleep(2)
        return False

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 2)

        for key, value in nodes_state.items():
            if 'volume' in key:
                self.assertTrue('external_name' in value['runtime_properties'])
                self.assertTrue('external_id' in value['runtime_properties'])
                self.assertTrue('external_type'
                                in value['runtime_properties'])
                self.assertEqual('volume',
                                 value['runtime_properties']['external_type'])
                self.assertEqual(value['state'], 'started')

                volume_id = value['runtime_properties']['external_id']
                volume_created = True
                try:
                    volume = self.cinderclient.volumes.get(volume_id)
                except cinder_exc.NotFound:
                    volume_created = False
                self.assertTrue(volume_created)
                self.assertEqual(self.VOLUME_SIZE, volume.size)
                self.assertEqual(1, len(volume.attachments))
                self.assertEqual(self.DEVICE_NAME,
                                 volume.attachments[0]['device'])

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 2)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 2)
        volume_node_verified = False
        for node in nodes:
            if node.id == 'test_volume':
                self.assertEqual(len(node.relationships), 1)
                self.assertTrue('device_name' in node.properties)
                volume_node_verified = True
        self.assertTrue(volume_node_verified)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
        self._check_executions(deployment)

    def _check_executions(self, deployment):
        executions = self.client.executions.list(deployment_id=deployment.id)

        self.assertEqual(len(executions), 2)

        execution_from_list = executions[0]
        execution_by_id = self.client.executions.get(execution_from_list.id)

        self.assertEqual(execution_from_list.id, execution_by_id.id)
        self.assertEqual(execution_from_list.workflow_id,
                         execution_by_id.workflow_id)
        self.assertEqual(execution_from_list['blueprint_id'],
                         execution_by_id['blueprint_id'])

        events, _ = self.client.events.get(execution_by_id.id)

        self.assertGreater(len(events), 0)

class CinderVolumeFSTest(CinderVolumeTestBase):

    FS_TYPE = 'ext4'
    FS_MOUNT_PATH = '/test-mount'

    def test_volume_file_system(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = blueprint_path / 'file-system-blueprint.yaml'

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'volume_size': self.VOLUME_SIZE,
            'fs_type': self.FS_TYPE,
            'fs_mount_path': self.FS_MOUNT_PATH
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 4)

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 4)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 4)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
<commit_msg>CFY-1761 fix flake<commit_after>########
# Copyright (c) 2014 GigaSpaces Technologies Ltd. All rights reserved
#
# 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 time

from cinderclient import exceptions as cinder_exc

from cosmo_tester.framework.testenv import TestCase

class CinderVolumeTestBase(TestCase):

    VOLUME_SIZE = 1
    DEVICE_NAME = '/dev/vdc'

    def setUp(self):
        super(CinderVolumeTestBase, self).setUp()
        _, _, self.cinderclient = self.env.handler.openstack_clients()

class CinderVolumeTest(CinderVolumeTestBase):

    def test_volume_create_new(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = \
            blueprint_path / 'create-new-volume-blueprint.yaml'

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'volume_size': self.VOLUME_SIZE,
            'device_name': self.DEVICE_NAME
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def test_volume_use_existing(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = \
            blueprint_path / 'use-existing-volume-blueprint.yaml'

        volume_name = 'volume-system-test'

        volume = self.cinderclient.volumes.create(size=self.VOLUME_SIZE,
                                                  display_name=volume_name)
        self.addCleanup(self.cinderclient.volumes.delete, volume.id)

        self._wait_for_volume_available(volume)

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'existing_volume_id': volume.id,
            'device_name': self.DEVICE_NAME
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _wait_for_volume_available(self, volume, num_tries=10):
        for _ in range(num_tries):
            v = self.cinderclient.volumes.get(volume.id)
            if v.status == 'available':
                return True
            else:
                time.sleep(2)
        return False

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 2)

        for key, value in nodes_state.items():
            if 'volume' in key:
                self.assertTrue('external_name' in value['runtime_properties'])
                self.assertTrue('external_id' in value['runtime_properties'])
                self.assertTrue('external_type'
                                in value['runtime_properties'])
                self.assertEqual('volume',
                                 value['runtime_properties']['external_type'])
                self.assertEqual(value['state'], 'started')

                volume_id = value['runtime_properties']['external_id']
                volume_created = True
                try:
                    volume = self.cinderclient.volumes.get(volume_id)
                except cinder_exc.NotFound:
                    volume_created = False
                self.assertTrue(volume_created)
                self.assertEqual(self.VOLUME_SIZE, volume.size)
                self.assertEqual(1, len(volume.attachments))
                self.assertEqual(self.DEVICE_NAME,
                                 volume.attachments[0]['device'])

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 2)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 2)
        volume_node_verified = False
        for node in nodes:
            if node.id == 'test_volume':
                self.assertEqual(len(node.relationships), 1)
                self.assertTrue('device_name' in node.properties)
                volume_node_verified = True
        self.assertTrue(volume_node_verified)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
        self._check_executions(deployment)

    def _check_executions(self, deployment):
        executions = self.client.executions.list(deployment_id=deployment.id)

        self.assertEqual(len(executions), 2)

        execution_from_list = executions[0]
        execution_by_id = self.client.executions.get(execution_from_list.id)

        self.assertEqual(execution_from_list.id, execution_by_id.id)
        self.assertEqual(execution_from_list.workflow_id,
                         execution_by_id.workflow_id)
        self.assertEqual(execution_from_list['blueprint_id'],
                         execution_by_id['blueprint_id'])

        events, _ = self.client.events.get(execution_by_id.id)

        self.assertGreater(len(events), 0)

class CinderVolumeFSTest(CinderVolumeTestBase):

    FS_TYPE = 'ext4'
    FS_MOUNT_PATH = '/test-mount'

    def test_volume_file_system(self):

        blueprint_path = self.copy_blueprint('openstack-cinder')
        self.blueprint_yaml = blueprint_path / 'file-system-blueprint.yaml'

        inputs = {
            'image_name': self.env.ubuntu_image_name,
            'flavor_name': self.env.flavor_name,
            'volume_size': self.VOLUME_SIZE,
            'fs_type': self.FS_TYPE,
            'fs_mount_path': self.FS_MOUNT_PATH
        }

        before, after = self.upload_deploy_and_execute_install(
            inputs=inputs
        )

        self._post_install_assertions(before, after)

        self.execute_uninstall()

        self._post_uninstall_assertions()

    def _post_install_assertions(self, before_state, after_state):
        delta = self.get_manager_state_delta(before_state, after_state)

        self.assertEqual(len(delta['deployment_nodes']), 1)

        self.assertEqual(len(delta['node_state']), 1)

        self._check_nodes(delta)
        self._check_blueprint(delta)
        self._check_deployment(delta)

        nodes_state = delta['node_state'].values()[0]
        self.assertEqual(len(nodes_state), 4)

    def _post_uninstall_assertions(self):
        nodes_instances = self.client.node_instances.list(self.deployment_id)
        self.assertEqual(len([node_ins for node_ins in nodes_instances if
                              node_ins.state != 'deleted']), 0)

    def _check_nodes(self, delta):
        self.assertEqual(len(delta['nodes']), 4)
        deployment = delta['deployments'].values()[0]
        nodes = self.client.nodes.list(deployment.id)
        self.assertEqual(len(nodes), 4)

    def _check_blueprint(self, delta):
        self.assertEqual(len(delta['blueprints']), 1)

    def _check_deployment(self, delta):
        self.assertEqual(len(delta['deployments']), 1)
        deployment_from_list = delta['deployments'].values()[0]
        deployment = self.client.deployments.get(deployment_from_list.id)
        self.assertEqual(deployment_from_list.id, deployment.id)
        self.deployment_id = deployment_from_list.id
<|endoftext|>


<commit_before>#!/usr/bin/env python3
###############################################################################
 #  Copyright (c) 2015 Jamis Hoo
 #  Distributed under the MIT license
 #  (See accompanying file LICENSE or copy at http://opensource.org/licenses/MIT)
 #
 #  Project:
 #  Filename: computing_node.py
 #  Version: 1.0
 #  Author: Jamis Hoo
 #  E-mail: hoojamis@gmail.com
 #  Date: Aug  2, 2015
 #  Time: 08:32:02
 #  Description:
###############################################################################

import multiprocessing

from new_process import *

search_index_path = "index"

local_addr = "0.0.0.0"
tcp_port = 10000
http_ports = [ 10001, 10002 , 10003, 10004 ]

master_addr = "localhost"
master_port = 10005

if __name__ == "__main__":
    SearchEngine.load_index(search_index_path)

    # processes set
    processes = []

    # TCP server
    processes.append(multiprocessing.Process(target = new_tcp_process, \
                                             args = (local_addr, tcp_port)))
    processes[-1].start()

    # HTTP servers
    for http_port in http_ports:
        processes.append(multiprocessing.Process(target = new_http_process, \
                                                 args = (local_addr, http_port)))
        processes[-1].start()

    # master detecting process
    processes.append(multiprocessing.Process(target = new_master_detect_process, \
                                             args = (master_addr, master_port, local_addr, http_ports, tcp_port)))
    processes[-1].start()

    for process in processes:
        try:
            process.join()
        except KeyboardInterrupt:
            process.join()
<commit_msg>update port<commit_after>#!/usr/bin/env python3
###############################################################################
 #  Copyright (c) 2015 Jamis Hoo
 #  Distributed under the MIT license
 #  (See accompanying file LICENSE or copy at http://opensource.org/licenses/MIT)
 #
 #  Project:
 #  Filename: computing_node.py
 #  Version: 1.0
 #  Author: Jamis Hoo
 #  E-mail: hoojamis@gmail.com
 #  Date: Aug  2, 2015
 #  Time: 08:32:02
 #  Description:
###############################################################################

import multiprocessing

from new_process import *

search_index_path = "index"

local_addr = "0.0.0.0"
tcp_port = 10000
http_ports = [ 10001, 10002 , 10003, 10004 ]

master_addr = "localhost"
master_port = 3000

if __name__ == "__main__":
    SearchEngine.load_index(search_index_path)

    # processes set
    processes = []

    # TCP server
    processes.append(multiprocessing.Process(target = new_tcp_process, \
                                             args = (local_addr, tcp_port)))
    processes[-1].start()

    # HTTP servers
    for http_port in http_ports:
        processes.append(multiprocessing.Process(target = new_http_process, \
                                                 args = (local_addr, http_port)))
        processes[-1].start()

    # master detecting process
    processes.append(multiprocessing.Process(target = new_master_detect_process, \
                                             args = (master_addr, master_port, local_addr, http_ports, tcp_port)))
    processes[-1].start()

    for process in processes:
        try:
            process.join()
        except KeyboardInterrupt:
            process.join()
<|endoftext|>


<commit_before>import numpy as np
from tabulate import tabulate
from dama.utils.logger import log_config

log = log_config(__name__)

def greater_is_better_fn(reverse, output):
    def view(fn):
        fn.reverse = reverse
        fn.output = output
        return fn
    return view

class MeasureBase(object):
    __slots__ = ['name', 'measures']

    def __init__(self, name: str = None):
        self.name = name
        self.measures = []

    def __iter__(self) -> iter:
        return iter(self.measures)

    def scores(self):
        return []

    def add(self, measure, greater_is_better: bool = True, output=None) -> None:
        self.measures.append(greater_is_better_fn(greater_is_better, output)(measure))

    def outputs(self):
        groups = {}
        for measure in self.measures:
            groups[str(measure.output)] = measure.output
        return groups.values()

    def to_list(self) -> 'ListMeasure':
        list_measure = ListMeasure(headers=[""]+[fn.__name__ for fn in self.measures],
                                   order=[True]+[fn.reverse for fn in self.measures],
                                   measures=[[self.name] + list(self.scores())])
        return list_measure

    @staticmethod
    def base_metrics(measure_cls, measures: str = None, discrete: bool = True):
        if measures is None and discrete is True:
            measure_cls.add(accuracy, greater_is_better=True, output='discrete')
            measure_cls.add(precision, greater_is_better=True, output='discrete')
            measure_cls.add(recall, greater_is_better=True, output='discrete')
            measure_cls.add(f1, greater_is_better=True, output='discrete')
            measure_cls.add(auc, greater_is_better=True, output='discrete')
            measure_cls.add(logloss, greater_is_better=False, output='n_dim')
        elif measures is None and discrete is False:
            measure_cls.add(mse, greater_is_better=False, output='n_dim')
            measure_cls.add(msle, greater_is_better=False, output='n_dim')
            measure_cls.add(gini_normalized, greater_is_better=True, output='n_dim')
        elif isinstance(measures, str):
            import sys
            m = sys.modules['ml.measures']
            if hasattr(m, measures) and measures != 'logloss':
                measure_cls.add(getattr(m, measures), greater_is_better=True,
                                output='discrete')
            elif measures == 'logloss':
                measure_cls.add(getattr(m, measures), greater_is_better=False,
                                output='n_dim')
        return measure_cls

class Measure(MeasureBase):
    __slots__ = ['score', 'name', 'measures']

    def __init__(self, name: str = None):
        super(Measure, self).__init__(name=name)
        self.score = {}

    def update(self, predictions, target) -> None:
        for measure_fn in self:
            self.update_fn(predictions, target, measure_fn)

    def update_fn(self, predictions, target, measure_fn) -> None:
        self.score[measure_fn.__name__] = measure_fn(target, predictions)

    def scores(self):
        for measure in self.measures:
            yield self.score[measure.__name__]

    def make_metrics(self, measures=None, discrete: bool = True) -> 'Measure':
        return MeasureBase.base_metrics(self, measures=measures, discrete=discrete)

class MeasureBatch(MeasureBase):
    __slots__ = ['score', 'name', 'measures', 'batch_size']

    def __init__(self, name: str = None, batch_size: int = 1):
        super(MeasureBatch, self).__init__(name=name)
        self.score = {}
        self.batch_size = batch_size

    def update(self, predictions, target) -> None:
        for measure_fn in self:
            self.update_fn(predictions, target, measure_fn)

    def update_fn(self, predictions, target, measure_fn) -> None:
        log.debug("Set measure {}".format(measure_fn.__name__))
        target_ = target.batch.to_ndarray()
        predictions_ = predictions.batch.to_ndarray()
        try:
            self.score[measure_fn.__name__][0] += measure_fn(target_, predictions_) * \
                                                  (len(predictions_) / self.batch_size)
            self.score[measure_fn.__name__][1] += (len(predictions_) / self.batch_size)
        except KeyError:
            self.score[measure_fn.__name__] = [measure_fn(target_, predictions_), 1]

    def scores(self):
        for measure in self.measures:
            value, size = self.score[measure.__name__]
            yield value / size

    def make_metrics(self, measures=None, discrete: bool = True) -> 'MeasureBatch':
        return MeasureBase.base_metrics(self, measures=measures, discrete=discrete)

def accuracy(labels, predictions):
    """
    measure for correct predictions, true positives and true negatives.
    """
    from sklearn.metrics import accuracy_score
    return accuracy_score(labels, predictions)

def precision(labels, predictions):
    """
    measure for false positives predictions.
    """
    from sklearn.metrics import precision_score
    return precision_score(labels, predictions, average="macro", pos_label=None)

def recall(labels, predictions):
    """
    measure from false negatives predictions.
    """
    from sklearn.metrics import recall_score
    return recall_score(labels, predictions, average="macro", pos_label=None)

def f1(labels, predictions):
    """
    weighted average presicion and recall
    """
    from sklearn.metrics import f1_score
    return f1_score(labels, predictions, average="macro", pos_label=None)

def auc(labels, predictions):
    """
    area under the curve of the reciver operating characteristic, measure for
    true positives rate and false positive rate
    """
    from sklearn.metrics import roc_auc_score
    try:
        return roc_auc_score(labels, predictions, average="macro")
    except ValueError:
        return None

def confusion_matrix(labels, predictions, base_labels=None):
    from sklearn.metrics import confusion_matrix
    cm = confusion_matrix(labels, predictions, labels=base_labels)
    return cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

def logloss(labels, predictions):
    """
    accuracy by penalising false classifications
    """
    from sklearn.metrics import log_loss
    return log_loss(labels, predictions)

def mse(labels, predictions):
    from sklearn.metrics import mean_squared_error
    return mean_squared_error(labels, predictions)

def msle(labels, predicitions):
    from sklearn.metrics import mean_squared_log_error
    return mean_squared_log_error(labels, predicitions)

def gini(actual, pred):
    if len(actual) == len(pred):
        raise AssertionError()
    actual = np.asarray(actual, dtype=np.float)
    n = actual.shape[0]
    a_s = actual[np.argsort(pred)]
    a_c = a_s.cumsum()
    gini_sum = a_c.sum() / a_s.sum() - (n + 1) / 2.0
    return gini_sum / n

def gini_normalized(a, p):
    if p.ndim == 2:
        p = p[:, 1]  # just pick class 1 if is a binary array
    return gini(a, p) / gini(a, a)

class ListMeasure(object):
    """
    Class for save distincts measures

    :type headers: list
    :param headers: lists of headers

    :type measures: list
    :param measures: list of values

    list_measure = ListMeasure(headers=["classif", "f1"], measures=[["test", 0.5], ["test2", 0.6]])
    """

    def __init__(self, headers=None, measures=None, order=None):
        if headers is None:
            headers = []
        if measures is None:
            measures = [[]]
        if order is None:
            order = [False for _ in headers]

        self.headers = headers
        self.measures = measures
        self.order = order

    def add_measure(self, name: str, value, i: int = 0, reverse: bool = False):
        self.headers.append(name)
        try:
            self.measures[i].append(value)
        except IndexError:
            self.measures.append([])
            self.measures[len(self.measures) - 1].append(value)
        self.order.append(reverse)

    def get_measure(self, name: str):
        return self.measures_to_dict().get(name, None)

    def measures_to_dict(self):
        """
        convert the matrix to a dictionary
        """
        from collections import defaultdict
        measures = defaultdict(dict)
        for i, header in enumerate(self.headers, 0):
            measures[header] = {"values": [], "reverse": self.order[i]}
            for measure in self.measures:
                measures[header]["values"].append(measure[i])
        return measures

    @classmethod
    def dict_to_measures(cls, data_dict):
        headers = data_dict.keys()
        measures = [[v["values"][0] for k, v in data_dict.items()]]
        order = [v["reverse"] for k, v in data_dict.items()]
        return ListMeasure(headers=headers, measures=measures, order=order)

    def to_tabulate(self):
        headers, measures, _ = self.drop_empty_columns()
        return tabulate(measures, headers)

    def __str__(self):
        return self.to_tabulate()

    def empty_columns(self):
        """
        return a set of indexes of empty columns
        """
        empty_cols = {}
        for row in self.measures:
            for i, col in enumerate(row):
                if col is None or col == '':
                    empty_cols.setdefault(i, 0)
                    empty_cols[i] += 1

        return set([col for col, counter in empty_cols.items() if counter == len(self.measures)])

    def drop_empty_columns(self):
        """
        drop empty columns
        """
        empty_columns = self.empty_columns()
        not_empty_columns = [i for i in range(len(self.headers)) if i not in empty_columns]
        if not isinstance(self.measures, np.ndarray):
            measures = np.array(self.measures)
        else:
            measures = self.measures.copy()
        headers = list(self.headers)
        order = list(self.order)
        for counter, c in enumerate(empty_columns):
            del headers[c-counter]
            del order[c-counter]
        measures = measures[:, not_empty_columns]
        return headers, measures, order

    def print_matrix(self, labels):
        for _, measure in enumerate(self.measures):
            print("******")
            print(measure[0])
            print("******")
            print(tabulate(np.c_[labels.T, measure[1]], list(labels)))

    def __add__(self, other):
        for hs, ho in zip(self.headers, other.headers):
            if hs != ho:
                raise Exception("Could not add new headers to the table")

        if len(self.headers) == 0:
            headers = other.headers
            this_measures = []
            other_measures = other.measures
            order = other.order
        elif len(other.headers) == 0:
            headers = self.headers
            this_measures = self.measures
            other_measures = []
            order = self.order
        else:
            diff_len = abs(len(self.headers) - len(other.headers))
            if len(self.headers) < len(other.headers):
                headers = other.headers
                this_measures = [m + ([""] * diff_len) for m in self.measures]
                other_measures = other.measures
                order = other.order
            elif len(self.headers) > len(other.headers):
                headers = self.headers
                this_measures = self.measures
                other_measures = [m + ([""] * diff_len) for m in other.measures]
                order = self.order
            else:
                headers = self.headers
                this_measures = self.measures
                other_measures = other.measures
                order = self.order

        list_measure = ListMeasure(
            headers=headers,
            measures=this_measures+other_measures,
            order=order)
        return list_measure

    @staticmethod
    def exclude_columns(headers, exclude_list) -> list:
        headers_set = set(headers).difference(exclude_list)
        return [header for header in headers if header in headers_set]<commit_msg>fixed assertion error<commit_after>import numpy as np
from tabulate import tabulate
from dama.utils.logger import log_config

log = log_config(__name__)

def greater_is_better_fn(reverse, output):
    def view(fn):
        fn.reverse = reverse
        fn.output = output
        return fn
    return view

class MeasureBase(object):
    __slots__ = ['name', 'measures']

    def __init__(self, name: str = None):
        self.name = name
        self.measures = []

    def __iter__(self) -> iter:
        return iter(self.measures)

    def scores(self):
        return []

    def add(self, measure, greater_is_better: bool = True, output=None) -> None:
        self.measures.append(greater_is_better_fn(greater_is_better, output)(measure))

    def outputs(self):
        groups = {}
        for measure in self.measures:
            groups[str(measure.output)] = measure.output
        return groups.values()

    def to_list(self) -> 'ListMeasure':
        list_measure = ListMeasure(headers=[""]+[fn.__name__ for fn in self.measures],
                                   order=[True]+[fn.reverse for fn in self.measures],
                                   measures=[[self.name] + list(self.scores())])
        return list_measure

    @staticmethod
    def base_metrics(measure_cls, measures: str = None, discrete: bool = True):
        if measures is None and discrete is True:
            measure_cls.add(accuracy, greater_is_better=True, output='discrete')
            measure_cls.add(precision, greater_is_better=True, output='discrete')
            measure_cls.add(recall, greater_is_better=True, output='discrete')
            measure_cls.add(f1, greater_is_better=True, output='discrete')
            measure_cls.add(auc, greater_is_better=True, output='discrete')
            measure_cls.add(logloss, greater_is_better=False, output='n_dim')
        elif measures is None and discrete is False:
            measure_cls.add(mse, greater_is_better=False, output='n_dim')
            measure_cls.add(msle, greater_is_better=False, output='n_dim')
            measure_cls.add(gini_normalized, greater_is_better=True, output='n_dim')
        elif isinstance(measures, str):
            import sys
            m = sys.modules['ml.measures']
            if hasattr(m, measures) and measures != 'logloss':
                measure_cls.add(getattr(m, measures), greater_is_better=True,
                                output='discrete')
            elif measures == 'logloss':
                measure_cls.add(getattr(m, measures), greater_is_better=False,
                                output='n_dim')
        return measure_cls

class Measure(MeasureBase):
    __slots__ = ['score', 'name', 'measures']

    def __init__(self, name: str = None):
        super(Measure, self).__init__(name=name)
        self.score = {}

    def update(self, predictions, target) -> None:
        for measure_fn in self:
            self.update_fn(predictions, target, measure_fn)

    def update_fn(self, predictions, target, measure_fn) -> None:
        self.score[measure_fn.__name__] = measure_fn(target, predictions)

    def scores(self):
        for measure in self.measures:
            yield self.score[measure.__name__]

    def make_metrics(self, measures=None, discrete: bool = True) -> 'Measure':
        return MeasureBase.base_metrics(self, measures=measures, discrete=discrete)

class MeasureBatch(MeasureBase):
    __slots__ = ['score', 'name', 'measures', 'batch_size']

    def __init__(self, name: str = None, batch_size: int = 1):
        super(MeasureBatch, self).__init__(name=name)
        self.score = {}
        self.batch_size = batch_size

    def update(self, predictions, target) -> None:
        for measure_fn in self:
            self.update_fn(predictions, target, measure_fn)

    def update_fn(self, predictions, target, measure_fn) -> None:
        log.debug("Set measure {}".format(measure_fn.__name__))
        target_ = target.batch.to_ndarray()
        predictions_ = predictions.batch.to_ndarray()
        try:
            self.score[measure_fn.__name__][0] += measure_fn(target_, predictions_) * \
                                                  (len(predictions_) / self.batch_size)
            self.score[measure_fn.__name__][1] += (len(predictions_) / self.batch_size)
        except KeyError:
            self.score[measure_fn.__name__] = [measure_fn(target_, predictions_), 1]

    def scores(self):
        for measure in self.measures:
            value, size = self.score[measure.__name__]
            yield value / size

    def make_metrics(self, measures=None, discrete: bool = True) -> 'MeasureBatch':
        return MeasureBase.base_metrics(self, measures=measures, discrete=discrete)

def accuracy(labels, predictions):
    """
    measure for correct predictions, true positives and true negatives.
    """
    from sklearn.metrics import accuracy_score
    return accuracy_score(labels, predictions)

def precision(labels, predictions):
    """
    measure for false positives predictions.
    """
    from sklearn.metrics import precision_score
    return precision_score(labels, predictions, average="macro", pos_label=None)

def recall(labels, predictions):
    """
    measure from false negatives predictions.
    """
    from sklearn.metrics import recall_score
    return recall_score(labels, predictions, average="macro", pos_label=None)

def f1(labels, predictions):
    """
    weighted average presicion and recall
    """
    from sklearn.metrics import f1_score
    return f1_score(labels, predictions, average="macro", pos_label=None)

def auc(labels, predictions):
    """
    area under the curve of the reciver operating characteristic, measure for
    true positives rate and false positive rate
    """
    from sklearn.metrics import roc_auc_score
    try:
        return roc_auc_score(labels, predictions, average="macro")
    except ValueError:
        return None

def confusion_matrix(labels, predictions, base_labels=None):
    from sklearn.metrics import confusion_matrix
    cm = confusion_matrix(labels, predictions, labels=base_labels)
    return cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

def logloss(labels, predictions):
    """
    accuracy by penalising false classifications
    """
    from sklearn.metrics import log_loss
    return log_loss(labels, predictions)

def mse(labels, predictions):
    from sklearn.metrics import mean_squared_error
    return mean_squared_error(labels, predictions)

def msle(labels, predicitions):
    from sklearn.metrics import mean_squared_log_error
    return mean_squared_log_error(labels, predicitions)

def gini(actual, pred):
    if len(actual) != len(pred):
        raise AssertionError("The arrays must have the same length")
    actual = np.asarray(actual, dtype=np.float)
    n = actual.shape[0]
    a_s = actual[np.argsort(pred)]
    a_c = a_s.cumsum()
    gini_sum = a_c.sum() / a_s.sum() - (n + 1) / 2.0
    return gini_sum / n

def gini_normalized(a, p):
    if p.ndim == 2:
        p = p[:, 1]  # just pick class 1 if is a binary array
    return gini(a, p) / gini(a, a)

class ListMeasure(object):
    """
    Class for save distincts measures

    :type headers: list
    :param headers: lists of headers

    :type measures: list
    :param measures: list of values

    list_measure = ListMeasure(headers=["classif", "f1"], measures=[["test", 0.5], ["test2", 0.6]])
    """

    def __init__(self, headers=None, measures=None, order=None):
        if headers is None:
            headers = []
        if measures is None:
            measures = [[]]
        if order is None:
            order = [False for _ in headers]

        self.headers = headers
        self.measures = measures
        self.order = order

    def add_measure(self, name: str, value, i: int = 0, reverse: bool = False):
        self.headers.append(name)
        try:
            self.measures[i].append(value)
        except IndexError:
            self.measures.append([])
            self.measures[len(self.measures) - 1].append(value)
        self.order.append(reverse)

    def get_measure(self, name: str):
        return self.measures_to_dict().get(name, None)

    def measures_to_dict(self):
        """
        convert the matrix to a dictionary
        """
        from collections import defaultdict
        measures = defaultdict(dict)
        for i, header in enumerate(self.headers, 0):
            measures[header] = {"values": [], "reverse": self.order[i]}
            for measure in self.measures:
                measures[header]["values"].append(measure[i])
        return measures

    @classmethod
    def dict_to_measures(cls, data_dict):
        headers = data_dict.keys()
        measures = [[v["values"][0] for k, v in data_dict.items()]]
        order = [v["reverse"] for k, v in data_dict.items()]
        return ListMeasure(headers=headers, measures=measures, order=order)

    def to_tabulate(self):
        headers, measures, _ = self.drop_empty_columns()
        return tabulate(measures, headers)

    def __str__(self):
        return self.to_tabulate()

    def empty_columns(self):
        """
        return a set of indexes of empty columns
        """
        empty_cols = {}
        for row in self.measures:
            for i, col in enumerate(row):
                if col is None or col == '':
                    empty_cols.setdefault(i, 0)
                    empty_cols[i] += 1

        return set([col for col, counter in empty_cols.items() if counter == len(self.measures)])

    def drop_empty_columns(self):
        """
        drop empty columns
        """
        empty_columns = self.empty_columns()
        not_empty_columns = [i for i in range(len(self.headers)) if i not in empty_columns]
        if not isinstance(self.measures, np.ndarray):
            measures = np.array(self.measures)
        else:
            measures = self.measures.copy()
        headers = list(self.headers)
        order = list(self.order)
        for counter, c in enumerate(empty_columns):
            del headers[c-counter]
            del order[c-counter]
        measures = measures[:, not_empty_columns]
        return headers, measures, order

    def print_matrix(self, labels):
        for _, measure in enumerate(self.measures):
            print("******")
            print(measure[0])
            print("******")
            print(tabulate(np.c_[labels.T, measure[1]], list(labels)))

    def __add__(self, other):
        for hs, ho in zip(self.headers, other.headers):
            if hs != ho:
                raise Exception("Could not add new headers to the table")

        if len(self.headers) == 0:
            headers = other.headers
            this_measures = []
            other_measures = other.measures
            order = other.order
        elif len(other.headers) == 0:
            headers = self.headers
            this_measures = self.measures
            other_measures = []
            order = self.order
        else:
            diff_len = abs(len(self.headers) - len(other.headers))
            if len(self.headers) < len(other.headers):
                headers = other.headers
                this_measures = [m + ([""] * diff_len) for m in self.measures]
                other_measures = other.measures
                order = other.order
            elif len(self.headers) > len(other.headers):
                headers = self.headers
                this_measures = self.measures
                other_measures = [m + ([""] * diff_len) for m in other.measures]
                order = self.order
            else:
                headers = self.headers
                this_measures = self.measures
                other_measures = other.measures
                order = self.order

        list_measure = ListMeasure(
            headers=headers,
            measures=this_measures+other_measures,
            order=order)
        return list_measure

    @staticmethod
    def exclude_columns(headers, exclude_list) -> list:
        headers_set = set(headers).difference(exclude_list)
        return [header for header in headers if header in headers_set]<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Arne Neumann <arne-neumann@web.de>

"""
This module converts an RS3 XML file (used by RSTTool to annotate
rhetorical structure) into a networkx-based directed graph
(``DiscourseDocumentGraph``).
"""

from __future__ import print_function
import os
from lxml import etree

from discoursegraphs import DiscourseDocumentGraph, EdgeTypes
from discoursegraphs.util import sanitize_string
from discoursegraphs.readwrite.generic import generic_converter_cli

class RSTGraph(DiscourseDocumentGraph):
    """
    A directed graph with multiple edges (based on a networkx
    MultiDiGraph) that represents the rhetorical structure of a
    document.

    Attributes
    ----------
    name : str
        name, ID of the document or file name of the input file
    ns : str
        the namespace of the document (default: rst)
    relations : dict of (str, str)
        dictionary containing all legal RST relations of that file,
        with relation names as keys (str) and relation types
        (either 'rst' or 'multinuc') as values (str).
    segments : list of str
        list of segment node IDs (i.e. leaf nodes in a RST
        tree that represent segments of text).
    root : str
        name of the document root node ID
    tokens : list of str
        sorted list of all token node IDs contained in this document graph
    tokenized : bool
        True (default), if the segments have been tokenized (after they were
        imported from an RS3 file) and have outgoing edges to nodes
        representing tokens.
        False, if the segments represent untokenized text.
    """
    def __init__(self, rs3_filepath, name=None, namespace='rst',
                 tokenize=True, precedence=False):
        """
        Creates an RSTGraph from a RS3 XML file and adds metadata to it
        (filename of the RS3 file, names and types of allowed RST
        relations).

        Parameters
        ----------
        rs3_filepath : str
            absolute or relative path to the RS3 file to be parsed.
        name : str or None
            the name or ID of the graph to be generated. If no name is
            given, the basename of the input file is used.
        namespace : str
            the namespace of the document (default: rst)
        tokenize : bool
            If True, the RST segments (i.e. nuclei and satellites) will
            be tokenized and added as additonal token nodes to the
            document graph (with edges from the respective RST segments).
        precedence : bool
            If True (and if tokenize == True), add precedence relation edges
            (root precedes token1, which precedes token2 etc.)
        """
        # super calls __init__() of base class DiscourseDocumentGraph
        super(RSTGraph, self).__init__()

        self.name = name if name else os.path.basename(rs3_filepath)
        self.ns = namespace
        self.root = None  # __rst2graph() will find/set the root node

        self.segments = []
        self.tokenized = False
        self.tokens = []

        utf8_parser = etree.XMLParser(encoding="utf-8")
        rs3_xml_tree = etree.parse(rs3_filepath, utf8_parser)
        self.relations = extract_relationtypes(rs3_xml_tree)
        self.__rst2graph(rs3_xml_tree)

        if tokenize:
            self.__tokenize_segments()
            self.tokenized = True
            if precedence:
                self.add_precedence_relations()

    def __rst2graph(self, rs3_xml_tree):
        """
        Reads an RST tree (from an ElementTree representation of an RS3
        XML file) and adds all segments (nodes representing text) and
        groups (nonterminal nodes in an RST tree) as well as the
        relationships that hold between them (typed edges) to this
        RSTGraph.

        Parameters
        ----------
        rs3_xml_tree : lxml.etree._ElementTree
            lxml ElementTree representation of an RS3 XML file
        """
        rst_xml_root = rs3_xml_tree.getroot()

        # adds a node to the graph for each RST segment (nucleus or satellite)
        for segment in rst_xml_root.iterfind('./body/segment'):
            segment_node_id = int(segment.attrib['id'])
            self.add_node(
                segment_node_id, layers={self.ns, self.ns+':segment'},
                attr_dict={self.ns+':text': sanitize_string(segment.text)},
                label='{0}:segment:{1}'.format(self.ns, segment.attrib['id']))
            self.segments.append(segment_node_id)

            # adds an edge from the parent node to the segment
            if 'parent' in segment.attrib:
                # node has an outgoing edge,
                # i.e. segment is in an RST relation
                parent_node_id = int(segment.attrib['parent'])
                # if the parent node is not in graph yet, we'll add it first
                if parent_node_id not in self:
                    self.add_node(parent_node_id,
                                  layers={self.ns, self.ns+':segment'})

                segment_rel = self.relations.get(segment.attrib['relname'],
                                                 'span')
                if segment_rel in ('multinuc', 'span'):
                    source_node = parent_node_id
                    target_node = segment_node_id
                else:  # if segment_rel == 'rst'
                    source_node = segment_node_id
                    target_node = parent_node_id

                self.add_edge(
                    source_node, target_node,
                    layers={self.ns, self.ns+':relation'},
                    attr_dict={self.ns+':relname': segment.attrib['relname'],
                               'label': self.ns+':'+segment.attrib['relname']},
                    edge_type=EdgeTypes.dominance_relation)

        for group in rst_xml_root.iterfind('./body/group'):
            group_node_id = int(group.attrib['id'])
            node_type = group.attrib['type']
            if group_node_id in self:  # group node already exists
                self.node[group_node_id].update(
                    {self.ns+':reltype': node_type,
                     'label': '{0}:group:{1}:{2}'.format(self.ns,
                                                         node_type,
                                                         group_node_id)})
            else:
                self.add_node(
                    group_node_id, layers={self.ns, self.ns+':group'},
                    attr_dict={self.ns+':reltype': node_type,
                               'label': '{0}:{1}:{2}'.format(self.ns,
                                                             node_type,
                                                             group_node_id)})

            if 'parent' in group.attrib:
                # node has an outgoing edge, i.e. group is not the
                # topmost element in an RST tree
                parent_node_id = int(group.attrib['parent'])
                if parent_node_id not in self:  # node not in graph, yet
                    self.add_node(
                        parent_node_id, layers={self.ns, self.ns+':group'},
                        label='{0}:{1}'.format(self.ns, parent_node_id))

                group_rel = self.relations.get(group.attrib['relname'], 'span')
                if group_rel in ('multinuc', 'span'):
                    source_node = parent_node_id
                    target_node = group_node_id
                else:  # if group_rel == 'rst'
                    source_node = group_node_id
                    target_node = parent_node_id

                self.add_edge(
                    source_node, target_node,
                    layers={self.ns, self.ns+':relation'},
                    attr_dict={self.ns+':relname': group.attrib['relname'],
                               'label': self.ns+':'+group.attrib['relname']},
                    edge_type=EdgeTypes.dominance_relation)

            else:  # group node is the root of an RST tree
                self.root = group_node_id
                existing_layers = self.node[group_node_id]['layers']
                all_layers = existing_layers.union({self.ns+':root'})
                self.node[group_node_id].update(
                    {'layers': all_layers,
                     'label': '{0}:root:{1}'.format(self.ns, group_node_id)})

    def __tokenize_segments(self):
        """
        tokenizes every RS3 segment (i.e. an RST nucleus or satellite).
        for each token, a node is added to the graph, as well as an edge from
        the segment node to the token node. the token node IDs are also added
        to ``self.tokens``.
        """
        for seg_node_id in self.segments:
            segment_toks = self.node[seg_node_id][self.ns+':text'].split()
            for i, tok in enumerate(segment_toks):
                tok_node_id = '{0}:{1}_{2}'.format(self.ns, seg_node_id, i)
                self.add_node(tok_node_id, layers={self.ns, self.ns+':token'},
                              attr_dict={self.ns+':token': tok, 'label': tok})
                self.tokens.append(tok_node_id)
                self.add_edge(seg_node_id, tok_node_id,
                              layers={'rst', 'rst:token'},
                              edge_type=EdgeTypes.spanning_relation)

    def __str__(self):
        """
        string representation of an RSTGraph (contains filename,
        allowed relations and tokenization status).
        """
        ret_str = '(file) name: {}\n'.format(self.name)
        ret_str += 'number of segments: {}\n'.format(len(self.segments))
        ret_str += 'is tokenized: {}\n'.format(self.tokenized)
        ret_str += 'allowed relations: {}\n'.format(self.relations)
        return ret_str

def extract_relationtypes(rs3_xml_tree):
    """
    extracts the allowed RST relation names and relation types from
    an RS3 XML file.

    Parameters
    ----------
    rs3_xml_tree : lxml.etree._ElementTree
        lxml ElementTree representation of an RS3 XML file

    Returns
    -------
    relations : dict of (str, str)
        Returns a dictionary with RST relation names as keys (str)
        and relation types (either 'rst' or 'multinuc') as values
        (str).
    """
    return {rel.attrib['name']: rel.attrib['type']
            for rel in rs3_xml_tree.iterfind('//header/relations/rel')}

if __name__ == '__main__':
    generic_converter_cli(RSTGraph, 'RST (rhetorical structure)')
<commit_msg>rst: fixed #95, hopefully - parents of segments were created ad-hoc and erroneously labeled as segments<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Arne Neumann <arne-neumann@web.de>

"""
This module converts an RS3 XML file (used by RSTTool to annotate
rhetorical structure) into a networkx-based directed graph
(``DiscourseDocumentGraph``).
"""

from __future__ import print_function
import os
from lxml import etree

from discoursegraphs import DiscourseDocumentGraph, EdgeTypes
from discoursegraphs.util import sanitize_string
from discoursegraphs.readwrite.generic import generic_converter_cli

class RSTGraph(DiscourseDocumentGraph):
    """
    A directed graph with multiple edges (based on a networkx
    MultiDiGraph) that represents the rhetorical structure of a
    document.

    Attributes
    ----------
    name : str
        name, ID of the document or file name of the input file
    ns : str
        the namespace of the document (default: rst)
    relations : dict of (str, str)
        dictionary containing all legal RST relations of that file,
        with relation names as keys (str) and relation types
        (either 'rst' or 'multinuc') as values (str).
    segments : list of str
        list of segment node IDs (i.e. leaf nodes in a RST
        tree that represent segments of text).
    root : str
        name of the document root node ID
    tokens : list of str
        sorted list of all token node IDs contained in this document graph
    tokenized : bool
        True (default), if the segments have been tokenized (after they were
        imported from an RS3 file) and have outgoing edges to nodes
        representing tokens.
        False, if the segments represent untokenized text.
    """
    def __init__(self, rs3_filepath, name=None, namespace='rst',
                 tokenize=True, precedence=False):
        """
        Creates an RSTGraph from a RS3 XML file and adds metadata to it
        (filename of the RS3 file, names and types of allowed RST
        relations).

        Parameters
        ----------
        rs3_filepath : str
            absolute or relative path to the RS3 file to be parsed.
        name : str or None
            the name or ID of the graph to be generated. If no name is
            given, the basename of the input file is used.
        namespace : str
            the namespace of the document (default: rst)
        tokenize : bool
            If True, the RST segments (i.e. nuclei and satellites) will
            be tokenized and added as additonal token nodes to the
            document graph (with edges from the respective RST segments).
        precedence : bool
            If True (and if tokenize == True), add precedence relation edges
            (root precedes token1, which precedes token2 etc.)
        """
        # super calls __init__() of base class DiscourseDocumentGraph
        super(RSTGraph, self).__init__()

        self.name = name if name else os.path.basename(rs3_filepath)
        self.ns = namespace
        self.root = None  # __rst2graph() will find/set the root node

        self.segments = []
        self.tokenized = False
        self.tokens = []

        utf8_parser = etree.XMLParser(encoding="utf-8")
        rs3_xml_tree = etree.parse(rs3_filepath, utf8_parser)
        self.relations = extract_relationtypes(rs3_xml_tree)
        self.__rst2graph(rs3_xml_tree)

        if tokenize:
            self.__tokenize_segments()
            self.tokenized = True
            if precedence:
                self.add_precedence_relations()

    def __rst2graph(self, rs3_xml_tree):
        """
        Reads an RST tree (from an ElementTree representation of an RS3
        XML file) and adds all segments (nodes representing text) and
        groups (nonterminal nodes in an RST tree) as well as the
        relationships that hold between them (typed edges) to this
        RSTGraph.

        Parameters
        ----------
        rs3_xml_tree : lxml.etree._ElementTree
            lxml ElementTree representation of an RS3 XML file
        """
        rst_xml_root = rs3_xml_tree.getroot()

        # adds a node to the graph for each RST segment (nucleus or satellite)
        for segment in rst_xml_root.iterfind('./body/segment'):
            segment_node_id = int(segment.attrib['id'])
            self.add_node(
                segment_node_id, layers={self.ns, self.ns+':segment'},
                attr_dict={self.ns+':text': sanitize_string(segment.text)},
                label='{0}:segment:{1}'.format(self.ns, segment.attrib['id']))
            self.segments.append(segment_node_id)

            # adds an edge from the parent node to the segment
            if 'parent' in segment.attrib:
                # node has an outgoing edge,
                # i.e. segment is in an RST relation
                parent_node_id = int(segment.attrib['parent'])
                # if the parent node is not in graph yet, we'll add it first
                if parent_node_id not in self:
                    self.add_node(parent_node_id,
                                  layers={self.ns, self.ns+':group'})

                segment_rel = self.relations.get(segment.attrib['relname'],
                                                 'span')
                if segment_rel in ('multinuc', 'span'):
                    source_node = parent_node_id
                    target_node = segment_node_id
                else:  # if segment_rel == 'rst'
                    source_node = segment_node_id
                    target_node = parent_node_id

                self.add_edge(
                    source_node, target_node,
                    layers={self.ns, self.ns+':relation'},
                    attr_dict={self.ns+':relname': segment.attrib['relname'],
                               'label': self.ns+':'+segment.attrib['relname']},
                    edge_type=EdgeTypes.dominance_relation)

        for group in rst_xml_root.iterfind('./body/group'):
            group_node_id = int(group.attrib['id'])
            node_type = group.attrib['type']
            if group_node_id in self:  # group node already exists
                self.node[group_node_id].update(
                    {self.ns+':reltype': node_type,
                     'label': '{0}:group:{1}:{2}'.format(self.ns,
                                                         node_type,
                                                         group_node_id)})
            else:
                self.add_node(
                    group_node_id, layers={self.ns, self.ns+':group'},
                    attr_dict={self.ns+':reltype': node_type,
                               'label': '{0}:{1}:{2}'.format(self.ns,
                                                             node_type,
                                                             group_node_id)})

            if 'parent' in group.attrib:
                # node has an outgoing edge, i.e. group is not the
                # topmost element in an RST tree
                parent_node_id = int(group.attrib['parent'])
                if parent_node_id not in self:  # node not in graph, yet
                    self.add_node(
                        parent_node_id, layers={self.ns, self.ns+':group'},
                        label='{0}:{1}'.format(self.ns, parent_node_id))

                group_rel = self.relations.get(group.attrib['relname'], 'span')
                if group_rel in ('multinuc', 'span'):
                    source_node = parent_node_id
                    target_node = group_node_id
                else:  # if group_rel == 'rst'
                    source_node = group_node_id
                    target_node = parent_node_id

                self.add_edge(
                    source_node, target_node,
                    layers={self.ns, self.ns+':relation'},
                    attr_dict={self.ns+':relname': group.attrib['relname'],
                               'label': self.ns+':'+group.attrib['relname']},
                    edge_type=EdgeTypes.dominance_relation)

            else:  # group node is the root of an RST tree
                self.root = group_node_id
                existing_layers = self.node[group_node_id]['layers']
                all_layers = existing_layers.union({self.ns+':root'})
                self.node[group_node_id].update(
                    {'layers': all_layers,
                     'label': '{0}:root:{1}'.format(self.ns, group_node_id)})

    def __tokenize_segments(self):
        """
        tokenizes every RS3 segment (i.e. an RST nucleus or satellite).
        for each token, a node is added to the graph, as well as an edge from
        the segment node to the token node. the token node IDs are also added
        to ``self.tokens``.
        """
        for seg_node_id in self.segments:
            segment_toks = self.node[seg_node_id][self.ns+':text'].split()
            for i, tok in enumerate(segment_toks):
                tok_node_id = '{0}:{1}_{2}'.format(self.ns, seg_node_id, i)
                self.add_node(tok_node_id, layers={self.ns, self.ns+':token'},
                              attr_dict={self.ns+':token': tok, 'label': tok})
                self.tokens.append(tok_node_id)
                self.add_edge(seg_node_id, tok_node_id,
                              layers={'rst', 'rst:token'},
                              edge_type=EdgeTypes.spanning_relation)

    def __str__(self):
        """
        string representation of an RSTGraph (contains filename,
        allowed relations and tokenization status).
        """
        ret_str = '(file) name: {}\n'.format(self.name)
        ret_str += 'number of segments: {}\n'.format(len(self.segments))
        ret_str += 'is tokenized: {}\n'.format(self.tokenized)
        ret_str += 'allowed relations: {}\n'.format(self.relations)
        return ret_str

def extract_relationtypes(rs3_xml_tree):
    """
    extracts the allowed RST relation names and relation types from
    an RS3 XML file.

    Parameters
    ----------
    rs3_xml_tree : lxml.etree._ElementTree
        lxml ElementTree representation of an RS3 XML file

    Returns
    -------
    relations : dict of (str, str)
        Returns a dictionary with RST relation names as keys (str)
        and relation types (either 'rst' or 'multinuc') as values
        (str).
    """
    return {rel.attrib['name']: rel.attrib['type']
            for rel in rs3_xml_tree.iterfind('//header/relations/rel')}

if __name__ == '__main__':
    generic_converter_cli(RSTGraph, 'RST (rhetorical structure)')
<|endoftext|>