<commit_before>"""
Extract metadata for three materials data publication services from a common payload.

Services:
    Citrine:                    <link>
    Materials Data Facility:    <link>
    Materials Commons:          <link>
"""

from pypif import pif
import pypif.obj as pobj
import json
import datetime

def get_common_payload_template():
    """
    Get a template dictionary that can be used to create a payload object.

    TODO: implement
    """
    return {
        'template': {
            'title': 'string',
            'source': {
                'name': 'string',
                'producer': 'string',
                'url': 'url string',
                'tags': ['string']
            },
            'data_contacts': [
                {
                    'given_name': 'string',
                    'family_name': 'string',
                    'title': 'string',
                    'orcid': 'TBD',
                    'email': 'TBD',
                    'tags': ['string']
                }
            ],
            'data_contributors': [
                {
                    'given_name': 'string',
                    'family_name': 'string',
                    'title': 'string',
                    'orcid': 'TBD',
                    'email': 'TBD',
                    'tags': ['string']
                }
            ],
            'links': 'TBD',
            'authors': [
                {
                    'given_name': 'string',
                    'family_name': 'string',
                    'title': 'string',
                    'orcid': 'TBD',
                    'email': 'TBD',
                    'tags': ['string']
                }
            ],
            'licenses': [
                {
                    'name': 'string',
                    'description': 'string',
                    'url': 'string',
                    'tags': ['string']
                }
            ],
            'citations': 'TBD',
            'repository': 'TBD',
            'collection': 'TBD',
            'tags': ['string'],
            'description': 'string',
            'raw': 'TBD',
            'year': 'integer',
            'composition': 'TBD'
        },
        'required': [
            'title',
            'source (name only)',
            'data_contacts',
            'data_contributors',
            'links'
        ],
        'usage': 'payload = <PublishablePayload_subclass>(**input_dictionary); metadata = payload.metapayload'
    }

class PublishablePayload(dict):
    def __init__(self, title, source, data_contacts, data_contributors, links, **kwargs):
        """
        Parameters
        ----------
        title:              str
        source:             dict: { 'name': str, 'producer': str, 'url': str, 'tags': list(str)}
        data_contact:       Human
        data_contributor:   Human
        links:              dict
        """
        super(PublishablePayload, self).__init__()
        self.__dict__ = self

        self._optionalkeys = ['licenses',    # list of dicts: {
                                            #     'name': str,
                                            #     'description': str,
                                            #     'url': str,
                                            #     'tags': list of str
                                            # }
                              'citations',  # list of str
                              'authors',  # list of human
                              'repository',  # str
                              'collection',  # str
                              'tags',  # list of str
                              'description',  # str
                              'raw',  # str
                              'links',  # dict
                              'year',  # int
                              'composition']  # str
        for prop in self._optionalkeys:
            self[prop] = kwargs.get(prop, None)

        # validate the source parameter
        # (require MDF field (name), but not the citrine fields)
        expected_source_keys = ['name',]
        for key in expected_source_keys:
            if key not in source:
                raise Exception('source requires a %s field' % key)

        self['title'] = title
        self['source'] = source
        self['data_contacts'] = data_contacts
        self['data_contributors'] = data_contributors
        self['links'] = links
        self['additionalProperties'] = kwargs

    @property
    def metapayload(self):
        raise NotImplementedError

    def ingest(self):
        raise NotImplementedError

class CITPayload(PublishablePayload):
    """
    Construct payload for POST call to Citrine service.

    Implemented fields (so far) include contacts and source.
    These roughly correspond to the fields required by MDF.

    Examples
    --------
    >>> scripty = Human(given_name='Totally', family_name='NotARobot', email='a@a.com', institution='Earth')
    >>> payload = CITPayload(title='Test Payload', source={'name': 'whatever', 'producer': 'test producer', 'url': 'http://www.testurl.org', 'tags': ['these', 'are', 'source', 'tags']}, data_contacts=[scripty], data_contributors=[scripty], links={'landing_page':'http://www.globus.org'}, licenses=[{'name': 'license name', 'url': 'http://www.licenseurl.org', 'description': 'license description', 'tags': ['license', 'tags']},])
    >>> payload.metapayload
    {'email': 'a@a.com', 'name': {'family': 'NotARobot', 'given': 'Totally', 'title': ''}, 'tags': ['contributor']}], 'licenses': [{'description': 'license description', 'name': 'license name', 'tags': ['license', 'tags'], 'url': 'http://www.licenseurl.org'}], 'source': {'producer': 'test producer', 'tags': ['these', 'are', 'source', 'tags'], 'url': 'http://www.testurl.org'}}

    """

    def __init__(self, *args, **kwargs):
        super(CITPayload, self).__init__(*args, **kwargs)
        self.metadata = pobj.System()
        self._add_source()
        self._add_people()
        self._add_licenses()

    @property
    def metapayload(self):
        return json.loads(pif.dumps(self.metadata))

    def _add_source(self):
        if 'producer' in self['source']:
            producer = self['source']['producer']
        if 'url' in self['source']:
            url = self['source']['url']
        if 'tags' in self['source']:
            tags = self['source']['tags']

        self.metadata.source = pobj.Source(
            producer=producer,
            url=url,
            tags=tags
        )

    def _add_people(self):
        people = []

        def add_to_people(person_list, tags):
            for person in person_list:
                citrine_name_info = {
                    'given': person.get('given_name', ''),
                    'family': person.get('family_name', ''),
                    'title': person.get('title', ''),
                }
                citrine_name = pobj.Name(**citrine_name_info)
                citrine_person_info = {
                    'name': citrine_name,
                    'orcid': person.get('orcid', None),
                    'email': person.get('email', None),
                    'tags': tags
                }
                citrine_person = pobj.Person(**citrine_person_info)
                people.append(citrine_person)

        if 'authors' in self and isinstance(self['authors'], list):
            add_to_people(person_list=self['authors'], tags=['author'])
        if 'data_contacts' in self and isinstance(self['data_contacts'], list):
            add_to_people(person_list=self['data_contacts'], tags=['contact'])
        if 'data_contributors' in self and isinstance(self['data_contributors'], list):
            add_to_people(person_list=self['data_contributors'], tags=['contributor'])

        self.metadata.contacts = people

    def _add_licenses(self):
        if 'licenses' not in self or not isinstance(self['licenses'], list):
            return

        citrine_licenses = []

        for license in self['licenses']:
            try:
                citrine_license = pobj.License(**license)
                citrine_licenses.append(citrine_license)
            except Exception as ex:
                print(ex)

        self.metadata.licenses = citrine_licenses

class MDFPayload(PublishablePayload):
    """

    Examples
    --------
    >>> scripty = Human(given_name='Totally', family_name='NotARobot', email='a@a.com', institution='Earth')
    >>> payload = MDFPayload(title='Test Payload', source={'name': 'Doctest Example Script'}, data_contacts=[scripty], data_contributors=[scripty], links={'landing_page':'http://www.globus.org'})
    >>> payload.metapayload
    {'Doctest Example Script': {}, 'dc': {}, 'mdf': {'acl': ['public'],  'citations': None,  'data_contact': [{'email': 'a@a.com', 'family_name': 'NotARobot', 'given_name': 'Totally', 'institution': 'Earth'}], 'data_contributor': [{'email': 'a@a.com', 'family_name': 'NotARobot', 'given_name': 'Totally', 'institution': 'Earth'}], 'links': {'landing_page': 'http://www.globus.org'}, 'source_name': 'Doctest Example Script', 'title': 'Test Payload'}}

    """

    def __init__(self, *args, **kwargs):
        super(MDFPayload, self).__init__(*args, **kwargs)

    @property
    def metapayload(self):
        dataset = {
            "mdf": {
                "title": self.title,
                "acl": ["public"],  # TODO: allow list of globus auth uuids or "public"
                "source_name": self.source['name'],
                "citations": self.citations,
                "links": self.links,
                "data_contact": self.data_contacts,
                "data_contributor": [dict(data_contributor) for data_contributor in self.data_contributors],
                # "ingest_date": datetime.datetime.now().strftime('%b %d, %Y'),  # Note: these commented keys are
                # "metadata_version": "1.1",                                     # populated automatically!
                # "mdf_id": "1",
                # "resource_type": "dataset",
                # "additionalProperties": self.additionalProperties
            },
            "dc": {},  # TODO: allow datacite keys to go here?
            # TODO: review the following.
            self.source['name']: self.additionalProperties
        }

        # Populate optional keys if they have been set
        for key in self._optionalkeys:
            val = getattr(self, key)
            if val is not None:
                dataset['mdf'][key] = val

        return dataset

class MCPayload(PublishablePayload):
    pass

class Human(dict):
    def __init__(self, given_name, family_name, email='', institution=''):
        super(Human, self).__init__()
        self.__dict__ = self
        self['given_name'] = given_name
        self['family_name'] = family_name
        self['email'] = email
        self['institution'] = institution
<commit_msg>added MCPayload (Materials Commons) implementation<commit_after>"""
Extract metadata for three materials data publication services from a common payload.

Services:
    Citrine:                    <link>
    Materials Data Facility:    <link>
    Materials Commons:          <link>
"""

from pypif import pif
import pypif.obj as pobj
import json
import datetime

def get_common_payload_template():
    """
    Get a template dictionary that can be used to create a payload object.

    TODO: implement
    """
    return {
        'template': {
            'title': 'string',
            'source': {
                'name': 'string',
                'producer': 'string',
                'url': 'url string',
                'tags': ['string']
            },
            'data_contacts': [
                {
                    'given_name': 'string',
                    'family_name': 'string',
                    'title': 'string',
                    'orcid': 'TBD',
                    'email': 'TBD',
                    'tags': ['string']
                }
            ],
            'data_contributors': [
                {
                    'given_name': 'string',
                    'family_name': 'string',
                    'title': 'string',
                    'orcid': 'TBD',
                    'email': 'TBD',
                    'tags': ['string']
                }
            ],
            'links': 'TBD',
            'authors': [
                {
                    'given_name': 'string',
                    'family_name': 'string',
                    'title': 'string',
                    'orcid': 'TBD',
                    'email': 'TBD',
                    'tags': ['string']
                }
            ],
            'licenses': [
                {
                    'name': 'string',
                    'description': 'string',
                    'url': 'string',
                    'tags': ['string']
                }
            ],
            'citations': 'TBD',
            'repository': 'TBD',
            'collection': 'TBD',
            'tags': ['string'],
            'description': 'string',
            'raw': 'TBD',
            'year': 'integer',
            'composition': 'TBD'
        },
        'required': [
            'title',
            'source (name only)',
            'data_contacts',
            'data_contributors',
            'links',
            'description'
        ],
        'usage': 'payload = <PublishablePayload_subclass>(**input_dictionary); metadata = payload.metapayload'
    }

class PublishablePayload(dict):
    def __init__(self, title, source, data_contacts, data_contributors, links, **kwargs):
        """
        Parameters
        ----------
        title:              str
        source:             dict: { 'name': str, 'producer': str, 'url': str, 'tags': list(str)}
        data_contact:       Human
        data_contributor:   Human
        links:              dict
        """
        super(PublishablePayload, self).__init__()
        self.__dict__ = self

        self._optionalkeys = ['licenses',    # list of dicts: {
                                            #     'name': str,
                                            #     'description': str,
                                            #     'url': str,
                                            #     'tags': list of str
                                            # }
                              'citations',  # list of str
                              'authors',  # list of human
                              'repository',  # str
                              'collection',  # str
                              'tags',  # list of str
                              'description',  # str
                              'raw',  # str
                              'links',  # dict
                              'year',  # int
                              'composition']  # str
        for prop in self._optionalkeys:
            self[prop] = kwargs.get(prop, None)

        # validate the source parameter
        # (require MDF field (name), but not the citrine fields)
        expected_source_keys = ['name',]
        for key in expected_source_keys:
            if key not in source:
                raise Exception('source requires a %s field' % key)

        self['title'] = title
        self['source'] = source
        self['data_contacts'] = data_contacts
        self['data_contributors'] = data_contributors
        self['links'] = links
        self['additionalProperties'] = kwargs

    @property
    def metapayload(self):
        raise NotImplementedError

    def ingest(self):
        raise NotImplementedError

class CITPayload(PublishablePayload):
    """
    Construct payload for POST call to Citrine service.

    Implemented fields (so far) include contacts and source.
    These roughly correspond to the fields required by MDF.

    Examples
    --------
    >>> scripty = Human(given_name='Totally', family_name='NotARobot', email='a@a.com', institution='Earth')
    >>> payload = CITPayload(title='Test Payload', source={'name': 'whatever', 'producer': 'test producer', 'url': 'http://www.testurl.org', 'tags': ['these', 'are', 'source', 'tags']}, data_contacts=[scripty], data_contributors=[scripty], links={'landing_page':'http://www.globus.org'}, licenses=[{'name': 'license name', 'url': 'http://www.licenseurl.org', 'description': 'license description', 'tags': ['license', 'tags']},])
    >>> payload.metapayload
    {'email': 'a@a.com', 'name': {'family': 'NotARobot', 'given': 'Totally', 'title': ''}, 'tags': ['contributor']}], 'licenses': [{'description': 'license description', 'name': 'license name', 'tags': ['license', 'tags'], 'url': 'http://www.licenseurl.org'}], 'source': {'producer': 'test producer', 'tags': ['these', 'are', 'source', 'tags'], 'url': 'http://www.testurl.org'}}

    """

    def __init__(self, *args, **kwargs):
        super(CITPayload, self).__init__(*args, **kwargs)
        self.metadata = pobj.System()
        self._add_source()
        self._add_people()
        self._add_licenses()

    @property
    def metapayload(self):
        return json.loads(pif.dumps(self.metadata))

    def _add_source(self):
        if 'producer' in self['source']:
            producer = self['source']['producer']
        if 'url' in self['source']:
            url = self['source']['url']
        if 'tags' in self['source']:
            tags = self['source']['tags']

        self.metadata.source = pobj.Source(
            producer=producer,
            url=url,
            tags=tags
        )

    def _add_people(self):
        people = []

        def add_to_people(person_list, tags):
            for person in person_list:
                citrine_name_info = {
                    'given': person.get('given_name', ''),
                    'family': person.get('family_name', ''),
                    'title': person.get('title', ''),
                }
                citrine_name = pobj.Name(**citrine_name_info)
                citrine_person_info = {
                    'name': citrine_name,
                    'orcid': person.get('orcid', None),
                    'email': person.get('email', None),
                    'tags': tags
                }
                citrine_person = pobj.Person(**citrine_person_info)
                people.append(citrine_person)

        if 'authors' in self and isinstance(self['authors'], list):
            add_to_people(person_list=self['authors'], tags=['author'])
        if 'data_contacts' in self and isinstance(self['data_contacts'], list):
            add_to_people(person_list=self['data_contacts'], tags=['contact'])
        if 'data_contributors' in self and isinstance(self['data_contributors'], list):
            add_to_people(person_list=self['data_contributors'], tags=['contributor'])

        self.metadata.contacts = people

    def _add_licenses(self):
        if 'licenses' not in self or not isinstance(self['licenses'], list):
            return

        citrine_licenses = []

        for license in self['licenses']:
            try:
                citrine_license = pobj.License(**license)
                citrine_licenses.append(citrine_license)
            except Exception as ex:
                print(ex)

        self.metadata.licenses = citrine_licenses

class MDFPayload(PublishablePayload):
    """

    Examples
    --------
    >>> scripty = Human(given_name='Totally', family_name='NotARobot', email='a@a.com', institution='Earth')
    >>> payload = MDFPayload(title='Test Payload', source={'name': 'Doctest Example Script'}, data_contacts=[scripty], data_contributors=[scripty], links={'landing_page':'http://www.globus.org'})
    >>> payload.metapayload
    {'Doctest Example Script': {}, 'dc': {}, 'mdf': {'acl': ['public'],  'citations': None,  'data_contact': [{'email': 'a@a.com', 'family_name': 'NotARobot', 'given_name': 'Totally', 'institution': 'Earth'}], 'data_contributor': [{'email': 'a@a.com', 'family_name': 'NotARobot', 'given_name': 'Totally', 'institution': 'Earth'}], 'links': {'landing_page': 'http://www.globus.org'}, 'source_name': 'Doctest Example Script', 'title': 'Test Payload'}}

    """

    def __init__(self, *args, **kwargs):
        super(MDFPayload, self).__init__(*args, **kwargs)

    @property
    def metapayload(self):
        dataset = {
            "mdf": {
                "title": self.title,
                "acl": ["public"],  # TODO: allow list of globus auth uuids or "public"
                "source_name": self.source['name'],
                "citations": self.citations,
                "links": self.links,
                "data_contact": self.data_contacts,
                "data_contributor": [dict(data_contributor) for data_contributor in self.data_contributors],
                # "ingest_date": datetime.datetime.now().strftime('%b %d, %Y'),  # Note: these commented keys are
                # "metadata_version": "1.1",                                     # populated automatically!
                # "mdf_id": "1",
                # "resource_type": "dataset",
                # "additionalProperties": self.additionalProperties
            },
            "dc": {},  # TODO: allow datacite keys to go here?
            # TODO: review the following.
            self.source['name']: self.additionalProperties
        }

        # Populate optional keys if they have been set
        for key in self._optionalkeys:
            val = getattr(self, key)
            if val is not None:
                dataset['mdf'][key] = val

        return dataset

class MCPayload(PublishablePayload):
    """
    Examples
    --------
    >>> scripty = Human(given_name='Totally', family_name='NotARobot', email='a@a.com', institution='Earth')
    >>> payload = MCPayload(title='Test Payload', source={'name': 'whatever', 'producer': 'test producer', 'url': 'http://www.testurl.org', 'tags': ['these', 'are', 'source', 'tags']}, data_contacts=[scripty], data_contributors=[scripty], links={'landing_page':'http://www.globus.org'}, licenses=[{'name': 'license name', 'url': 'http://www.licenseurl.org', 'description': 'license description', 'tags': ['license', 'tags']},], 'description': 'material description')
    >>> payload.metapayload
    {'email': 'a@a.com', 'name': {'family': 'NotARobot', 'given': 'Totally', 'title': ''}, 'tags': ['contributor']}], 'licenses': [{'description': 'license description', 'name': 'license name', 'tags': ['license', 'tags'], 'url': 'http://www.licenseurl.org'}], 'source': {'producer': 'test producer', 'tags': ['these', 'are', 'source', 'tags'], 'url': 'http://www.testurl.org'}}

    """

    def __init__(self, *args, **kwargs):
        if 'description' not in kwargs:
            raise Exception('description is required for Materials Commons')
        super(MCPayload, self).__init__(*args, **kwargs)

    @property
    def metapayload(self):
        return {
            'name': self['source']['name'],
            'description': self['description']
        }

class Human(dict):
    def __init__(self, given_name, family_name, email='', institution=''):
        super(Human, self).__init__()
        self.__dict__ = self
        self['given_name'] = given_name
        self['family_name'] = family_name
        self['email'] = email
        self['institution'] = institution
<|endoftext|>


<commit_before>import time
from threading import Event

from cassandra import ConsistencyLevel as CL
from cassandra import ReadFailure
from cassandra.query import SimpleStatement
from nose.tools import timed

from assertions import assert_invalid
from dtest import Tester, debug
from tools import known_failure, no_vnodes, since

class NotificationWaiter(object):
    """
    A helper class for waiting for pushed notifications from
    Cassandra over the native protocol.
    """

    def __init__(self, tester, node, notification_types, keyspace=None):
        """
        `address` should be a ccmlib.node.Node instance
        `notification_types` should be a list of
        "TOPOLOGY_CHANGE", "STATUS_CHANGE", and "SCHEMA_CHANGE".
        """
        self.node = node
        self.address = node.network_interfaces['binary'][0]
        self.notification_types = notification_types
        self.keyspace = keyspace

        # get a single, new connection
        session = tester.patient_exclusive_cql_connection(node)
        connection = session.cluster.connection_factory(self.address, is_control_connection=True)

        # coordinate with an Event
        self.event = Event()

        # the pushed notification
        self.notifications = []

        # register a callback for the notification type
        for notification_type in notification_types:
            connection.register_watcher(notification_type, self.handle_notification, register_timeout=5.0)

    def handle_notification(self, notification):
        """
        Called when a notification is pushed from Cassandra.
        """
        debug("Source {} sent {}".format(self.address, notification))

        if self.keyspace and notification['keyspace'] and self.keyspace != notification['keyspace']:
            return  # we are not interested in this schema change

        self.notifications.append(notification)
        self.event.set()

    def wait_for_notifications(self, timeout, num_notifications=1):
        """
        Waits up to `timeout` seconds for notifications from Cassandra. If
        passed `num_notifications`, stop waiting when that many notifications
        are observed.
        """

        deadline = time.time() + timeout
        while time.time() < deadline:
            self.event.wait(deadline - time.time())
            self.event.clear()
            if len(self.notifications) >= num_notifications:
                break

        return self.notifications

    def clear_notifications(self):
        self.notifications = []
        self.event.clear()

class TestPushedNotifications(Tester):
    """
    Tests for pushed native protocol notification from Cassandra.
    """

    @no_vnodes()
    def move_single_node_test(self):
        """
        @jira_ticket CASSANDRA-8516
        Moving a token should result in NODE_MOVED notifications.
        """
        self.cluster.populate(3).start(wait_for_binary_proto=True, wait_other_notice=True)

        # Despite waiting for each node to see the other nodes as UP, there is apparently
        # still a race condition that can result in NEW_NODE events being sent.  We don't
        # want to accidentally collect those, so for now we will just sleep a few seconds.
        time.sleep(3)

        waiters = [NotificationWaiter(self, node, ["TOPOLOGY_CHANGE"])
                   for node in self.cluster.nodes.values()]

        node1 = self.cluster.nodes.values()[0]
        node1.move("123")

        for waiter in waiters:
            debug("Waiting for notification from {}".format(waiter.address,))
            notifications = waiter.wait_for_notifications(60.0)
            self.assertEquals(1, len(notifications), notifications)
            notification = notifications[0]
            change_type = notification["change_type"]
            address, port = notification["address"]
            self.assertEquals("MOVED_NODE", change_type)
            self.assertEquals(self.get_ip_from_node(node1), address)

    @no_vnodes()
    def move_single_node_localhost_test(self):
        """
        @jira_ticket  CASSANDRA-10052
        Test that we don't get NODE_MOVED notifications from nodes other than the local one,
        when rpc_address is set to localhost (127.0.0.1).

        To set-up this test we override the rpc_address to "localhost (127.0.0.1)" for all nodes, and
        therefore we must change the rpc port or else processes won't start.
        """
        cluster = self.cluster
        cluster.populate(3)
        node1, node2, node3 = cluster.nodelist()

        # change node's 'rpc_address' from '127.0.0.x' to 'localhost (127.0.0.1)', increase port numbers
        i = 0
        for node in cluster.nodelist():
            # See CASSANDRA-11057 for IPv6 issues and why we set 127.0.0.1 for CCM
            debug('Set 127.0.0.1 to prevent IPv6 java prefs, set rpc_address: localhost in cassandra.yaml')
            node.network_interfaces['thrift'] = ('127.0.0.1', node.network_interfaces['thrift'][1] + i)
            node.network_interfaces['binary'] = ('127.0.0.1', node.network_interfaces['thrift'][1] + 1)
            node.import_config_files()  # this regenerates the yaml file and sets 'rpc_address' to the 'thrift' address
            node.set_configuration_options(values={'rpc_address': 'localhost'})
            debug(node.show())
            i = i + 2

        cluster.start(wait_for_binary_proto=True, wait_other_notice=True)

        # Despite waiting for each node to see the other nodes as UP, there is apparently
        # still a race condition that can result in NEW_NODE events being sent.  We don't
        # want to accidentally collect those, so for now we will just sleep a few seconds.
        time.sleep(3)

        waiters = [NotificationWaiter(self, node, ["TOPOLOGY_CHANGE"])
                   for node in self.cluster.nodes.values()]

        node1 = self.cluster.nodes.values()[0]
        node1.move("123")

        for waiter in waiters:
            debug("Waiting for notification from {}".format(waiter.address,))
            notifications = waiter.wait_for_notifications(30.0)
            self.assertEquals(1 if waiter.node is node1 else 0, len(notifications), notifications)

    @known_failure(failure_source='cassandra',
                   jira_url='https://issues.apache.org/jira/browse/CASSANDRA-11360',
                   flaky=True,
                   notes='2 different failures on 2.1 offheap memtables jobs. Fails on main 2.1 as well')
    def restart_node_test(self):
        """
        @jira_ticket CASSANDRA-7816
        Restarting a node should generate exactly one DOWN and one UP notification
        """

        self.cluster.populate(2).start(wait_for_binary_proto=True, wait_other_notice=True)
        node1, node2 = self.cluster.nodelist()
        # need to block so that these notifications don't confuse the state below
        waiter = NotificationWaiter(self, node1, ["STATUS_CHANGE", "TOPOLOGY_CHANGE"])
        waiter.wait_for_notifications(timeout=60, num_notifications=2)
        waiter.clear_notifications()

        for i in range(5):
            debug("Restarting second node...")
            node2.stop(wait_other_notice=True)
            node2.start(wait_other_notice=True)
            debug("Waiting for notifications from {}".format(waiter.address))
            notifications = waiter.wait_for_notifications(timeout=60.0, num_notifications=3)
            self.assertEquals(3, len(notifications), notifications)
            for notification in notifications:
                self.assertEquals(self.get_ip_from_node(node2), notification["address"][0])
            self.assertEquals("DOWN", notifications[0]["change_type"])
            self.assertEquals("UP", notifications[1]["change_type"])
            self.assertEquals("NEW_NODE", notifications[2]["change_type"])
            waiter.clear_notifications()

    def restart_node_localhost_test(self):
        """
        Test that we don't get client notifications when rpc_address is set to localhost.
        @jira_ticket  CASSANDRA-10052

        To set-up this test we override the rpc_address to "localhost" for all nodes, and
        therefore we must change the rpc port or else processes won't start.
        """
        cluster = self.cluster
        cluster.populate(2)
        node1, node2 = cluster.nodelist()

        i = 0  # change 'rpc_address' from '127.0.0.x' to 'localhost' and diversify port numbers
        for node in cluster.nodelist():
            node.network_interfaces['thrift'] = ('localhost', node.network_interfaces['thrift'][1] + i)
            node.network_interfaces['binary'] = ('localhost', node.network_interfaces['thrift'][1] + 1)
            node.import_config_files()  # this regenerates the yaml file and sets 'rpc_address' to the 'thrift' address
            debug(node.show())
            i = i + 2

        cluster.start(wait_for_binary_proto=True)

        # register for notification with node1
        waiter = NotificationWaiter(self, node1, ["STATUS_CHANGE", "TOPOLOGY_CHANGE"])

        # restart node 2
        debug("Restarting second node...")
        node2.stop(wait_other_notice=True)
        node2.start(wait_other_notice=True)

        # check that node1 did not send UP or DOWN notification for node2
        debug("Waiting for notifications from {}".format(waiter.address,))
        notifications = waiter.wait_for_notifications(timeout=30.0, num_notifications=3)
        self.assertEquals(0, len(notifications), notifications)

    @since("3.0")
    def schema_changes_test(self):
        """
        @jira_ticket CASSANDRA-10328
        Creating, updating and dropping a keyspace, a table and a materialized view
        will generate the correct schema change notifications.
        """

        self.cluster.populate(2).start(wait_for_binary_proto=True)
        node1, node2 = self.cluster.nodelist()

        session = self.patient_cql_connection(node1)
        waiter = NotificationWaiter(self, node2, ["SCHEMA_CHANGE"], keyspace='ks')

        self.create_ks(session, 'ks', 3)
        session.execute("create TABLE t (k int PRIMARY KEY , v int)")
        session.execute("alter TABLE t add v1 int;")

        session.execute("create MATERIALIZED VIEW mv as select * from t WHERE v IS NOT NULL AND v1 IS NOT NULL PRIMARY KEY (v, k)")
        session.execute(" alter materialized view mv with min_index_interval = 100")

        session.execute("drop MATERIALIZED VIEW mv")
        session.execute("drop TABLE t")
        session.execute("drop KEYSPACE ks")

        debug("Waiting for notifications from {}".format(waiter.address,))
        notifications = waiter.wait_for_notifications(timeout=60.0, num_notifications=8)
        self.assertEquals(8, len(notifications), notifications)
        self.assertDictContainsSubset({'change_type': u'CREATED', 'target_type': u'KEYSPACE'}, notifications[0])
        self.assertDictContainsSubset({'change_type': u'CREATED', 'target_type': u'TABLE', u'table': u't'}, notifications[1])
        self.assertDictContainsSubset({'change_type': u'UPDATED', 'target_type': u'TABLE', u'table': u't'}, notifications[2])
        self.assertDictContainsSubset({'change_type': u'CREATED', 'target_type': u'TABLE', u'table': u'mv'}, notifications[3])
        self.assertDictContainsSubset({'change_type': u'UPDATED', 'target_type': u'TABLE', u'table': u'mv'}, notifications[4])
        self.assertDictContainsSubset({'change_type': u'DROPPED', 'target_type': u'TABLE', u'table': u'mv'}, notifications[5])
        self.assertDictContainsSubset({'change_type': u'DROPPED', 'target_type': u'TABLE', u'table': u't'}, notifications[6])
        self.assertDictContainsSubset({'change_type': u'DROPPED', 'target_type': u'KEYSPACE'}, notifications[7])

class TestVariousNotifications(Tester):
    """
    Tests for various notifications/messages from Cassandra.
    """

    @known_failure(failure_source='test',
                   jira_url='https://issues.apache.org/jira/browse/CASSANDRA-11476')
    @since('2.2')
    def tombstone_failure_threshold_message_test(self):
        """
        Ensure nodes return an error message in case of TombstoneOverwhelmingExceptions rather
        than dropping the request. A drop makes the coordinator waits for the specified
        read_request_timeout_in_ms.
        @jira_ticket CASSANDRA-7886
        """

        self.allow_log_errors = True
        self.cluster.set_configuration_options(
            values={
                'tombstone_failure_threshold': 500,
                'read_request_timeout_in_ms': 30000,  # 30 seconds
                'range_request_timeout_in_ms': 40000
            }
        )
        self.cluster.populate(3).start()
        node1, node2, node3 = self.cluster.nodelist()
        session = self.patient_cql_connection(node1)

        self.create_ks(session, 'test', 3)
        session.execute(
            "CREATE TABLE test ( "
            "id int, mytext text, col1 int, col2 int, col3 int, "
            "PRIMARY KEY (id, mytext) )"
        )

        # Add data with tombstones
        values = map(lambda i: str(i), range(1000))
        for value in values:
            session.execute(SimpleStatement(
                "insert into test (id, mytext, col1) values (1, '{}', null) ".format(
                    value
                ),
                consistency_level=CL.ALL
            ))

        failure_msg = ("Scanned over.* tombstones.* query aborted")

        @timed(25)
        def read_failure_query():
            assert_invalid(
                session, SimpleStatement("select * from test where id in (1,2,3,4,5)", consistency_level=CL.ALL),
                expected=ReadFailure
            )

        read_failure_query()

        failure = (node1.grep_log(failure_msg) or
                   node2.grep_log(failure_msg) or
                   node3.grep_log(failure_msg))

        self.assertTrue(failure, ("Cannot find tombstone failure threshold error in log "
                                  "after failed query"))
        mark1 = node1.mark_log()
        mark2 = node2.mark_log()
        mark3 = node3.mark_log()

        @timed(35)
        def range_request_failure_query():
            assert_invalid(
                session, SimpleStatement("select * from test", consistency_level=CL.ALL),
                expected=ReadFailure
            )

        range_request_failure_query()

        failure = (node1.watch_log_for(failure_msg, from_mark=mark1, timeout=5) or
                   node2.watch_log_for(failure_msg, from_mark=mark2, timeout=5) or
                   node3.watch_log_for(failure_msg, from_mark=mark3, timeout=5))

        self.assertTrue(failure, ("Cannot find tombstone failure threshold error in log "
                                  "after range_request_timeout_query"))
<commit_msg>Fix 11476. Watch node logs instead of grepping.<commit_after>import time
from threading import Event

from cassandra import ConsistencyLevel as CL
from cassandra import ReadFailure
from cassandra.query import SimpleStatement
from ccmlib.node import TimeoutError
from nose.tools import timed

from assertions import assert_invalid
from dtest import Tester, debug
from tools import known_failure, no_vnodes, since

class NotificationWaiter(object):
    """
    A helper class for waiting for pushed notifications from
    Cassandra over the native protocol.
    """

    def __init__(self, tester, node, notification_types, keyspace=None):
        """
        `address` should be a ccmlib.node.Node instance
        `notification_types` should be a list of
        "TOPOLOGY_CHANGE", "STATUS_CHANGE", and "SCHEMA_CHANGE".
        """
        self.node = node
        self.address = node.network_interfaces['binary'][0]
        self.notification_types = notification_types
        self.keyspace = keyspace

        # get a single, new connection
        session = tester.patient_exclusive_cql_connection(node)
        connection = session.cluster.connection_factory(self.address, is_control_connection=True)

        # coordinate with an Event
        self.event = Event()

        # the pushed notification
        self.notifications = []

        # register a callback for the notification type
        for notification_type in notification_types:
            connection.register_watcher(notification_type, self.handle_notification, register_timeout=5.0)

    def handle_notification(self, notification):
        """
        Called when a notification is pushed from Cassandra.
        """
        debug("Source {} sent {}".format(self.address, notification))

        if self.keyspace and notification['keyspace'] and self.keyspace != notification['keyspace']:
            return  # we are not interested in this schema change

        self.notifications.append(notification)
        self.event.set()

    def wait_for_notifications(self, timeout, num_notifications=1):
        """
        Waits up to `timeout` seconds for notifications from Cassandra. If
        passed `num_notifications`, stop waiting when that many notifications
        are observed.
        """

        deadline = time.time() + timeout
        while time.time() < deadline:
            self.event.wait(deadline - time.time())
            self.event.clear()
            if len(self.notifications) >= num_notifications:
                break

        return self.notifications

    def clear_notifications(self):
        self.notifications = []
        self.event.clear()

class TestPushedNotifications(Tester):
    """
    Tests for pushed native protocol notification from Cassandra.
    """

    @no_vnodes()
    def move_single_node_test(self):
        """
        @jira_ticket CASSANDRA-8516
        Moving a token should result in NODE_MOVED notifications.
        """
        self.cluster.populate(3).start(wait_for_binary_proto=True, wait_other_notice=True)

        # Despite waiting for each node to see the other nodes as UP, there is apparently
        # still a race condition that can result in NEW_NODE events being sent.  We don't
        # want to accidentally collect those, so for now we will just sleep a few seconds.
        time.sleep(3)

        waiters = [NotificationWaiter(self, node, ["TOPOLOGY_CHANGE"])
                   for node in self.cluster.nodes.values()]

        node1 = self.cluster.nodes.values()[0]
        node1.move("123")

        for waiter in waiters:
            debug("Waiting for notification from {}".format(waiter.address,))
            notifications = waiter.wait_for_notifications(60.0)
            self.assertEquals(1, len(notifications), notifications)
            notification = notifications[0]
            change_type = notification["change_type"]
            address, port = notification["address"]
            self.assertEquals("MOVED_NODE", change_type)
            self.assertEquals(self.get_ip_from_node(node1), address)

    @no_vnodes()
    def move_single_node_localhost_test(self):
        """
        @jira_ticket  CASSANDRA-10052
        Test that we don't get NODE_MOVED notifications from nodes other than the local one,
        when rpc_address is set to localhost (127.0.0.1).

        To set-up this test we override the rpc_address to "localhost (127.0.0.1)" for all nodes, and
        therefore we must change the rpc port or else processes won't start.
        """
        cluster = self.cluster
        cluster.populate(3)
        node1, node2, node3 = cluster.nodelist()

        # change node's 'rpc_address' from '127.0.0.x' to 'localhost (127.0.0.1)', increase port numbers
        i = 0
        for node in cluster.nodelist():
            # See CASSANDRA-11057 for IPv6 issues and why we set 127.0.0.1 for CCM
            debug('Set 127.0.0.1 to prevent IPv6 java prefs, set rpc_address: localhost in cassandra.yaml')
            node.network_interfaces['thrift'] = ('127.0.0.1', node.network_interfaces['thrift'][1] + i)
            node.network_interfaces['binary'] = ('127.0.0.1', node.network_interfaces['thrift'][1] + 1)
            node.import_config_files()  # this regenerates the yaml file and sets 'rpc_address' to the 'thrift' address
            node.set_configuration_options(values={'rpc_address': 'localhost'})
            debug(node.show())
            i = i + 2

        cluster.start(wait_for_binary_proto=True, wait_other_notice=True)

        # Despite waiting for each node to see the other nodes as UP, there is apparently
        # still a race condition that can result in NEW_NODE events being sent.  We don't
        # want to accidentally collect those, so for now we will just sleep a few seconds.
        time.sleep(3)

        waiters = [NotificationWaiter(self, node, ["TOPOLOGY_CHANGE"])
                   for node in self.cluster.nodes.values()]

        node1 = self.cluster.nodes.values()[0]
        node1.move("123")

        for waiter in waiters:
            debug("Waiting for notification from {}".format(waiter.address,))
            notifications = waiter.wait_for_notifications(30.0)
            self.assertEquals(1 if waiter.node is node1 else 0, len(notifications), notifications)

    @known_failure(failure_source='cassandra',
                   jira_url='https://issues.apache.org/jira/browse/CASSANDRA-11360',
                   flaky=True,
                   notes='2 different failures on 2.1 offheap memtables jobs. Fails on main 2.1 as well')
    def restart_node_test(self):
        """
        @jira_ticket CASSANDRA-7816
        Restarting a node should generate exactly one DOWN and one UP notification
        """

        self.cluster.populate(2).start(wait_for_binary_proto=True, wait_other_notice=True)
        node1, node2 = self.cluster.nodelist()
        # need to block so that these notifications don't confuse the state below
        waiter = NotificationWaiter(self, node1, ["STATUS_CHANGE", "TOPOLOGY_CHANGE"])
        waiter.wait_for_notifications(timeout=60, num_notifications=2)
        waiter.clear_notifications()

        for i in range(5):
            debug("Restarting second node...")
            node2.stop(wait_other_notice=True)
            node2.start(wait_other_notice=True)
            debug("Waiting for notifications from {}".format(waiter.address))
            notifications = waiter.wait_for_notifications(timeout=60.0, num_notifications=3)
            self.assertEquals(3, len(notifications), notifications)
            for notification in notifications:
                self.assertEquals(self.get_ip_from_node(node2), notification["address"][0])
            self.assertEquals("DOWN", notifications[0]["change_type"])
            self.assertEquals("UP", notifications[1]["change_type"])
            self.assertEquals("NEW_NODE", notifications[2]["change_type"])
            waiter.clear_notifications()

    def restart_node_localhost_test(self):
        """
        Test that we don't get client notifications when rpc_address is set to localhost.
        @jira_ticket  CASSANDRA-10052

        To set-up this test we override the rpc_address to "localhost" for all nodes, and
        therefore we must change the rpc port or else processes won't start.
        """
        cluster = self.cluster
        cluster.populate(2)
        node1, node2 = cluster.nodelist()

        i = 0  # change 'rpc_address' from '127.0.0.x' to 'localhost' and diversify port numbers
        for node in cluster.nodelist():
            node.network_interfaces['thrift'] = ('localhost', node.network_interfaces['thrift'][1] + i)
            node.network_interfaces['binary'] = ('localhost', node.network_interfaces['thrift'][1] + 1)
            node.import_config_files()  # this regenerates the yaml file and sets 'rpc_address' to the 'thrift' address
            debug(node.show())
            i = i + 2

        cluster.start(wait_for_binary_proto=True)

        # register for notification with node1
        waiter = NotificationWaiter(self, node1, ["STATUS_CHANGE", "TOPOLOGY_CHANGE"])

        # restart node 2
        debug("Restarting second node...")
        node2.stop(wait_other_notice=True)
        node2.start(wait_other_notice=True)

        # check that node1 did not send UP or DOWN notification for node2
        debug("Waiting for notifications from {}".format(waiter.address,))
        notifications = waiter.wait_for_notifications(timeout=30.0, num_notifications=3)
        self.assertEquals(0, len(notifications), notifications)

    @since("3.0")
    def schema_changes_test(self):
        """
        @jira_ticket CASSANDRA-10328
        Creating, updating and dropping a keyspace, a table and a materialized view
        will generate the correct schema change notifications.
        """

        self.cluster.populate(2).start(wait_for_binary_proto=True)
        node1, node2 = self.cluster.nodelist()

        session = self.patient_cql_connection(node1)
        waiter = NotificationWaiter(self, node2, ["SCHEMA_CHANGE"], keyspace='ks')

        self.create_ks(session, 'ks', 3)
        session.execute("create TABLE t (k int PRIMARY KEY , v int)")
        session.execute("alter TABLE t add v1 int;")

        session.execute("create MATERIALIZED VIEW mv as select * from t WHERE v IS NOT NULL AND v1 IS NOT NULL PRIMARY KEY (v, k)")
        session.execute(" alter materialized view mv with min_index_interval = 100")

        session.execute("drop MATERIALIZED VIEW mv")
        session.execute("drop TABLE t")
        session.execute("drop KEYSPACE ks")

        debug("Waiting for notifications from {}".format(waiter.address,))
        notifications = waiter.wait_for_notifications(timeout=60.0, num_notifications=8)
        self.assertEquals(8, len(notifications), notifications)
        self.assertDictContainsSubset({'change_type': u'CREATED', 'target_type': u'KEYSPACE'}, notifications[0])
        self.assertDictContainsSubset({'change_type': u'CREATED', 'target_type': u'TABLE', u'table': u't'}, notifications[1])
        self.assertDictContainsSubset({'change_type': u'UPDATED', 'target_type': u'TABLE', u'table': u't'}, notifications[2])
        self.assertDictContainsSubset({'change_type': u'CREATED', 'target_type': u'TABLE', u'table': u'mv'}, notifications[3])
        self.assertDictContainsSubset({'change_type': u'UPDATED', 'target_type': u'TABLE', u'table': u'mv'}, notifications[4])
        self.assertDictContainsSubset({'change_type': u'DROPPED', 'target_type': u'TABLE', u'table': u'mv'}, notifications[5])
        self.assertDictContainsSubset({'change_type': u'DROPPED', 'target_type': u'TABLE', u'table': u't'}, notifications[6])
        self.assertDictContainsSubset({'change_type': u'DROPPED', 'target_type': u'KEYSPACE'}, notifications[7])

class TestVariousNotifications(Tester):
    """
    Tests for various notifications/messages from Cassandra.
    """

    @since('2.2')
    def tombstone_failure_threshold_message_test(self):
        """
        Ensure nodes return an error message in case of TombstoneOverwhelmingExceptions rather
        than dropping the request. A drop makes the coordinator waits for the specified
        read_request_timeout_in_ms.
        @jira_ticket CASSANDRA-7886
        """

        self.allow_log_errors = True
        self.cluster.set_configuration_options(
            values={
                'tombstone_failure_threshold': 500,
                'read_request_timeout_in_ms': 30000,  # 30 seconds
                'range_request_timeout_in_ms': 40000
            }
        )
        self.cluster.populate(3).start()
        node1, node2, node3 = self.cluster.nodelist()
        session = self.patient_cql_connection(node1)

        self.create_ks(session, 'test', 3)
        session.execute(
            "CREATE TABLE test ( "
            "id int, mytext text, col1 int, col2 int, col3 int, "
            "PRIMARY KEY (id, mytext) )"
        )

        # Add data with tombstones
        values = map(lambda i: str(i), range(1000))
        for value in values:
            session.execute(SimpleStatement(
                "insert into test (id, mytext, col1) values (1, '{}', null) ".format(
                    value
                ),
                consistency_level=CL.ALL
            ))

        failure_msg = ("Scanned over.* tombstones.* query aborted")

        @timed(25)
        def read_failure_query():
            assert_invalid(
                session, SimpleStatement("select * from test where id in (1,2,3,4,5)", consistency_level=CL.ALL),
                expected=ReadFailure
            )

        read_failure_query()

        # In almost all cases, we should find the failure message on node1 within a few seconds.
        # If it is not on node1, we grep all logs, as it *absolutely* should be somewhere.
        # If we still cannot find it then, we fail the test, as this is a problem.
        try:
            node1.watch_log_for(failure_msg, timeout=5)
        except TimeoutError:
            failure = (node1.grep_log(failure_msg) or
                       node2.grep_log(failure_msg) or
                       node3.grep_log(failure_msg))

            self.assertTrue(failure, ("Cannot find tombstone failure threshold error in log "
                                      "after failed query"))

        mark1 = node1.mark_log()
        mark2 = node2.mark_log()
        mark3 = node3.mark_log()

        @timed(35)
        def range_request_failure_query():
            assert_invalid(
                session, SimpleStatement("select * from test", consistency_level=CL.ALL),
                expected=ReadFailure
            )

        range_request_failure_query()

        # In almost all cases, we should find the failure message on node1 within a few seconds.
        # If it is not on node1, we grep all logs, as it *absolutely* should be somewhere.
        # If we still cannot find it then, we fail the test, as this is a problem.
        try:
            node1.watch_log_for(failure_msg, from_mark=mark1, timeout=5)
        except TimeoutError:
            failure = (node1.grep_log(failure_msg, from_mark=mark1) or
                       node2.grep_log(failure_msg, from_mark=mark2) or
                       node3.grep_log(failure_msg, from_mark=mark3))

            self.assertTrue(failure, ("Cannot find tombstone failure threshold error in log "
                                      "after range_request_timeout_query"))
<|endoftext|>


<commit_before>from __future__ import print_function

import sys

import numpy as np

from astropy.wcs import WCSSUB_CELESTIAL
from astropy.wcs.utils import wcs_to_celestial_frame
from astropy.coordinates import BaseCoordinateFrame

from ..utils.wcs_utils import get_spatial_scale

class Polygon(object):
    def __init__(self, x, y):
        self.x = x
        self.y = y

def segment_angles(x, y):

    dx = np.diff(x)
    dy = np.diff(y)

    d = np.hypot(dx, dy)

    cos_theta = (-dx[:-1] * dx[1:] - dy[:-1] * dy[1:]) / (d[:-1] * d[1:])
    cos_theta = np.clip(cos_theta, -1., 1.)

    sin_theta = (-dx[:-1] * dy[1:] + dy[:-1] * dx[1:]) / (d[:-1] * d[1:])
    sin_theta = np.clip(sin_theta, -1., 1.)

    theta = np.arctan2(sin_theta, cos_theta)

    theta[0] = np.pi
    theta[-1] = np.pi

    return theta

def get_endpoints(x, y, width):

    # Pad with same values at ends, to find slope of perpendicular end
    # lines.
    try:
        xp = np.pad(x, 1, mode='edge')
        yp = np.pad(y, 1, mode='edge')
    except AttributeError:  # Numpy < 1.7
        xp = np.hstack([x[0], x, x[-1]])
        yp = np.hstack([y[0], y, y[-1]])

    dx = np.diff(xp)
    dy = np.diff(yp)

    alpha = segment_angles(xp, yp) / 2.
    beta = np.arctan2(dy, dx)[:-1]
    beta[0] = beta[1]
    gamma = -(np.pi - alpha - beta)

    dx = np.cos(gamma)
    dy = np.sin(gamma)

    angles = segment_angles(xp, yp) / 2.

    # Find points offset from main curve, on bisecting lines
    x1 = x - dx * width * 0.5 / np.sin(angles)
    x2 = x + dx * width * 0.5 / np.sin(angles)
    y1 = y - dy * width * 0.5 / np.sin(angles)
    y2 = y + dy * width * 0.5 / np.sin(angles)

    return x1, y1, x2, y2

class Path(object):
    """
    A curved path that may have a non-zero width and is used to extract
    slices from cubes.

    Parameters
    ----------
    xy_or_coords : list or Astropy coordinates
        The points defining the path. This can be passed as a list of (x, y)
        tuples, which is interpreted as being pixel positions, or it can be
        an Astropy coordinate object containing an array of 2 or more
        coordinates.
    width : None or float or :class:`~astropy.units.Quantity`
        The width of the path. If ``coords`` is passed as a list of pixel
        positions, the width should be given (if passed) as a floating-point
        value in pixels. If ``coords`` is a coordinate object, the width
        should be passed as a :class:`~astropy.units.Quantity` instance with
        units of angle. If None, interpolation is used at the position of the
        path.
    """

    def __init__(self, xy_or_coords, width=None):
        if isinstance(xy_or_coords, list):
            self._xy = xy_or_coords
            self._coords = None
        elif sys.version_info[0] > 2 and isinstance(xy_or_coords, zip):
            self._xy = list(xy_or_coords)
            self._coords = None
        else:
            self._xy = None
            self._coords = xy_or_coords
        self.width = width

    def add_point(self, xy_or_coord):
        """
        Add a point to the path

        Parameters
        ----------
        xy_or_coord : tuple or Astropy coordinate
            A tuple (x, y) containing the coordinates of the point to add (if
            the path is defined in pixel space), or an Astropy coordinate
            object (if it is defined in world coordinates).
        """
        if self._xy is not None:
            if isinstance(xy_or_coord, tuple):
                self._xy.append(xy_or_coord)
            else:
                raise TypeError("Path is defined as a list of pixel "
                                "coordinates, so `xy_or_coord` should be "
                                "a tuple of `(x,y)` pixel coordinates.")
        else:
            if isinstance(xy_or_coord, BaseCoordinateFrame):
                raise NotImplementedError("Cannot yet append world coordinates to path")
            else:
                raise TypeError("Path is defined in world coordinates, "
                                "so `xy_or_coord` should be an Astropy "
                                "coordinate object.")

    def get_xy(self, wcs=None):
        """
        Return the pixel coordinates of the path.

        If the path is defined in world coordinates, the appropriate WCS
        transformation should be passed.

        Parameters
        ----------
        wcs : :class:`~astropy.wcs.WCS`
            The WCS transformation to assume in order to transform the path
            to pixel coordinates.
        """
        if self._xy is not None:
            return self._xy
        else:
            if wcs is None:
                raise ValueError("`wcs` is needed in order to compute "
                                 "the pixel coordinates")
            else:

                # Extract the celestial component of the WCS
                wcs_sky = wcs.sub([WCSSUB_CELESTIAL])

                # Find the astropy name for the coordinates
                celestial_system = wcs_to_celestial_frame(wcs_sky)

                world_coords = self._coords.transform_to(celestial_system)

                xw, yw = world_coords.spherical.lon.degree, world_coords.spherical.lat.degree

                return list(zip(*wcs_sky.wcs_world2pix(xw, yw, 0)))

    def sample_points_edges(self, spacing, wcs=None):

        x, y = zip(*self.get_xy(wcs=wcs))

        # Find the distance interval between all pairs of points
        dx = np.diff(x)
        dy = np.diff(y)
        dd = np.hypot(dx, dy)

        # Find the total displacement along the broken curve
        d = np.hstack([0., np.cumsum(dd)])

        # Figure out the number of points to sample, and stop short of the
        # last point.
        n_points = np.floor(d[-1] / spacing)

        if n_points == 0:
            raise ValueError("Path is shorter than spacing")

        d_sampled = np.linspace(0., n_points * spacing, n_points + 1)

        x_sampled = np.interp(d_sampled, d, x)
        y_sampled = np.interp(d_sampled, d, y)

        return d_sampled, x_sampled, y_sampled

    def sample_points(self, spacing, wcs=None):

        d_sampled, x_sampled, y_sampled = self.sample_points_edges(spacing, wcs=wcs)

        x_sampled = 0.5 * (x_sampled[:-1] + x_sampled[1:])
        y_sampled = 0.5 * (y_sampled[:-1] + y_sampled[1:])

        return x_sampled, y_sampled

    def sample_polygons(self, spacing, wcs=None):

        x, y = zip(*self.get_xy(wcs=wcs))

        d_sampled, x_sampled, y_sampled = self.sample_points_edges(spacing, wcs=wcs)

        # Find the distance interval between all pairs of points
        dx = np.diff(x)
        dy = np.diff(y)
        dd = np.hypot(dx, dy)

        # Normalize to find unit vectors
        dx = dx / dd
        dy = dy / dd

        # Find the total displacement along the broken curve
        d = np.hstack([0., np.cumsum(dd)])

        interval = np.searchsorted(d, d_sampled) - 1
        interval[0] = 0

        dx = dx[interval]
        dy = dy[interval]

        polygons = []

        x_beg = x_sampled
        x_end = x_sampled + dx * spacing

        y_beg = y_sampled
        y_end = y_sampled + dy * spacing

        if hasattr(self.width, 'unit'):
            scale = get_spatial_scale(wcs)
            width = (self.width / scale).decompose()
        else:
            width = self.width

        x1 = x_beg - dy * width * 0.5
        y1 = y_beg + dx * width * 0.5

        x2 = x_end - dy * width * 0.5
        y2 = y_end + dx * width * 0.5

        x3 = x_end + dy * width * 0.5
        y3 = y_end - dx * width * 0.5

        x4 = x_beg + dy * width * 0.5
        y4 = y_beg - dx * width * 0.5

        for i in range(len(x_sampled) - 1):
            p = Polygon([x1[i], x2[i], x3[i], x4[i]], [y1[i], y2[i], y3[i], y4[i]])
            polygons.append(p)

        return polygons

    def to_patches(self, spacing, wcs=None, **kwargs):
        from matplotlib.patches import Polygon as MPLPolygon
        patches = []
        for poly in self.sample_polygons(spacing, wcs=wcs):
            patches.append(MPLPolygon(zip(poly.x, poly.y), **kwargs))
        return patches
<commit_msg>Fix compatibility with Python 3 and avoid passing Quantity to matplotlib<commit_after>from __future__ import print_function

import sys

import numpy as np

from astropy.wcs import WCSSUB_CELESTIAL
from astropy.wcs.utils import wcs_to_celestial_frame
from astropy.coordinates import BaseCoordinateFrame

from ..utils.wcs_utils import get_spatial_scale

class Polygon(object):
    def __init__(self, x, y):
        self.x = x
        self.y = y

def segment_angles(x, y):

    dx = np.diff(x)
    dy = np.diff(y)

    d = np.hypot(dx, dy)

    cos_theta = (-dx[:-1] * dx[1:] - dy[:-1] * dy[1:]) / (d[:-1] * d[1:])
    cos_theta = np.clip(cos_theta, -1., 1.)

    sin_theta = (-dx[:-1] * dy[1:] + dy[:-1] * dx[1:]) / (d[:-1] * d[1:])
    sin_theta = np.clip(sin_theta, -1., 1.)

    theta = np.arctan2(sin_theta, cos_theta)

    theta[0] = np.pi
    theta[-1] = np.pi

    return theta

def get_endpoints(x, y, width):

    # Pad with same values at ends, to find slope of perpendicular end
    # lines.
    try:
        xp = np.pad(x, 1, mode='edge')
        yp = np.pad(y, 1, mode='edge')
    except AttributeError:  # Numpy < 1.7
        xp = np.hstack([x[0], x, x[-1]])
        yp = np.hstack([y[0], y, y[-1]])

    dx = np.diff(xp)
    dy = np.diff(yp)

    alpha = segment_angles(xp, yp) / 2.
    beta = np.arctan2(dy, dx)[:-1]
    beta[0] = beta[1]
    gamma = -(np.pi - alpha - beta)

    dx = np.cos(gamma)
    dy = np.sin(gamma)

    angles = segment_angles(xp, yp) / 2.

    # Find points offset from main curve, on bisecting lines
    x1 = x - dx * width * 0.5 / np.sin(angles)
    x2 = x + dx * width * 0.5 / np.sin(angles)
    y1 = y - dy * width * 0.5 / np.sin(angles)
    y2 = y + dy * width * 0.5 / np.sin(angles)

    return x1, y1, x2, y2

class Path(object):
    """
    A curved path that may have a non-zero width and is used to extract
    slices from cubes.

    Parameters
    ----------
    xy_or_coords : list or Astropy coordinates
        The points defining the path. This can be passed as a list of (x, y)
        tuples, which is interpreted as being pixel positions, or it can be
        an Astropy coordinate object containing an array of 2 or more
        coordinates.
    width : None or float or :class:`~astropy.units.Quantity`
        The width of the path. If ``coords`` is passed as a list of pixel
        positions, the width should be given (if passed) as a floating-point
        value in pixels. If ``coords`` is a coordinate object, the width
        should be passed as a :class:`~astropy.units.Quantity` instance with
        units of angle. If None, interpolation is used at the position of the
        path.
    """

    def __init__(self, xy_or_coords, width=None):
        if isinstance(xy_or_coords, list):
            self._xy = xy_or_coords
            self._coords = None
        elif sys.version_info[0] > 2 and isinstance(xy_or_coords, zip):
            self._xy = list(xy_or_coords)
            self._coords = None
        else:
            self._xy = None
            self._coords = xy_or_coords
        self.width = width

    def add_point(self, xy_or_coord):
        """
        Add a point to the path

        Parameters
        ----------
        xy_or_coord : tuple or Astropy coordinate
            A tuple (x, y) containing the coordinates of the point to add (if
            the path is defined in pixel space), or an Astropy coordinate
            object (if it is defined in world coordinates).
        """
        if self._xy is not None:
            if isinstance(xy_or_coord, tuple):
                self._xy.append(xy_or_coord)
            else:
                raise TypeError("Path is defined as a list of pixel "
                                "coordinates, so `xy_or_coord` should be "
                                "a tuple of `(x,y)` pixel coordinates.")
        else:
            if isinstance(xy_or_coord, BaseCoordinateFrame):
                raise NotImplementedError("Cannot yet append world coordinates to path")
            else:
                raise TypeError("Path is defined in world coordinates, "
                                "so `xy_or_coord` should be an Astropy "
                                "coordinate object.")

    def get_xy(self, wcs=None):
        """
        Return the pixel coordinates of the path.

        If the path is defined in world coordinates, the appropriate WCS
        transformation should be passed.

        Parameters
        ----------
        wcs : :class:`~astropy.wcs.WCS`
            The WCS transformation to assume in order to transform the path
            to pixel coordinates.
        """
        if self._xy is not None:
            return self._xy
        else:
            if wcs is None:
                raise ValueError("`wcs` is needed in order to compute "
                                 "the pixel coordinates")
            else:

                # Extract the celestial component of the WCS
                wcs_sky = wcs.sub([WCSSUB_CELESTIAL])

                # Find the astropy name for the coordinates
                celestial_system = wcs_to_celestial_frame(wcs_sky)

                world_coords = self._coords.transform_to(celestial_system)

                xw, yw = world_coords.spherical.lon.degree, world_coords.spherical.lat.degree

                return list(zip(*wcs_sky.wcs_world2pix(xw, yw, 0)))

    def sample_points_edges(self, spacing, wcs=None):

        x, y = zip(*self.get_xy(wcs=wcs))

        # Find the distance interval between all pairs of points
        dx = np.diff(x)
        dy = np.diff(y)
        dd = np.hypot(dx, dy)

        # Find the total displacement along the broken curve
        d = np.hstack([0., np.cumsum(dd)])

        # Figure out the number of points to sample, and stop short of the
        # last point.
        n_points = np.floor(d[-1] / spacing)

        if n_points == 0:
            raise ValueError("Path is shorter than spacing")

        d_sampled = np.linspace(0., n_points * spacing, n_points + 1)

        x_sampled = np.interp(d_sampled, d, x)
        y_sampled = np.interp(d_sampled, d, y)

        return d_sampled, x_sampled, y_sampled

    def sample_points(self, spacing, wcs=None):

        d_sampled, x_sampled, y_sampled = self.sample_points_edges(spacing, wcs=wcs)

        x_sampled = 0.5 * (x_sampled[:-1] + x_sampled[1:])
        y_sampled = 0.5 * (y_sampled[:-1] + y_sampled[1:])

        return x_sampled, y_sampled

    def sample_polygons(self, spacing, wcs=None):

        x, y = zip(*self.get_xy(wcs=wcs))

        d_sampled, x_sampled, y_sampled = self.sample_points_edges(spacing, wcs=wcs)

        # Find the distance interval between all pairs of points
        dx = np.diff(x)
        dy = np.diff(y)
        dd = np.hypot(dx, dy)

        # Normalize to find unit vectors
        dx = dx / dd
        dy = dy / dd

        # Find the total displacement along the broken curve
        d = np.hstack([0., np.cumsum(dd)])

        interval = np.searchsorted(d, d_sampled) - 1
        interval[0] = 0

        dx = dx[interval]
        dy = dy[interval]

        polygons = []

        x_beg = x_sampled
        x_end = x_sampled + dx * spacing

        y_beg = y_sampled
        y_end = y_sampled + dy * spacing

        if hasattr(self.width, 'unit'):
            scale = get_spatial_scale(wcs)
            width = (self.width / scale).decompose().value
        else:
            width = self.width

        x1 = x_beg - dy * width * 0.5
        y1 = y_beg + dx * width * 0.5

        x2 = x_end - dy * width * 0.5
        y2 = y_end + dx * width * 0.5

        x3 = x_end + dy * width * 0.5
        y3 = y_end - dx * width * 0.5

        x4 = x_beg + dy * width * 0.5
        y4 = y_beg - dx * width * 0.5

        for i in range(len(x_sampled) - 1):
            p = Polygon([x1[i], x2[i], x3[i], x4[i]], [y1[i], y2[i], y3[i], y4[i]])
            polygons.append(p)

        return polygons

    def to_patches(self, spacing, wcs=None, **kwargs):
        from matplotlib.patches import Polygon as MPLPolygon
        patches = []
        for poly in self.sample_polygons(spacing, wcs=wcs):
            patches.append(MPLPolygon(list(zip(poly.x, poly.y)), **kwargs))
        return patches
<|endoftext|>


<commit_before>def implode(glue, pieces):
    return str(glue).join(map(str, pieces))

def explode(delimiter, string, limit = None):
    if limit is None:
        return string.split(delimiter)
    else:
        if limit < 0:
            return string.split(delimiter, limit)[:limit]
        else:
            return string.split(delimiter, limit)
            <commit_msg>added join function (alias of implode<commit_after>def explode(delimiter, string, limit = None):
    if limit is None:
        return string.split(delimiter)
    else:
        if limit < 0:
            return string.split(delimiter, limit)[:limit]
        else:
            return string.split(delimiter, limit)

def implode(glue, pieces):
    return str(glue).join(map(str, pieces))

def join(glue, pieces):
    return implode(glue, pieces)<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
desispec.scripts.night
======================

Entry points for start/update/end night scripts.
"""
from __future__ import (absolute_import, division, print_function,
    unicode_literals)

import sys
import os
import argparse
import subprocess as sp
import time

from desiutil.log import get_logger

from .. import io

from .. import pipeline as pipe

def stage_from_command():
    """
    Extract the processing stage from the executable command.

    Returns:
        str: The extracted stage.

    Raises:
        KeyError: If the extracted string does not match the set of stages.
    """
    stage = os.path.basename(sys.argv[0]).split("_")[1]
    if stage not in ("start", "update", "end"):
        raise KeyError("Command does not match one of the stages!")
    return stage

def parse(stage, options=None):
    """
    Parse command-line options for start/update/end night scripts.

    Args:
        stage (str): The stage of the launch, one of "start", "update", "end".
        options : iterable

    Returns:
        :class:`argparse.Namespace`: The parsed command-line options.
    """
    desc = {"start": "Begin DESI pipeline processing for a particular night.",
            "update": "Run DESI pipeline on new data for a particular night.",
            "end": "Conclude DESI pipeline processing for a particular night."}

    parser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]),
        description=desc[stage])

    parser.add_argument("--night", required=True, type=str,
        help="Night ID in the form YYYYMMDD.")

    parser.add_argument("--nersc", required=False, default=None,
        help="run jobs on this NERSC system (edison | cori-haswell "
        "| cori-knl)")

    parser.add_argument("--nersc_queue", required=False, default="regular",
        help="use this NERSC queue (debug | regular)")

    parser.add_argument("--nersc_runtime", required=False, type=int,
        default=30, help="Then maximum run time (in minutes) for a single "
        " NERSC job.  If the list of tasks cannot be run in this time,"
        " multiple job scripts will be written")

    parser.add_argument("--nersc_shifter", required=False, default=None,
        help="The shifter image to use for NERSC jobs")

    parser.add_argument("--mpi_procs", required=False, type=int, default=1,
        help="The number of MPI processes to use for non-NERSC shell "
        "scripts (default 1)")

    parser.add_argument("--mpi_run", required=False, type=str,
        default="", help="The command to launch MPI programs "
        "for non-NERSC shell scripts (default do not use MPI)")

    parser.add_argument("--procs_per_node", required=False, type=int,
        default=0, help="The number of processes to use per node.  If not "
        "specified it uses a default value for each machine.")

    parser.add_argument("--debug", required=False, default=False,
        action="store_true", help="debugging messages in job logs")

    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)

    status = 0
    log = get_logger()

    int_night = 0
    if status == 0:
        try:
            int_night = int(args.night)
        except ValueError:
            log.critical("Could not convert night \"{}\" to an integer!"\
                .format(args.night))
            status = 3

    if status == 0:
        year = int_night // 10000
        month = (int_night - year*10000) // 100
        day = int_night - year*10000 - month*100
        try:
            assert 1969 < year < 2038
            assert 0 < month < 13
            assert 0 < day < 32
        except AssertionError:
            log.critical("Value for night \"{}\" is not a valid calendar "
                "date!".format(args.night))
            status = 4

    if status == 0:
        try:
            raw = io.rawdata_root()
        except AssertionError:
            log.critical("Raw data location not set")
            status = 5

    if status == 0:
        try:
            proddir = io.specprod_root()
        except AssertionError:
            log.critical("Production location not set")
            status = 6

    return args, status

def main():
    """
    Entry point for :command:`desi_start_night`,
    :command:`desi_update_night` and :command:`desi_end_night`.

    Returns:
        :class:`int`: An integer suitable for passing to :func:`sys.exit`.
    """
    # Find stage
    try:
        stage = stage_from_command()
    except:
        return 2

    # Parse and check arguments
    args, status = parse(stage)
    if status != 0:
        return status

    # Perform a production update to get any new raw data

    pipe.update_prod(nightstr=args.night)

    # Call desi_pipe chain.
    # FIXME: the functions in desi_pipe should be split off into functions
    # that we can call here.

    states = "waiting,ready"

    chaincom = ["desi_pipe", "chain", "--night", "{}".format(args.night)]
    chaincom.extend(["--states", states])
    chaincom.extend(["--nosubmitted"])

    ttlist = None
    if stage == "start":
        ttlist = ["preproc", "psf", "psfnight", "traceshift",
            "extract", "fiberflat", "fiberflatnight"]
        ttstr = ",".join(ttlist)
        chaincom.extend(["--tasktypes", ttstr])

    elif stage == "update":
        ttlist = ["preproc", "psf", "psfnight", "traceshift",
            "extract", "fiberflat", "fiberflatnight", "sky", "starfit",
            "fluxcalib", "cframe"]
        ttstr = ",".join(ttlist)
        chaincom.extend(["--tasktypes", ttstr])

    elif stage == "end":
        ttlist = ["preproc", "psf", "psfnight", "traceshift",
            "extract", "fiberflat", "fiberflatnight", "sky", "starfit",
            "fluxcalib", "cframe", "spectra", "redshift"]
        ttstr = ",".join(ttlist)
        chaincom.extend(["--tasktypes", ttstr])

    else:
        return 2

    machine = None

    if args.nersc is not None:
        if args.nersc == "shell":
            # override to use plain bash scripts
            machine = "shell"
        else:
            machine = args.nersc
    else:
        if os.getenv("NERSC_HOST") == "edison":
            machine = "edison"
        elif os.getenv("NERSC_HOST") == "cori":
            machine = "cori"
        else:
            print("NERSC_HOST environment variable not set, and --nersc "
                "option not given")
            return 1

    if machine == "shell":
        if args.mpi_procs > 1:
            chaincom.extend(["--mpi_procs", "{}".format(args.mpi_procs)])
        if args.mpi_run != "":
            chaincom.extend(["--mpi_run", args.mpi_run])
    else:
        machprops = pipe.scriptgen.nersc_machine(machine, args.nersc_queue)
        if len(ttlist) > machprops["submitlimit"]:
            print("Queue {} on machine {} limited to {} jobs."\
                .format(args.nersc_queue, machine,
                machprops["submitlimit"]))
            print("Use a different queue or shorter chains of tasks.")
            return 1
        chaincom.extend(["--nersc", machine])
        chaincom.extend(["--nersc_queue", args.nersc_queue])
        chaincom.extend(["--nersc_runtime", "{}".format(args.nersc_runtime)])
        if args.nersc_shifter is not None:
            chaincom.extend(["--nersc_shifter", args.nersc_shifter])
        if args.procs_per_node > 0:
            chaincom.extend(["--procs_per_node",
                "{}".format(args.procs_per_node)])

    if args.debug:
        chaincom.extend(["--debug"])

    # Call desi_pipe chain

    log = get_logger()
    proddir = io.specprod_root()
    scrdir = io.get_pipe_scriptdir()
    ntdir = os.path.join(proddir, io.get_pipe_rundir(),
        io.get_pipe_scriptdir(), io.get_pipe_nightdir())

    if not os.path.isdir(ntdir):
        os.makedirs(ntdir)

    thisnight = os.path.join(ntdir, "{}".format(args.night))

    if not os.path.isdir(thisnight):
        os.makedirs(thisnight)

    reldir = os.path.join(io.get_pipe_nightdir(), "{}".format(args.night))

    chaincom.extend(["--outdir", reldir])

    log.info("Production {}:".format(proddir))
    log.info("  Running stage {} for night {}.".format(stage, args.night))
    log.info("  Chain command:  {}".format(" ".join(chaincom)))
    sys.stdout.flush()

    jobstr = sp.check_output(" ".join(chaincom), shell=True,
        universal_newlines=True)
    jobs = jobstr.split(",")

    if len(jobs) > 0:
        log.info("  Final job ID(s) are:  {}".format(":".join(jobs)))
        sys.stdout.flush()

    return 0
<commit_msg>fix broken desi_start_night options<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
desispec.scripts.night
======================

Entry points for start/update/end night scripts.
"""
from __future__ import (absolute_import, division, print_function,
    unicode_literals)

import sys
import os
import argparse
import subprocess as sp
import time

from desiutil.log import get_logger

from .. import io

from .. import pipeline as pipe

def stage_from_command():
    """
    Extract the processing stage from the executable command.

    Returns:
        str: The extracted stage.

    Raises:
        KeyError: If the extracted string does not match the set of stages.
    """
    stage = os.path.basename(sys.argv[0]).split("_")[1]
    if stage not in ("start", "update", "end"):
        raise KeyError("Command does not match one of the stages!")
    return stage

def parse(stage, options=None):
    """
    Parse command-line options for start/update/end night scripts.

    Args:
        stage (str): The stage of the launch, one of "start", "update", "end".
        options : iterable

    Returns:
        :class:`argparse.Namespace`: The parsed command-line options.
    """
    desc = {"start": "Begin DESI pipeline processing for a particular night.",
            "update": "Run DESI pipeline on new data for a particular night.",
            "end": "Conclude DESI pipeline processing for a particular night."}

    parser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]),
        description=desc[stage])

    parser.add_argument("--night", required=True, type=str,
        help="Night ID in the form YYYYMMDD.")

    parser.add_argument("--nersc", required=False, default=None,
        help="run jobs on this NERSC system (edison | cori-haswell "
        "| cori-knl)")

    parser.add_argument("--nersc_queue", required=False, default="regular",
        help="use this NERSC queue (debug | regular)")

    parser.add_argument("--nersc_maxtime", required=False, type=int,
        default=30, help="Then maximum run time (in minutes) for a single "
        " NERSC job.  If the list of tasks cannot be run in this time,"
        " multiple job scripts will be written")

    parser.add_argument("--nersc_shifter", required=False, default=None,
        help="The shifter image to use for NERSC jobs")

    parser.add_argument("--mpi_procs", required=False, type=int, default=1,
        help="The number of MPI processes to use for non-NERSC shell "
        "scripts (default 1)")

    parser.add_argument("--mpi_run", required=False, type=str,
        default="", help="The command to launch MPI programs "
        "for non-NERSC shell scripts (default do not use MPI)")

    parser.add_argument("--procs_per_node", required=False, type=int,
        default=0, help="The number of processes to use per node.  If not "
        "specified it uses a default value for each machine.")

    parser.add_argument("--debug", required=False, default=False,
        action="store_true", help="debugging messages in job logs")

    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)

    status = 0
    log = get_logger()

    int_night = 0
    if status == 0:
        try:
            int_night = int(args.night)
        except ValueError:
            log.critical("Could not convert night \"{}\" to an integer!"\
                .format(args.night))
            status = 3

    if status == 0:
        year = int_night // 10000
        month = (int_night - year*10000) // 100
        day = int_night - year*10000 - month*100
        try:
            assert 1969 < year < 2038
            assert 0 < month < 13
            assert 0 < day < 32
        except AssertionError:
            log.critical("Value for night \"{}\" is not a valid calendar "
                "date!".format(args.night))
            status = 4

    if status == 0:
        try:
            raw = io.rawdata_root()
        except AssertionError:
            log.critical("Raw data location not set")
            status = 5

    if status == 0:
        try:
            proddir = io.specprod_root()
        except AssertionError:
            log.critical("Production location not set")
            status = 6

    return args, status

def main():
    """
    Entry point for :command:`desi_start_night`,
    :command:`desi_update_night` and :command:`desi_end_night`.

    Returns:
        :class:`int`: An integer suitable for passing to :func:`sys.exit`.
    """
    # Find stage
    try:
        stage = stage_from_command()
    except:
        return 2

    # Parse and check arguments
    args, status = parse(stage)
    if status != 0:
        return status

    # Perform a production update to get any new raw data

    pipe.update_prod(nightstr=args.night)

    # Call desi_pipe chain.
    # FIXME: the functions in desi_pipe should be split off into functions
    # that we can call here.

    states = "waiting,ready"

    chaincom = ["desi_pipe", "chain", "--night", "{}".format(args.night)]
    chaincom.extend(["--states", states])
    chaincom.extend(["--nosubmitted"])

    ttlist = None
    if stage == "start":
        ttlist = ["preproc", "psf", "psfnight", "traceshift",
            "extract", "fiberflat", "fiberflatnight"]
        ttstr = ",".join(ttlist)
        chaincom.extend(["--tasktypes", ttstr])

    elif stage == "update":
        ttlist = ["preproc", "psf", "psfnight", "traceshift",
            "extract", "fiberflat", "fiberflatnight", "sky", "starfit",
            "fluxcalib", "cframe"]
        ttstr = ",".join(ttlist)
        chaincom.extend(["--tasktypes", ttstr])

    elif stage == "end":
        ttlist = ["preproc", "psf", "psfnight", "traceshift",
            "extract", "fiberflat", "fiberflatnight", "sky", "starfit",
            "fluxcalib", "cframe", "spectra", "redshift"]
        ttstr = ",".join(ttlist)
        chaincom.extend(["--tasktypes", ttstr])

    else:
        return 2

    machine = None

    if args.nersc is not None:
        if args.nersc == "shell":
            # override to use plain bash scripts
            machine = "shell"
        else:
            machine = args.nersc
    else:
        if os.getenv("NERSC_HOST") == "edison":
            machine = "edison"
        elif os.getenv("NERSC_HOST") == "cori":
            machine = "cori-haswell"
        else:
            print("NERSC_HOST environment variable not set, and --nersc "
                "option not given")
            return 1

    if machine == "shell":
        if args.mpi_procs > 1:
            chaincom.extend(["--mpi_procs", "{}".format(args.mpi_procs)])
        if args.mpi_run != "":
            chaincom.extend(["--mpi_run", args.mpi_run])
    else:
        machprops = pipe.scriptgen.nersc_machine(machine, args.nersc_queue)
        if len(ttlist) > machprops["submitlimit"]:
            print("Queue {} on machine {} limited to {} jobs."\
                .format(args.nersc_queue, machine,
                machprops["submitlimit"]))
            print("Use a different queue or shorter chains of tasks.")
            return 1
        chaincom.extend(["--nersc", machine])
        chaincom.extend(["--nersc_queue", args.nersc_queue])
        chaincom.extend(["--nersc_maxtime", "{}".format(args.nersc_maxtime)])
        if args.nersc_shifter is not None:
            chaincom.extend(["--nersc_shifter", args.nersc_shifter])
        if args.procs_per_node > 0:
            chaincom.extend(["--procs_per_node",
                "{}".format(args.procs_per_node)])

    if args.debug:
        chaincom.extend(["--debug"])

    # Call desi_pipe chain

    log = get_logger()
    proddir = io.specprod_root()
    scrdir = io.get_pipe_scriptdir()
    ntdir = os.path.join(proddir, io.get_pipe_rundir(),
        io.get_pipe_scriptdir(), io.get_pipe_nightdir())

    if not os.path.isdir(ntdir):
        os.makedirs(ntdir)

    thisnight = os.path.join(ntdir, "{}".format(args.night))

    if not os.path.isdir(thisnight):
        os.makedirs(thisnight)

    reldir = os.path.join(io.get_pipe_nightdir(), "{}".format(args.night))

    chaincom.extend(["--outdir", reldir])

    log.info("Production {}:".format(proddir))
    log.info("  Running stage {} for night {}.".format(stage, args.night))
    log.info("  Chain command:  {}".format(" ".join(chaincom)))
    sys.stdout.flush()

    jobstr = sp.check_output(" ".join(chaincom), shell=True,
        universal_newlines=True)
    jobs = jobstr.split(",")

    if len(jobs) > 0:
        log.info("  Final job ID(s) are:  {}".format(":".join(jobs)))
        sys.stdout.flush()

    return 0
<|endoftext|>


<commit_before>import os
import signal
import logging
import pika
import simplejson
import threading
import Queue

class PublishSpooler(threading.Thread):
    def __init__(self, amqpURL):
        super(PublishSpooler, self).__init__()
        self.daemon = True
        self._queue = Queue.Queue()
        self._amqpURL = amqpURL
        threading.Thread.start(self)

    def run(self):
        self._declaredExchanges = set()
        try:
            self._connect()
        except:
            logging.exception("Could not connect to RabbitMQ broker. Commiting suicide.")
            os.kill(os.getpid(), signal.SIGTERM)
        while True:
            try:
                finishedEvent, command, kwargs, returnValue = self._queue.get(block=True, timeout=10)
            except Queue.Empty:
                self._connection.process_data_events()
                continue
            try:
                returnValue.data = command(**kwargs)
                self._connection.process_data_events()
            except Exception as e:
                returnValue.exception = e
                logging.exception("Got an exception while handling command.")
            finally:
                finishedEvent.set()

    def _executeCommand(self, function, **kwargs):
        class ReturnValue(object):
            def __init__(self):
                self.data = None
                self.exception = None

        finishedEvent = threading.Event()
        returnValue = ReturnValue()
        self._queue.put((finishedEvent, function, kwargs, returnValue), block=True)
        finishedEvent.wait()
        if returnValue.exception is not None:
            raise returnValue.exception
        return returnValue.data

    def _publish(self, exchange, message):
        if exchange not in self._declaredExchanges:
            self._channel.exchange_declare(exchange=exchange, type='fanout')
            self._declaredExchanges.add(exchange)
        self._channel.basic_publish(exchange=exchange, routing_key='', body=simplejson.dumps(message))

    def _cleanupAllocationPublishResources(self, allocationID):
        exchange = self.allocationExchange(allocationID)
        if exchange not in self._declaredExchanges:
            logging.warn("Tried to delete an unfamiliar exchange %(exchange)s.", dict(exchange=exchange))
            return
        self._declaredExchanges.remove(exchange)
        self._channel.exchange_delete(exchange=exchange)

    def _connect(self):
        logging.info("Rackattack event publisher connects to rabbit MQ %(url)s", dict(url=self._amqpURL))
        parameters = pika.URLParameters(self._amqpURL)
        self._connection = pika.BlockingConnection(parameters)
        self._channel = self._connection.channel()

    @classmethod
    def allocationExchange(cls, id):
        return "allocation_status__%s" % id

class Publish(PublishSpooler):
    ALL_HOSTS_ALLOCATIONS_EXCHANGE_NAME = 'allHostsAllocations'

    def __init__(self, amqpURL):
        super(Publish, self).__init__(amqpURL)

    def allocationChangedState(self, allocationID):
        self._publishAllocationStatus('changedState', allocationID)

    def cleanupAllocationPublishResources(self, allocationID):
        self._executeCommand(self._cleanupAllocationPublishResources, allocationID=allocationID)

    def allocationProviderMessage(self, allocationID, message):
        self._publishAllocationStatus('providerMessage', allocationID, message)

    def allocationWithdraw(self, allocationID, message):
        self._publishAllocationStatus('withdrawn', allocationID, message)

    def allocationRequested(self, requirements, allocationInfo):
        self._publishMessage(self.ALL_HOSTS_ALLOCATIONS_EXCHANGE_NAME,
                             event='requested', requirements=requirements, allocationInfo=allocationInfo)

    def allocationCreated(self, allocationID, requirements, allocationInfo, allocated):
        allocatedIDs = {name: stateMachine.hostImplementation().id()
                        for name, stateMachine in allocated.iteritems()}
        self._publishMessage(self.ALL_HOSTS_ALLOCATIONS_EXCHANGE_NAME,
                             event='created', allocationID=allocationID, requirements=requirements,
                             allocationInfo=allocationInfo, allocated=allocatedIDs)

    def _publishMessage(self, exchange, **message):
        self._executeCommand(self._publish, exchange=exchange, message=message)

    def _publishAllocationStatus(self, event, allocationID, message=None):
        exchange = self.allocationExchange(allocationID)
        self._publishMessage(exchange, event=event, allocationID=allocationID, message=message)
<commit_msg>Don't wait for messages to be published<commit_after>import os
import signal
import logging
import pika
import simplejson
import threading
import Queue

class PublishSpooler(threading.Thread):
    def __init__(self, amqpURL):
        super(PublishSpooler, self).__init__()
        self.daemon = True
        self._queue = Queue.Queue()
        self._amqpURL = amqpURL
        threading.Thread.start(self)

    def run(self):
        self._declaredExchanges = set()
        try:
            self._connect()
        except:
            logging.exception("Could not connect to RabbitMQ broker. Commiting suicide.")
            os.kill(os.getpid(), signal.SIGTERM)
        while True:
            try:
                command, kwargs = self._queue.get(block=True, timeout=10)
            except Queue.Empty:
                self._connection.process_data_events()
                continue
            try:
                command(**kwargs)
                self._connection.process_data_events()
            except Exception as e:
                logging.exception("Got an exception while handling command.")

    def _executeCommand(self, function, **kwargs):
        self._queue.put((function, kwargs), block=True)

    def _publish(self, exchange, message):
        if exchange not in self._declaredExchanges:
            self._channel.exchange_declare(exchange=exchange, type='fanout')
            self._declaredExchanges.add(exchange)
        self._channel.basic_publish(exchange=exchange, routing_key='', body=simplejson.dumps(message))

    def _cleanupAllocationPublishResources(self, allocationID):
        exchange = self.allocationExchange(allocationID)
        if exchange not in self._declaredExchanges:
            logging.warn("Tried to delete an unfamiliar exchange %(exchange)s.", dict(exchange=exchange))
            return
        self._declaredExchanges.remove(exchange)
        self._channel.exchange_delete(exchange=exchange)

    def _connect(self):
        logging.info("Rackattack event publisher connects to rabbit MQ %(url)s", dict(url=self._amqpURL))
        parameters = pika.URLParameters(self._amqpURL)
        self._connection = pika.BlockingConnection(parameters)
        self._channel = self._connection.channel()

    @classmethod
    def allocationExchange(cls, id):
        return "allocation_status__%s" % id

class Publish(PublishSpooler):
    ALL_HOSTS_ALLOCATIONS_EXCHANGE_NAME = 'allHostsAllocations'

    def __init__(self, amqpURL):
        super(Publish, self).__init__(amqpURL)

    def allocationChangedState(self, allocationID):
        self._publishAllocationStatus('changedState', allocationID)

    def cleanupAllocationPublishResources(self, allocationID):
        self._executeCommand(self._cleanupAllocationPublishResources, allocationID=allocationID)

    def allocationProviderMessage(self, allocationID, message):
        self._publishAllocationStatus('providerMessage', allocationID, message)

    def allocationWithdraw(self, allocationID, message):
        self._publishAllocationStatus('withdrawn', allocationID, message)

    def allocationRequested(self, requirements, allocationInfo):
        self._publishMessage(self.ALL_HOSTS_ALLOCATIONS_EXCHANGE_NAME,
                             event='requested', requirements=requirements, allocationInfo=allocationInfo)

    def allocationCreated(self, allocationID, requirements, allocationInfo, allocated):
        allocatedIDs = {name: stateMachine.hostImplementation().id()
                        for name, stateMachine in allocated.iteritems()}
        self._publishMessage(self.ALL_HOSTS_ALLOCATIONS_EXCHANGE_NAME,
                             event='created', allocationID=allocationID, requirements=requirements,
                             allocationInfo=allocationInfo, allocated=allocatedIDs)

    def _publishMessage(self, exchange, **message):
        self._executeCommand(self._publish, exchange=exchange, message=message)

    def _publishAllocationStatus(self, event, allocationID, message=None):
        exchange = self.allocationExchange(allocationID)
        self._publishMessage(exchange, event=event, allocationID=allocationID, message=message)
<|endoftext|>