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

# Following code fixes 2 issues with kafka-python and
# The current release of eventlet (0.19.0) does not actually remove
# select.poll [1]. Because of kafka-python.selectors34 selects
# PollSelector instead of SelectSelector [2]. PollSelector relies on
# select.poll, which does not work when eventlet/greenlet is used. This
# bug in evenlet is fixed in the master branch [3], but there's no
# release of eventlet that includes this fix at this point.

import json
import threading

import kafka
from kafka.client_async import selectors
import kafka.errors
from oslo_log import log as logging
from oslo_utils import eventletutils
import tenacity

from oslo_messaging._drivers import base
from oslo_messaging._drivers import common as driver_common
from oslo_messaging._drivers import kafka_options
from oslo_messaging._drivers import pool as driver_pool
from oslo_messaging._i18n import _LE
from oslo_messaging._i18n import _LW
from oslo_serialization import jsonutils

import logging as l
l.basicConfig(level=l.INFO)
l.getLogger("kafka").setLevel(l.WARN)
l.getLogger("stevedore").setLevel(l.WARN)

if eventletutils.is_monkey_patched('select'):
    # monkeypatch the vendored SelectSelector._select like eventlet does
    # https://github.com/eventlet/eventlet/blob/master/eventlet/green/selectors.py#L32
    from eventlet.green import select
    selectors.SelectSelector._select = staticmethod(select.select)

    # Force to use the select selectors
    KAFKA_SELECTOR = selectors.SelectSelector
else:
    KAFKA_SELECTOR = selectors.DefaultSelector

LOG = logging.getLogger(__name__)

def unpack_message(msg):
    context = {}
    message = None
    msg = json.loads(msg)
    message = driver_common.deserialize_msg(msg)
    context = message['_context']
    del message['_context']
    return context, message

def pack_message(ctxt, msg):
    """Pack context into msg."""

    if isinstance(ctxt, dict):
        context_d = ctxt
    else:
        context_d = ctxt.to_dict()
    msg['_context'] = context_d

    msg = driver_common.serialize_msg(msg)

    return msg

def target_to_topic(target, priority=None):
    """Convert target into topic string

    :param target: Message destination target
    :type target: oslo_messaging.Target
    :param priority: Notification priority
    :type priority: string
    """
    if not priority:
        return target.topic
    return target.topic + '.' + priority

def retry_on_retriable_kafka_error(exc):
    return (isinstance(exc, kafka.errors.KafkaError) and exc.retriable)

def with_reconnect(retries=None):
    def decorator(func):
        @tenacity.retry(
            retry=tenacity.retry_if_exception(retry_on_retriable_kafka_error),
            wait=tenacity.wait_fixed(1),
            stop=tenacity.stop_after_attempt(retries),
            reraise=True
        )
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper
    return decorator

class Connection(object):

    def __init__(self, conf, url, purpose):

        self.client = None
        driver_conf = conf.oslo_messaging_kafka
        self.batch_size = driver_conf.producer_batch_size
        self.linger_ms = driver_conf.producer_batch_timeout * 1000
        self.conf = conf
        self.producer = None
        self.producer_lock = threading.Lock()
        self.consumer = None
        self.consumer_timeout = float(driver_conf.kafka_consumer_timeout)
        self.max_fetch_bytes = driver_conf.kafka_max_fetch_bytes
        self.group_id = driver_conf.consumer_group
        self.url = url
        self._parse_url()
        self._consume_loop_stopped = False

    def _parse_url(self):
        driver_conf = self.conf.oslo_messaging_kafka
        self.hostaddrs = []

        for host in self.url.hosts:
            if host.hostname:
                self.hostaddrs.append("%s:%s" % (
                    host.hostname,
                    host.port or driver_conf.kafka_default_port))

        if not self.hostaddrs:
            self.hostaddrs.append("%s:%s" % (driver_conf.kafka_default_host,
                                             driver_conf.kafka_default_port))

    def notify_send(self, topic, ctxt, msg, retry):
        """Send messages to Kafka broker.

        :param topic: String of the topic
        :param ctxt: context for the messages
        :param msg: messages for publishing
        :param retry: the number of retry
        """
        retry = retry if retry >= 0 else None
        message = pack_message(ctxt, msg)
        message = jsonutils.dumps(message)

        @with_reconnect(retries=retry)
        def wrapped_with_reconnect():
            self._ensure_producer()
            # NOTE(sileht): This returns a future, we can use get()
            # if we want to block like other driver
            self.producer.send(topic, message)

        try:
            wrapped_with_reconnect()
        except Exception:
            # NOTE(sileht): if something goes wrong close the producer
            # connection
            self._close_producer()
            raise

    @with_reconnect()
    def _poll_messages(self, timeout):
        messages = self.consumer.poll(timeout * 1000.0)
        messages = [record.value
                    for records in messages.values() if records
                    for record in records]
        if not messages:
            # NOTE(sileht): really ? you return payload but no messages...
            # simulate timeout to consume message again
            raise kafka.errors.ConsumerTimeout()
        return messages

    def consume(self, timeout=None):
        """Receive up to 'max_fetch_messages' messages.

        :param timeout: poll timeout in seconds
        """

        def _raise_timeout(exc):
            raise driver_common.Timeout(exc.message)

        timer = driver_common.DecayingTimer(duration=timeout)
        timer.start()

        poll_timeout = (self.consumer_timeout if timeout is None
                        else min(timeout, self.consumer_timeout))

        while True:
            if self._consume_loop_stopped:
                return
            try:
                return self._poll_messages(poll_timeout)
            except kafka.errors.ConsumerTimeout as exc:
                poll_timeout = timer.check_return(
                    _raise_timeout, exc, maximum=self.consumer_timeout)
            except Exception:
                LOG.exception(_LE("Failed to consume messages"))
                return

    def stop_consuming(self):
        self._consume_loop_stopped = True

    def reset(self):
        """Reset a connection so it can be used again."""
        pass

    def close(self):
        self._close_producer()
        if self.consumer:
            self.consumer.close()
            self.consumer = None

    def _close_producer(self):
        with self.producer_lock:
            if self.producer:
                self.producer.close()
                self.producer = None

    def _ensure_producer(self):
        if self.producer:
            return
        with self.producer_lock:
            if self.producer:
                return
            self.producer = kafka.KafkaProducer(
                bootstrap_servers=self.hostaddrs,
                linger_ms=self.linger_ms,
                batch_size=self.batch_size,
                selector=KAFKA_SELECTOR)

    @with_reconnect()
    def declare_topic_consumer(self, topics, group=None):
        # TODO(Support for manual/auto_commit functionality)
        # When auto_commit is False, consumer can manually notify
        # the completion of the subscription.
        # Currently we don't support for non auto commit option
        self.consumer = kafka.KafkaConsumer(
            *topics, group_id=(group or self.group_id),
            bootstrap_servers=self.hostaddrs,
            max_partition_fetch_bytes=self.max_fetch_bytes,
            selector=KAFKA_SELECTOR
        )

class OsloKafkaMessage(base.RpcIncomingMessage):

    def __init__(self, ctxt, message):
        super(OsloKafkaMessage, self).__init__(ctxt, message)

    def requeue(self):
        LOG.warning(_LW("requeue is not supported"))

    def reply(self, reply=None, failure=None):
        LOG.warning(_LW("reply is not supported"))

class KafkaListener(base.PollStyleListener):

    def __init__(self, conn):
        super(KafkaListener, self).__init__()
        self._stopped = threading.Event()
        self.conn = conn
        self.incoming_queue = []

    @base.batch_poll_helper
    def poll(self, timeout=None):
        while not self._stopped.is_set():
            if self.incoming_queue:
                return self.incoming_queue.pop(0)
            try:
                messages = self.conn.consume(timeout=timeout) or []
                for message in messages:
                    msg = OsloKafkaMessage(*unpack_message(message))
                    self.incoming_queue.append(msg)
            except driver_common.Timeout:
                return None

    def stop(self):
        self._stopped.set()
        self.conn.stop_consuming()

    def cleanup(self):
        self.conn.close()

class KafkaDriver(base.BaseDriver):
    """Note: Current implementation of this driver is experimental.
    We will have functional and/or integrated testing enabled for this driver.
    """

    def __init__(self, conf, url, default_exchange=None,
                 allowed_remote_exmods=None):
        super(KafkaDriver, self).__init__(
            conf, url, default_exchange, allowed_remote_exmods)

        kafka_options.register_opts(conf)
        # the pool configuration properties
        max_size = self.conf.oslo_messaging_kafka.pool_size
        min_size = self.conf.oslo_messaging_kafka.conn_pool_min_size
        ttl = self.conf.oslo_messaging_kafka.conn_pool_ttl

        self.connection_pool = driver_pool.ConnectionPool(
            self.conf, max_size, min_size, ttl,
            self._url, Connection)
        self.listeners = []

    def cleanup(self):
        for c in self.listeners:
            c.close()
        self.listeners = []

    def send(self, target, ctxt, message, wait_for_reply=None, timeout=None,
             retry=None):
        raise NotImplementedError(
            'The RPC implementation for Kafka is not implemented')

    def send_notification(self, target, ctxt, message, version, retry=None):
        """Send notification to Kafka brokers

        :param target: Message destination target
        :type target: oslo_messaging.Target
        :param ctxt: Message context
        :type ctxt: dict
        :param message: Message payload to pass
        :type message: dict
        :param version: Messaging API version (currently not used)
        :type version: str
        :param retry: an optional default kafka consumer retries configuration
                      None means to retry forever
                      0 means no retry
                      N means N retries
        :type retry: int
        """
        with self._get_connection(purpose=driver_common.PURPOSE_SEND) as conn:
            conn.notify_send(target_to_topic(target), ctxt, message, retry)

    def listen(self, target, batch_size, batch_timeout):
        raise NotImplementedError(
            'The RPC implementation for Kafka is not implemented')

    def listen_for_notifications(self, targets_and_priorities, pool,
                                 batch_size, batch_timeout):
        """Listen to a specified list of targets on Kafka brokers

        :param targets_and_priorities: List of pairs (target, priority)
                                       priority is not used for kafka driver
                                       target.exchange_target.topic is used as
                                       a kafka topic
        :type targets_and_priorities: list
        :param pool: consumer group of Kafka consumers
        :type pool: string
        """
        conn = self._get_connection(purpose=driver_common.PURPOSE_LISTEN)
        topics = set()
        for target, priority in targets_and_priorities:
            topics.add(target_to_topic(target, priority))

        conn.declare_topic_consumer(topics, pool)

        listener = KafkaListener(conn)
        return base.PollStyleListenerAdapter(listener, batch_size,
                                             batch_timeout)

    def _get_connection(self, purpose):
        return driver_common.ConnectionContext(self.connection_pool, purpose)
<commit_msg>kafka: fix python3 exception<commit_after># Copyright (C) 2015 Cisco Systems, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

# Following code fixes 2 issues with kafka-python and
# The current release of eventlet (0.19.0) does not actually remove
# select.poll [1]. Because of kafka-python.selectors34 selects
# PollSelector instead of SelectSelector [2]. PollSelector relies on
# select.poll, which does not work when eventlet/greenlet is used. This
# bug in evenlet is fixed in the master branch [3], but there's no
# release of eventlet that includes this fix at this point.

import json
import threading

import kafka
from kafka.client_async import selectors
import kafka.errors
from oslo_log import log as logging
from oslo_utils import eventletutils
import tenacity

from oslo_messaging._drivers import base
from oslo_messaging._drivers import common as driver_common
from oslo_messaging._drivers import kafka_options
from oslo_messaging._drivers import pool as driver_pool
from oslo_messaging._i18n import _LE
from oslo_messaging._i18n import _LW
from oslo_serialization import jsonutils

import logging as l
l.basicConfig(level=l.INFO)
l.getLogger("kafka").setLevel(l.WARN)
l.getLogger("stevedore").setLevel(l.WARN)

if eventletutils.is_monkey_patched('select'):
    # monkeypatch the vendored SelectSelector._select like eventlet does
    # https://github.com/eventlet/eventlet/blob/master/eventlet/green/selectors.py#L32
    from eventlet.green import select
    selectors.SelectSelector._select = staticmethod(select.select)

    # Force to use the select selectors
    KAFKA_SELECTOR = selectors.SelectSelector
else:
    KAFKA_SELECTOR = selectors.DefaultSelector

LOG = logging.getLogger(__name__)

def unpack_message(msg):
    context = {}
    message = None
    msg = json.loads(msg)
    message = driver_common.deserialize_msg(msg)
    context = message['_context']
    del message['_context']
    return context, message

def pack_message(ctxt, msg):
    """Pack context into msg."""

    if isinstance(ctxt, dict):
        context_d = ctxt
    else:
        context_d = ctxt.to_dict()
    msg['_context'] = context_d

    msg = driver_common.serialize_msg(msg)

    return msg

def target_to_topic(target, priority=None):
    """Convert target into topic string

    :param target: Message destination target
    :type target: oslo_messaging.Target
    :param priority: Notification priority
    :type priority: string
    """
    if not priority:
        return target.topic
    return target.topic + '.' + priority

def retry_on_retriable_kafka_error(exc):
    return (isinstance(exc, kafka.errors.KafkaError) and exc.retriable)

def with_reconnect(retries=None):
    def decorator(func):
        @tenacity.retry(
            retry=tenacity.retry_if_exception(retry_on_retriable_kafka_error),
            wait=tenacity.wait_fixed(1),
            stop=tenacity.stop_after_attempt(retries),
            reraise=True
        )
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper
    return decorator

class Connection(object):

    def __init__(self, conf, url, purpose):

        self.client = None
        driver_conf = conf.oslo_messaging_kafka
        self.batch_size = driver_conf.producer_batch_size
        self.linger_ms = driver_conf.producer_batch_timeout * 1000
        self.conf = conf
        self.producer = None
        self.producer_lock = threading.Lock()
        self.consumer = None
        self.consumer_timeout = float(driver_conf.kafka_consumer_timeout)
        self.max_fetch_bytes = driver_conf.kafka_max_fetch_bytes
        self.group_id = driver_conf.consumer_group
        self.url = url
        self._parse_url()
        self._consume_loop_stopped = False

    def _parse_url(self):
        driver_conf = self.conf.oslo_messaging_kafka
        self.hostaddrs = []

        for host in self.url.hosts:
            if host.hostname:
                self.hostaddrs.append("%s:%s" % (
                    host.hostname,
                    host.port or driver_conf.kafka_default_port))

        if not self.hostaddrs:
            self.hostaddrs.append("%s:%s" % (driver_conf.kafka_default_host,
                                             driver_conf.kafka_default_port))

    def notify_send(self, topic, ctxt, msg, retry):
        """Send messages to Kafka broker.

        :param topic: String of the topic
        :param ctxt: context for the messages
        :param msg: messages for publishing
        :param retry: the number of retry
        """
        retry = retry if retry >= 0 else None
        message = pack_message(ctxt, msg)
        message = jsonutils.dumps(message)

        @with_reconnect(retries=retry)
        def wrapped_with_reconnect():
            self._ensure_producer()
            # NOTE(sileht): This returns a future, we can use get()
            # if we want to block like other driver
            self.producer.send(topic, message)

        try:
            wrapped_with_reconnect()
        except Exception:
            # NOTE(sileht): if something goes wrong close the producer
            # connection
            self._close_producer()
            raise

    @with_reconnect()
    def _poll_messages(self, timeout):
        messages = self.consumer.poll(timeout * 1000.0)
        messages = [record.value
                    for records in messages.values() if records
                    for record in records]
        if not messages:
            # NOTE(sileht): really ? you return payload but no messages...
            # simulate timeout to consume message again
            raise kafka.errors.ConsumerTimeout()
        return messages

    def consume(self, timeout=None):
        """Receive up to 'max_fetch_messages' messages.

        :param timeout: poll timeout in seconds
        """

        def _raise_timeout(exc):
            raise driver_common.Timeout(str(exc))

        timer = driver_common.DecayingTimer(duration=timeout)
        timer.start()

        poll_timeout = (self.consumer_timeout if timeout is None
                        else min(timeout, self.consumer_timeout))

        while True:
            if self._consume_loop_stopped:
                return
            try:
                return self._poll_messages(poll_timeout)
            except kafka.errors.ConsumerTimeout as exc:
                poll_timeout = timer.check_return(
                    _raise_timeout, exc, maximum=self.consumer_timeout)
            except Exception:
                LOG.exception(_LE("Failed to consume messages"))
                return

    def stop_consuming(self):
        self._consume_loop_stopped = True

    def reset(self):
        """Reset a connection so it can be used again."""
        pass

    def close(self):
        self._close_producer()
        if self.consumer:
            self.consumer.close()
            self.consumer = None

    def _close_producer(self):
        with self.producer_lock:
            if self.producer:
                self.producer.close()
                self.producer = None

    def _ensure_producer(self):
        if self.producer:
            return
        with self.producer_lock:
            if self.producer:
                return
            self.producer = kafka.KafkaProducer(
                bootstrap_servers=self.hostaddrs,
                linger_ms=self.linger_ms,
                batch_size=self.batch_size,
                selector=KAFKA_SELECTOR)

    @with_reconnect()
    def declare_topic_consumer(self, topics, group=None):
        # TODO(Support for manual/auto_commit functionality)
        # When auto_commit is False, consumer can manually notify
        # the completion of the subscription.
        # Currently we don't support for non auto commit option
        self.consumer = kafka.KafkaConsumer(
            *topics, group_id=(group or self.group_id),
            bootstrap_servers=self.hostaddrs,
            max_partition_fetch_bytes=self.max_fetch_bytes,
            selector=KAFKA_SELECTOR
        )

class OsloKafkaMessage(base.RpcIncomingMessage):

    def __init__(self, ctxt, message):
        super(OsloKafkaMessage, self).__init__(ctxt, message)

    def requeue(self):
        LOG.warning(_LW("requeue is not supported"))

    def reply(self, reply=None, failure=None):
        LOG.warning(_LW("reply is not supported"))

class KafkaListener(base.PollStyleListener):

    def __init__(self, conn):
        super(KafkaListener, self).__init__()
        self._stopped = threading.Event()
        self.conn = conn
        self.incoming_queue = []

    @base.batch_poll_helper
    def poll(self, timeout=None):
        while not self._stopped.is_set():
            if self.incoming_queue:
                return self.incoming_queue.pop(0)
            try:
                messages = self.conn.consume(timeout=timeout) or []
                for message in messages:
                    msg = OsloKafkaMessage(*unpack_message(message))
                    self.incoming_queue.append(msg)
            except driver_common.Timeout:
                return None

    def stop(self):
        self._stopped.set()
        self.conn.stop_consuming()

    def cleanup(self):
        self.conn.close()

class KafkaDriver(base.BaseDriver):
    """Note: Current implementation of this driver is experimental.
    We will have functional and/or integrated testing enabled for this driver.
    """

    def __init__(self, conf, url, default_exchange=None,
                 allowed_remote_exmods=None):
        super(KafkaDriver, self).__init__(
            conf, url, default_exchange, allowed_remote_exmods)

        kafka_options.register_opts(conf)
        # the pool configuration properties
        max_size = self.conf.oslo_messaging_kafka.pool_size
        min_size = self.conf.oslo_messaging_kafka.conn_pool_min_size
        ttl = self.conf.oslo_messaging_kafka.conn_pool_ttl

        self.connection_pool = driver_pool.ConnectionPool(
            self.conf, max_size, min_size, ttl,
            self._url, Connection)
        self.listeners = []

    def cleanup(self):
        for c in self.listeners:
            c.close()
        self.listeners = []

    def send(self, target, ctxt, message, wait_for_reply=None, timeout=None,
             retry=None):
        raise NotImplementedError(
            'The RPC implementation for Kafka is not implemented')

    def send_notification(self, target, ctxt, message, version, retry=None):
        """Send notification to Kafka brokers

        :param target: Message destination target
        :type target: oslo_messaging.Target
        :param ctxt: Message context
        :type ctxt: dict
        :param message: Message payload to pass
        :type message: dict
        :param version: Messaging API version (currently not used)
        :type version: str
        :param retry: an optional default kafka consumer retries configuration
                      None means to retry forever
                      0 means no retry
                      N means N retries
        :type retry: int
        """
        with self._get_connection(purpose=driver_common.PURPOSE_SEND) as conn:
            conn.notify_send(target_to_topic(target), ctxt, message, retry)

    def listen(self, target, batch_size, batch_timeout):
        raise NotImplementedError(
            'The RPC implementation for Kafka is not implemented')

    def listen_for_notifications(self, targets_and_priorities, pool,
                                 batch_size, batch_timeout):
        """Listen to a specified list of targets on Kafka brokers

        :param targets_and_priorities: List of pairs (target, priority)
                                       priority is not used for kafka driver
                                       target.exchange_target.topic is used as
                                       a kafka topic
        :type targets_and_priorities: list
        :param pool: consumer group of Kafka consumers
        :type pool: string
        """
        conn = self._get_connection(purpose=driver_common.PURPOSE_LISTEN)
        topics = set()
        for target, priority in targets_and_priorities:
            topics.add(target_to_topic(target, priority))

        conn.declare_topic_consumer(topics, pool)

        listener = KafkaListener(conn)
        return base.PollStyleListenerAdapter(listener, batch_size,
                                             batch_timeout)

    def _get_connection(self, purpose):
        return driver_common.ConnectionContext(self.connection_pool, purpose)
<|endoftext|>


<commit_before>import logging

from django.db import models
from django.utils.datetime_safe import datetime
from django.utils.translation import ugettext_lazy as _
from rest_framework.reverse import reverse

from osmaxx.conversion_api import statuses
from osmaxx.conversion_api.formats import FORMAT_CHOICES
from osmaxx.conversion_api.statuses import FAILED, FINAL_STATUSES, FINISHED

logger = logging.getLogger(__name__)

class TimeStampModelMixin(models.Model):
    created = models.DateTimeField(_('created on'), default=datetime.now, blank=True, editable=False)
    updated = models.DateTimeField(_('updated on'), default=None, blank=True, editable=False, null=True)

    def save(self, *args, **kwargs):
        now = datetime.now()
        if not self.id:
            self.created = now
        self.updated = now
        super().save(*args, **kwargs)

    class Meta:
        abstract = True

class Export(TimeStampModelMixin, models.Model):
    """
    Frontend-side representation of both, a export procedure in progress (or concluded) *and* the result of exporting.

    Each ``Export`` instance corresponds to a specific, individual ``job`` sent to the conversion service and thus
    encompasses

    - the spatial selection ('clipping' or 'extraction') of the input data within one perimeter
      (``extraction_order.excerpt``)
    - the transformation of the data from the data sources' schemata (e.g. ``osm2pgsql`` schema) to the OSMaxx schema
    - the actual export to one specific GIS or navigation file format with one specific set of parameters
    """
    INITIAL = 'initial'
    INITIAL_CHOICE = (INITIAL, _('initial'))
    STATUS_CHOICES = (INITIAL_CHOICE,) + statuses.STATUS_CHOICES

    extraction_order = models.ForeignKey('excerptexport.ExtractionOrder', related_name='exports',
                                         verbose_name=_('extraction order'))
    file_format = models.CharField(choices=FORMAT_CHOICES, verbose_name=_('file format / data format'), max_length=10)
    conversion_service_job_id = models.IntegerField(verbose_name=_('conversion service job ID'), null=True)
    status = models.CharField(_('job status'), choices=STATUS_CHOICES, default=INITIAL, max_length=20)
    finished = models.DateTimeField(_('finished on'), default=None, blank=True, editable=False, null=True)

    def send_to_conversion_service(self, clipping_area_json, incoming_request):
        from osmaxx.api_client.conversion_api_client import ConversionApiClient
        api_client = ConversionApiClient()
        extraction_format = self.file_format
        gis_options = self.extraction_order.extraction_configuration['gis_options']
        out_srs = gis_options['coordinate_reference_system']
        parametrization_json = api_client.create_parametrization(boundary=clipping_area_json, out_format=extraction_format, out_srs=out_srs)
        job_json = api_client.create_job(parametrization_json, self.get_full_status_update_uri(incoming_request))
        self.conversion_service_job_id = job_json['id']
        self.status = job_json['status']
        self.save()
        return job_json

    def get_full_status_update_uri(self, request):
        return request.build_absolute_uri(self.status_update_url)

    @property
    def status_update_url(self):
        return reverse('job_progress:tracker', kwargs=dict(export_id=self.id))

    def set_and_handle_new_status(self, new_status, *, incoming_request):
        assert new_status in dict(self.STATUS_CHOICES)
        if self.status != new_status:
            self.status = new_status
            self.save()
            self._handle_changed_status(incoming_request=incoming_request)

    def _handle_changed_status(self, *, incoming_request):
        from osmaxx.utilities.shortcuts import Emissary
        emissary = Emissary(recipient=self.extraction_order.orderer)
        status_changed_message = self._get_export_status_changed_message()
        if self.status == FAILED:
            emissary.error(status_changed_message)
        elif self.status == FINISHED:
            from osmaxx.api_client.conversion_api_client import ResultFileNotAvailableError
            try:
                self._fetch_result_file()
                emissary.success(status_changed_message)
            except ResultFileNotAvailableError:
                logger.error(self._get_job_finished_but_result_file_missing_log_message())
                emissary.warn(_("{} But the result file is not available.").format(status_changed_message))
        else:
            emissary.info(status_changed_message)
        self.extraction_order.send_email_if_all_exports_done(incoming_request)

    def _get_export_status_changed_message(self):
        from django.template.loader import render_to_string
        view_context = dict(export=self)
        return render_to_string(
            'job_progress/messages/export_status_changed.unsave_text',
            context=view_context,
        ).strip()

    def _get_job_finished_but_result_file_missing_log_message(self):
        return 'Export {export_id}: Job {job_id} finished, but file not available.'.format(
            export_id=self.id,
            job_id=self.conversion_service_job_id,
        )

    def _fetch_result_file(self):
        from osmaxx.api_client import ConversionApiClient
        from . import OutputFile
        api_client = ConversionApiClient()
        file_content = api_client.get_result_file(self.conversion_service_job_id)
        now = datetime.now()
        of = OutputFile.objects.create(
            export=self,
            mime_type='application/zip',
            file_extension='zip',
        )
        of.file.save(
            of.download_file_name,
            file_content,
        )
        self.finished = now
        self.save()

    @property
    def is_status_final(self):
        return self.status in FINAL_STATUSES

    @property
    def css_status_class(self):
        """
        based on the status, returns the bootstrap 3 class

        Returns: the bootstrap css class
        """
        default_class = 'default'

        status_map = {
            statuses.RECEIVED: 'info',
            statuses.QUEUED: 'info',
            statuses.FINISHED: 'success',
            statuses.FAILED: 'danger',
            statuses.STARTED: 'info',
            statuses.DEFERRED: 'default',
        }
        return status_map.get(self.status, default_class)
<commit_msg>fix bug: test on None instead of all falsy values<commit_after>import logging

from django.db import models
from django.utils.datetime_safe import datetime
from django.utils.translation import ugettext_lazy as _
from rest_framework.reverse import reverse

from osmaxx.conversion_api import statuses
from osmaxx.conversion_api.formats import FORMAT_CHOICES
from osmaxx.conversion_api.statuses import FAILED, FINAL_STATUSES, FINISHED

logger = logging.getLogger(__name__)

class TimeStampModelMixin(models.Model):
    created = models.DateTimeField(_('created on'), default=datetime.now, blank=True, editable=False)
    updated = models.DateTimeField(_('updated on'), default=None, blank=True, editable=False, null=True)

    def save(self, *args, **kwargs):
        now = datetime.now()
        if self.id is None:
            self.created = now
        self.updated = now
        super().save(*args, **kwargs)

    class Meta:
        abstract = True

class Export(TimeStampModelMixin, models.Model):
    """
    Frontend-side representation of both, a export procedure in progress (or concluded) *and* the result of exporting.

    Each ``Export`` instance corresponds to a specific, individual ``job`` sent to the conversion service and thus
    encompasses

    - the spatial selection ('clipping' or 'extraction') of the input data within one perimeter
      (``extraction_order.excerpt``)
    - the transformation of the data from the data sources' schemata (e.g. ``osm2pgsql`` schema) to the OSMaxx schema
    - the actual export to one specific GIS or navigation file format with one specific set of parameters
    """
    INITIAL = 'initial'
    INITIAL_CHOICE = (INITIAL, _('initial'))
    STATUS_CHOICES = (INITIAL_CHOICE,) + statuses.STATUS_CHOICES

    extraction_order = models.ForeignKey('excerptexport.ExtractionOrder', related_name='exports',
                                         verbose_name=_('extraction order'))
    file_format = models.CharField(choices=FORMAT_CHOICES, verbose_name=_('file format / data format'), max_length=10)
    conversion_service_job_id = models.IntegerField(verbose_name=_('conversion service job ID'), null=True)
    status = models.CharField(_('job status'), choices=STATUS_CHOICES, default=INITIAL, max_length=20)
    finished = models.DateTimeField(_('finished on'), default=None, blank=True, editable=False, null=True)

    def send_to_conversion_service(self, clipping_area_json, incoming_request):
        from osmaxx.api_client.conversion_api_client import ConversionApiClient
        api_client = ConversionApiClient()
        extraction_format = self.file_format
        gis_options = self.extraction_order.extraction_configuration['gis_options']
        out_srs = gis_options['coordinate_reference_system']
        parametrization_json = api_client.create_parametrization(boundary=clipping_area_json, out_format=extraction_format, out_srs=out_srs)
        job_json = api_client.create_job(parametrization_json, self.get_full_status_update_uri(incoming_request))
        self.conversion_service_job_id = job_json['id']
        self.status = job_json['status']
        self.save()
        return job_json

    def get_full_status_update_uri(self, request):
        return request.build_absolute_uri(self.status_update_url)

    @property
    def status_update_url(self):
        return reverse('job_progress:tracker', kwargs=dict(export_id=self.id))

    def set_and_handle_new_status(self, new_status, *, incoming_request):
        assert new_status in dict(self.STATUS_CHOICES)
        if self.status != new_status:
            self.status = new_status
            self.save()
            self._handle_changed_status(incoming_request=incoming_request)

    def _handle_changed_status(self, *, incoming_request):
        from osmaxx.utilities.shortcuts import Emissary
        emissary = Emissary(recipient=self.extraction_order.orderer)
        status_changed_message = self._get_export_status_changed_message()
        if self.status == FAILED:
            emissary.error(status_changed_message)
        elif self.status == FINISHED:
            from osmaxx.api_client.conversion_api_client import ResultFileNotAvailableError
            try:
                self._fetch_result_file()
                emissary.success(status_changed_message)
            except ResultFileNotAvailableError:
                logger.error(self._get_job_finished_but_result_file_missing_log_message())
                emissary.warn(_("{} But the result file is not available.").format(status_changed_message))
        else:
            emissary.info(status_changed_message)
        self.extraction_order.send_email_if_all_exports_done(incoming_request)

    def _get_export_status_changed_message(self):
        from django.template.loader import render_to_string
        view_context = dict(export=self)
        return render_to_string(
            'job_progress/messages/export_status_changed.unsave_text',
            context=view_context,
        ).strip()

    def _get_job_finished_but_result_file_missing_log_message(self):
        return 'Export {export_id}: Job {job_id} finished, but file not available.'.format(
            export_id=self.id,
            job_id=self.conversion_service_job_id,
        )

    def _fetch_result_file(self):
        from osmaxx.api_client import ConversionApiClient
        from . import OutputFile
        api_client = ConversionApiClient()
        file_content = api_client.get_result_file(self.conversion_service_job_id)
        now = datetime.now()
        of = OutputFile.objects.create(
            export=self,
            mime_type='application/zip',
            file_extension='zip',
        )
        of.file.save(
            of.download_file_name,
            file_content,
        )
        self.finished = now
        self.save()

    @property
    def is_status_final(self):
        return self.status in FINAL_STATUSES

    @property
    def css_status_class(self):
        """
        based on the status, returns the bootstrap 3 class

        Returns: the bootstrap css class
        """
        default_class = 'default'

        status_map = {
            statuses.RECEIVED: 'info',
            statuses.QUEUED: 'info',
            statuses.FINISHED: 'success',
            statuses.FAILED: 'danger',
            statuses.STARTED: 'info',
            statuses.DEFERRED: 'default',
        }
        return status_map.get(self.status, default_class)
<|endoftext|>


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

class TrajWriter(object):
    '''
    A class for writing out trajectory traces as an xyz file, for subsequent
    visualization.
    '''
    def __init__(self, trace, w, filename='trace.xyz'):
        self.trace = trace
        self.w = w
        self.filename = filename
        self._write()

    def _get_coords(self, iteration, seg_id):
        self.w.iteration = iteration
        coords = self.w.current.auxdata['coord'][seg_id]
        return coords

    def _write(self):
        all_coords = []
        starting_iteration = self.w.iteration
        for i, iteration in enumerate(self.trace.iteration):
            seg_id = self.trace.seg_id[i]
            coords = self._get_coords(iteration, seg_id)
            # The last timepoint of one iteration is the same as the first
            # timepoint of the last, so skip the last timepoint of each
            # iteration
            coords = coords[:-1]
            all_coords.append(coords)
        self.w.iteration = starting_iteration
        all_coords = numpy.concatenate(all_coords)
        with open(self.filename, 'w') as outfile:
            for i, frame in enumerate(all_coords):
                outfile.write("2\n")
                outfile.write("{0}\n".format(i))
                outfile.write("SOD {0:9.5f} {1:9.5f} {2:9.5f}\n".format(
                  float(frame[0,0]), float(frame[0,1]), float(frame[0,2])))
                outfile.write("CLA {0:9.5f} {1:9.5f} {2:9.5f}\n".format(
                  float(frame[1,0]), float(frame[1,1]), float(frame[1,2])))

<commit_msg>Updated docstring<commit_after>#!/usr/bin/env python
import numpy

class TrajWriter(object):
    '''
    A class for writing out trajectory traces as an xyz file, for subsequent
    visualization.

    ---------
    Arguments
    ---------
    trace: A trace returned by w.trace.
    w: The w object loaded by w_ipa

    -----------------
    Keyword Arguments
    -----------------
    filename (default='trace.xyz'): A path designating the file to which the
      trajectory should be written, in xyz format.
    '''
    def __init__(self, trace, w, filename='trace.xyz'):
        self.trace = trace
        self.w = w
        self.filename = filename
        self._write()

    def _get_coords(self, iteration, seg_id):
        self.w.iteration = iteration
        coords = self.w.current.auxdata['coord'][seg_id]
        return coords

    def _write(self):
        all_coords = []
        starting_iteration = self.w.iteration
        for i, iteration in enumerate(self.trace.iteration):
            seg_id = self.trace.seg_id[i]
            coords = self._get_coords(iteration, seg_id)
            # The last timepoint of one iteration is the same as the first
            # timepoint of the last, so skip the last timepoint of each
            # iteration
            coords = coords[:-1]
            all_coords.append(coords)
        self.w.iteration = starting_iteration
        all_coords = numpy.concatenate(all_coords)
        with open(self.filename, 'w') as outfile:
            for i, frame in enumerate(all_coords):
                outfile.write("2\n")
                outfile.write("{0}\n".format(i))
                outfile.write("SOD {0:9.5f} {1:9.5f} {2:9.5f}\n".format(
                  float(frame[0,0]), float(frame[0,1]), float(frame[0,2])))
                outfile.write("CLA {0:9.5f} {1:9.5f} {2:9.5f}\n".format(
                  float(frame[1,0]), float(frame[1,1]), float(frame[1,2])))

<|endoftext|>


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

reload(sys)
sys.setdefaultencoding('utf8')

built_in = ['abs','dict','help','min','setattr','all','dir','hex','next',
'slice','any','divmod','id','object','sorted','ascii','enumerate','input',
'oct','staticmethod','bin','eval','int','open','str','bool','exec','isinstance',
'ord','sum','bytearray','filter','issubclass','pow','super','bytes','float',
'iter','print','tuple','callable','format','len','property','type','chr',
'frozenset','list','range','vars','classmethod','getattr','locals','repr',
'zip','compile','globals','map','reversed','import','complex','hasattr',
'max','round','delattr','hash','memoryview','set']

ignore_keys = ['The Python Tutorial']

bad_strings = ['\\000', '\\xe0']

class BadEntryException(Exception):
    """
    Thrown when entry data is invalid
    """
    pass

class Entry(object):
    """
    Represents a single entry in fathead output.txt datafile
    """
    def __init__(self, obj):
        # Type of entry. R = Redirect, A = Article.
        self.entry_type = ''

        # Unique key
        self.key = ''

        # Referenced entry key (only for redirects)
        self.reference = ''
        self.abstract = ''
        self.anchor = ''

        self.category = ''
        self.related = ''

        # Alternative keys used for redirects
        self.alternative_keys = []
        self.parse(obj)

    def __str__(self, obj):
        return self.alternative_keys

    def parse(self, input_obj):
        """
        Try to parse given input into a valid entry.
        Args:
            input_obj: TSV string or list of data.

        Returns:
            List of data
        Raises:
            Throws BadEntryException if data is invalid.
        """
        if isinstance(input_obj, str):
            processed = input_obj.split('\t')
            self.data = processed
        elif isinstance(input_obj, list):
            self.data = input_obj
        try:
            self.key = self.data[0].strip()
            self.entry_type = self.data[1].strip()
            self.reference = self.data[2].strip()
            if len(self.data) > 3:
                self.category = self.data[4].strip()
                self.related = self.data[6].strip()
                self.abstract = self.data[11].strip()
                self.anchor = self.data[12].strip()
            elif self.entry_type == 13 and len(self.data) != 13:
                raise BadEntryException
        except Exception as e:
            raise BadEntryException('Article had invalid number of elements.')

        if self.entry_type == 'A':
            self.parse_alternative_keys()

        return self.data

    def parse_alternative_keys(self):
        """
        Find alternative keys to use in generated redirects

        Returns:
            Set of possible redirect entries
        """
        self.alternative_keys = set()

        if '.' in self.key and self.entry_type == 'A':
            key_arr = self.key.split('.')
            method_name = key_arr[-1]
            key_arr_len = len(key_arr)

            # always add method name as a key
            self.alternative_keys.add(method_name)

            # add all permutations of package and class
            if key_arr_len >= 3:

                for l in range(key_arr_len-1):
                    permutations = itertools.permutations(key_arr[:key_arr_len-1], l+1)

                    for k in permutations:
                        new_key = "{} {}".format(' '.join(k), method_name)
                        self.alternative_keys.add(new_key)
            else:
                package_name = key_arr[0]
                new_key = "{} {}".format(package_name, method_name)
                self.alternative_keys.add(new_key)

        return self.alternative_keys

    def get_data(self):
        return self.data

    def get_key(self):
        return self.key

    def set_entry_type(self, new_type):
        self.entry_type = new_type

    def get_type(self):
        return self.entry_type

    def get_redirects(self):
        redirs = []
        for alt_key in list(self.alternative_keys):
            if alt_key in built_in:
                continue
            entry = Entry([
                    alt_key,
                    'R',
                    self.key
                ])
            redirs.append(entry)

        return redirs

    def get_alternatives(self):
        return self.alternative_keys

    def get_abstract(self):
        return self.abstract

    def get_entry(self):
        return '\t'.join([
            self.key,             # title / key
            self.entry_type,      # entry type
            self.reference,       # no redirect data
            '',                   # ignore
            self.category,        # categories
            '',                   # ignore
            self.related,                   # no related topics
            '',                   # ignore
            '',                   # add an external link back to page
            '',                   # no disambiguation
            '',                   # images
            self.abstract,        # abstract
            self.anchor           # anchor to specific section
        ])

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

def generate_redirects(f):
    output = dict()

    # For debugging purposes
    duplicate_count = 0
    disambiguations = 0
    for line in f.readlines():
        try:
            # Parse entry
            entry = Entry(line)

            if entry.get_type() == 'R' or entry.get_key() in ignore_keys:
                continue

            key = "'" + entry.get_key() + "'"

            # Do we have the entry yet?
            if key not in output or '3.5/library/functions.html' in str(entry):
                output[key] = str(entry)

            disambiguation_key = entry.get_key()
            splitted_key = disambiguation_key.split('.')
            package_name = '.'.join(splitted_key[0:len(splitted_key)-1])
            temp = "'" + package_name + "'"

            if len(splitted_key) >= 2:
                if temp in output and output[temp].startswith(package_name + '\t' + 'A'):
                    output[temp] = str(package_name) + '\t' + 'D' +'\t\t\t\t\t\t\t\t' + '*' + '[['+str(disambiguation_key)+']] ' + str(entry.get_abstract()) + '\\n'
                    if output[temp].startswith(package_name + '\t' + 'A'):
                        del output[temp]
                    disambiguations += 1
                elif temp in output and disambiguation_key not in output[temp] and output[temp].startswith(package_name + '\t' + 'D'):
                    output[temp] += '*' + '[['+str(disambiguation_key)+']] ' + str(entry.get_abstract()) + '\\n'

            # Get all possible redirect entries
            key = entry.get_key()
            key_length = len((re.findall(r"[\w']+", key)))
            if key_length > 1:
                redirects = entry.get_redirects()
                for redirect in redirects:
                    key = redirect.get_key()
                    built_in_key = '"' + redirect.get_key() + '"'
                    if key not in output and built_in_key not in built_in:
                        output[key] = str(redirect.get_entry())
                    else:
                        del output[key]
                        duplicate_count += 1

        except BadEntryException as e:
            pass  # Continue execution entry data is invalid.

    print("Duplicates: %s" % duplicate_count)
    print("Disambiguations: %s" % disambiguations)

    with open('output2.txt', 'w') as output_file:
        for key, line in output.items():
            if line.endswith('\\n'):
                line = line[:-2]
                line += '\t\t\t'
            line = bad_string_check(line)
            tsv = '{}\n'.format(line)
            output_file.write(tsv)

def bad_string_check(line):
    for string in bad_strings:
        if string in line:
            replacement = '\\' + string
            line = line.replace(string, replacement)
    return line

if __name__ == "__main__":
    # Open output file for reading and writing.
    with open('output.txt', 'r+') as output_file:
        generate_redirects(output_file)
<commit_msg>Add placeholder comment<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import itertools
import re
import sys

reload(sys)
sys.setdefaultencoding('utf8')

built_in = ['abs','dict','help','min','setattr','all','dir','hex','next',
'slice','any','divmod','id','object','sorted','ascii','enumerate','input',
'oct','staticmethod','bin','eval','int','open','str','bool','exec','isinstance',
'ord','sum','bytearray','filter','issubclass','pow','super','bytes','float',
'iter','print','tuple','callable','format','len','property','type','chr',
'frozenset','list','range','vars','classmethod','getattr','locals','repr',
'zip','compile','globals','map','reversed','import','complex','hasattr',
'max','round','delattr','hash','memoryview','set']

ignore_keys = ['The Python Tutorial']

bad_strings = ['\\000', '\\xe0']

class BadEntryException(Exception):
    """
    Thrown when entry data is invalid
    """
    pass

class Entry(object):
    """
    Represents a single entry in fathead output.txt datafile
    """
    def __init__(self, obj):
        # Type of entry. R = Redirect, A = Article.
        self.entry_type = ''

        # Unique key
        self.key = ''

        # Referenced entry key (only for redirects)
        self.reference = ''
        self.abstract = ''
        self.anchor = ''

        self.category = ''
        self.related = ''

        # Alternative keys used for redirects
        self.alternative_keys = []
        self.parse(obj)

    def __str__(self, obj):
        return self.alternative_keys

    def parse(self, input_obj):
        """
        Try to parse given input into a valid entry.
        Args:
            input_obj: TSV string or list of data.

        Returns:
            List of data
        Raises:
            Throws BadEntryException if data is invalid.
        """
        if isinstance(input_obj, str):
            processed = input_obj.split('\t')
            self.data = processed
        elif isinstance(input_obj, list):
            self.data = input_obj
        try:
            self.key = self.data[0].strip()
            self.entry_type = self.data[1].strip()
            self.reference = self.data[2].strip()
            if len(self.data) > 3:
                self.category = self.data[4].strip()
                self.related = self.data[6].strip()
                self.abstract = self.data[11].strip()
                self.anchor = self.data[12].strip()
            elif self.entry_type == 13 and len(self.data) != 13:
                raise BadEntryException
        except Exception as e:
            raise BadEntryException('Article had invalid number of elements.')

        if self.entry_type == 'A':
            self.parse_alternative_keys()

        return self.data

    def parse_alternative_keys(self):
        """
        Find alternative keys to use in generated redirects

        Returns:
            Set of possible redirect entries
        """
        self.alternative_keys = set()

        if '.' in self.key and self.entry_type == 'A':
            key_arr = self.key.split('.')
            method_name = key_arr[-1]
            key_arr_len = len(key_arr)

            # always add method name as a key
            self.alternative_keys.add(method_name)

            # add all permutations of package and class
            if key_arr_len >= 3:

                for l in range(key_arr_len-1):
                    permutations = itertools.permutations(key_arr[:key_arr_len-1], l+1)

                    for k in permutations:
                        new_key = "{} {}".format(' '.join(k), method_name)
                        self.alternative_keys.add(new_key)
            else:
                package_name = key_arr[0]
                new_key = "{} {}".format(package_name, method_name)
                self.alternative_keys.add(new_key)

        return self.alternative_keys

    def get_data(self):
        return self.data

    def get_key(self):
        return self.key

    def set_entry_type(self, new_type):
        self.entry_type = new_type

    def get_type(self):
        return self.entry_type

    def get_redirects(self):
        redirs = []
        for alt_key in list(self.alternative_keys):
            if alt_key in built_in:
                continue
            entry = Entry([
                    alt_key,
                    'R',
                    self.key
                ])
            redirs.append(entry)

        return redirs

    def get_alternatives(self):
        return self.alternative_keys

    def get_abstract(self):
        return self.abstract

    def get_entry(self):
        return '\t'.join([
            self.key,             # title / key
            self.entry_type,      # entry type
            self.reference,       # no redirect data
            '',                   # ignore
            self.category,        # categories
            '',                   # ignore
            self.related,                   # no related topics
            '',                   # ignore
            '',                   # add an external link back to page
            '',                   # no disambiguation
            '',                   # images
            self.abstract,        # abstract
            self.anchor           # anchor to specific section
        ])

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

def generate_redirects(f):
    output = dict()

    # For debugging purposes
    duplicate_count = 0
    disambiguations = 0
    for line in f.readlines():
        try:
            # Parse entry
            entry = Entry(line)

            if entry.get_type() == 'R' or entry.get_key() in ignore_keys:
                continue

            key = "'" + entry.get_key() + "'"

            # Do we have the entry yet?
            if key not in output or '3.5/library/functions.html' in str(entry):
                output[key] = str(entry)

            disambiguation_key = entry.get_key()
            splitted_key = disambiguation_key.split('.')
            package_name = '.'.join(splitted_key[0:len(splitted_key)-1])
            #temp doesn't sound like a good variable name
            temp = "'" + package_name + "'"

            if len(splitted_key) >= 2:
                if temp in output and output[temp].startswith(package_name + '\t' + 'A'):
                    output[temp] = str(package_name) + '\t' + 'D' +'\t\t\t\t\t\t\t\t' + '*' + '[['+str(disambiguation_key)+']] ' + str(entry.get_abstract()) + '\\n'
                    if output[temp].startswith(package_name + '\t' + 'A'):
                        del output[temp]
                    disambiguations += 1
                elif temp in output and disambiguation_key not in output[temp] and output[temp].startswith(package_name + '\t' + 'D'):
                    output[temp] += '*' + '[['+str(disambiguation_key)+']] ' + str(entry.get_abstract()) + '\\n'

            # Get all possible redirect entries
            key = entry.get_key()
            key_length = len((re.findall(r"[\w']+", key)))
            if key_length > 1:
                redirects = entry.get_redirects()
                for redirect in redirects:
                    key = redirect.get_key()
                    built_in_key = '"' + redirect.get_key() + '"'
                    if key not in output and built_in_key not in built_in:
                        output[key] = str(redirect.get_entry())
                    else:
                        del output[key]
                        duplicate_count += 1

        except BadEntryException as e:
            pass  # Continue execution entry data is invalid.

    print("Duplicates: %s" % duplicate_count)
    print("Disambiguations: %s" % disambiguations)

    with open('output2.txt', 'w') as output_file:
        for key, line in output.items():
            if line.endswith('\\n'):
                line = line[:-2]
                line += '\t\t\t'
            line = bad_string_check(line)
            tsv = '{}\n'.format(line)
            output_file.write(tsv)

def bad_string_check(line):
    for string in bad_strings:
        if string in line:
            replacement = '\\' + string
            line = line.replace(string, replacement)
    return line

if __name__ == "__main__":
    # Open output file for reading and writing.
    with open('output.txt', 'r+') as output_file:
        generate_redirects(output_file)
<|endoftext|>


<commit_before>import binaryninja as binja
from binaryninja.enums import (
    Endianness, SegmentFlag
)
import inspect
import logging
import magic
import re
import struct

LOGNAME = 'binja.cfg'
log = logging.getLogger(LOGNAME)

class StackFormatter(logging.Formatter):
    def __init__(self, fmt=None, datefmt=None):
        logging.Formatter.__init__(self, fmt, datefmt)
        self.stack_base = len(inspect.stack()) + 7

    def format(self, record):
        record.indent = '  ' * (len(inspect.stack()) - self.stack_base)
        res = logging.Formatter.format(self, record)
        del record.indent
        return res

def init_logger(log_file):
    formatter = StackFormatter('[%(levelname)s] %(indent)s%(message)s')
    handler = logging.FileHandler(log_file)
    handler.setFormatter(formatter)

    log.addHandler(handler)
    log.setLevel(logging.DEBUG)

ENDIAN_TO_STRUCT = {
    Endianness.LittleEndian: '<',
    Endianness.BigEndian: '>'
}

def read_dword(bv, addr):
    # type: (binja.BinaryView, int) -> int
    data = bv.read(addr, 4)
    fmt = '{}L'.format(ENDIAN_TO_STRUCT[bv.endianness])
    return struct.unpack(fmt, data)[0]

def read_qword(bv, addr):
    # type: (binja.BinaryView, int) -> int
    data = bv.read(addr, 8)
    fmt = '{}Q'.format(ENDIAN_TO_STRUCT[bv.endianness])
    return struct.unpack(fmt, data)[0]

def load_binary(path):
    magic_type = magic.from_file(path)
    if 'ELF' in magic_type:
        bv_type = binja.BinaryViewType['ELF']
    elif 'PE32' in magic_type:
        bv_type = binja.BinaryViewType['PE']
    elif 'Mach-O' in magic_type:
        bv_type = binja.BinaryViewType['Mach-O']
    else:
        bv_type = binja.BinaryViewType['Raw']

        # Can't do anything with Raw type
        log.fatal('Unknown binary type: "{}", exiting'.format(magic_type))
        exit(1)

    log.debug('Loading binary in binja...')
    bv = bv_type.open(path)
    bv.update_analysis_and_wait()

    # NOTE: at the moment binja will not load a binary
    # that doesn't have an entry point
    if len(bv) == 0:
        log.error('Binary could not be loaded in binja, is it linked?')
        exit(1)

    return bv

def find_symbol_name(bv, addr):
    """Attempt to find a symbol for a given address

    Args:
        bv (binja.BinaryView)
        addr (int): Address the symbol should point to

    Returns:
        (str): Symbol name if found, empty string otherwise

    """
    sym = bv.get_symbol_at(addr)
    if sym is not None:
        return sym.name
    return ''

def get_func_containing(bv, addr):
    """ Finds the function, if any, containing the given address
    Args:
        bv (binja.BinaryView)
        addr (int)

    Returns:
        binja.Function
    """
    funcs = bv.get_functions_containing(addr)
    return funcs[0] if funcs is not None else None

def is_external_ref(bv, addr):
    sym = bv.get_symbol_at(addr)
    return sym is not None and 'Import' in sym.type.name

def is_valid_addr(bv, addr):
    return bv.get_segment_at(addr) is not None

def is_code(bv, addr):
    """Returns `True` if the given address lies in an executable segment"""
    return (bv.get_segment_at(addr).flags & SegmentFlag.SegmentExecutable) != 0

def is_readable(bv, addr):
    """Returns `True` if the given address lies in a readable segment"""
    return (bv.get_segment_at(addr).flags & SegmentFlag.SegmentReadable) != 0

def is_writeable(bv, addr):
    """Returns `True` if the given address lies in a writable segment"""
    return (bv.get_segment_at(addr).flags & SegmentFlag.SegmentWritable) != 0

def is_ELF(bv):
    return bv.view_type == 'ELF'

def is_PE(bv):
    return bv.view_type == 'PE'

def clamp(val, vmin, vmax):
    return min(vmax, max(vmin, val))

# Caching results of is_section_external
_EXT_SECTIONS = set()
_INT_SECTIONS = set()

def is_section_external(bv, sect):
    """Returns `True` if the given section contains only external references

    Args:
        bv (binja.BinaryView)
        sect (binja.binaryview.Section)
    """
    if sect.start in _EXT_SECTIONS:
        return True

    if sect.start in _INT_SECTIONS:
        return False

    if is_ELF(bv):
        if re.search(r'\.(got|plt)', sect.name):
            _EXT_SECTIONS.add(sect.start)
            return True

    if is_PE(bv):
        if '.idata' in sect.name:
            _EXT_SECTIONS.add(sect.start)
            return True

    _INT_SECTIONS.add(sect.start)
    return False

def is_tls_section(bv, addr):
    sect_names = (sect.name for sect in bv.get_sections_at(addr))
    return any(sect in ['.tbss', '.tdata', '.tls'] for sect in sect_names)
<commit_msg>API update for checking segment flags<commit_after>import binaryninja as binja
from binaryninja.enums import (
    Endianness, SegmentFlag
)
import inspect
import logging
import magic
import re
import struct

LOGNAME = 'binja.cfg'
log = logging.getLogger(LOGNAME)

class StackFormatter(logging.Formatter):
    def __init__(self, fmt=None, datefmt=None):
        logging.Formatter.__init__(self, fmt, datefmt)
        self.stack_base = len(inspect.stack()) + 7

    def format(self, record):
        record.indent = '  ' * (len(inspect.stack()) - self.stack_base)
        res = logging.Formatter.format(self, record)
        del record.indent
        return res

def init_logger(log_file):
    formatter = StackFormatter('[%(levelname)s] %(indent)s%(message)s')
    handler = logging.FileHandler(log_file)
    handler.setFormatter(formatter)

    log.addHandler(handler)
    log.setLevel(logging.DEBUG)

ENDIAN_TO_STRUCT = {
    Endianness.LittleEndian: '<',
    Endianness.BigEndian: '>'
}

def read_dword(bv, addr):
    # type: (binja.BinaryView, int) -> int
    data = bv.read(addr, 4)
    fmt = '{}L'.format(ENDIAN_TO_STRUCT[bv.endianness])
    return struct.unpack(fmt, data)[0]

def read_qword(bv, addr):
    # type: (binja.BinaryView, int) -> int
    data = bv.read(addr, 8)
    fmt = '{}Q'.format(ENDIAN_TO_STRUCT[bv.endianness])
    return struct.unpack(fmt, data)[0]

def load_binary(path):
    magic_type = magic.from_file(path)
    if 'ELF' in magic_type:
        bv_type = binja.BinaryViewType['ELF']
    elif 'PE32' in magic_type:
        bv_type = binja.BinaryViewType['PE']
    elif 'Mach-O' in magic_type:
        bv_type = binja.BinaryViewType['Mach-O']
    else:
        bv_type = binja.BinaryViewType['Raw']

        # Can't do anything with Raw type
        log.fatal('Unknown binary type: "{}", exiting'.format(magic_type))
        exit(1)

    log.debug('Loading binary in binja...')
    bv = bv_type.open(path)
    bv.update_analysis_and_wait()

    # NOTE: at the moment binja will not load a binary
    # that doesn't have an entry point
    if len(bv) == 0:
        log.error('Binary could not be loaded in binja, is it linked?')
        exit(1)

    return bv

def find_symbol_name(bv, addr):
    """Attempt to find a symbol for a given address

    Args:
        bv (binja.BinaryView)
        addr (int): Address the symbol should point to

    Returns:
        (str): Symbol name if found, empty string otherwise

    """
    sym = bv.get_symbol_at(addr)
    if sym is not None:
        return sym.name
    return ''

def get_func_containing(bv, addr):
    """ Finds the function, if any, containing the given address
    Args:
        bv (binja.BinaryView)
        addr (int)

    Returns:
        binja.Function
    """
    funcs = bv.get_functions_containing(addr)
    return funcs[0] if funcs is not None else None

def is_external_ref(bv, addr):
    sym = bv.get_symbol_at(addr)
    return sym is not None and 'Import' in sym.type.name

def is_valid_addr(bv, addr):
    return bv.get_segment_at(addr) is not None

def is_code(bv, addr):
    """Returns `True` if the given address lies in an executable segment"""
    seg = bv.get_segment_at(addr)
    return seg is not None and seg.executable

def is_readable(bv, addr):
    """Returns `True` if the given address lies in a readable segment"""
    seg = bv.get_segment_at(addr)
    return seg is not None and seg.writable

def is_writeable(bv, addr):
    """Returns `True` if the given address lies in a writable segment"""
    seg = bv.get_segment_at(addr)
    return seg is not None and seg.readable

def is_ELF(bv):
    return bv.view_type == 'ELF'

def is_PE(bv):
    return bv.view_type == 'PE'

def clamp(val, vmin, vmax):
    return min(vmax, max(vmin, val))

# Caching results of is_section_external
_EXT_SECTIONS = set()
_INT_SECTIONS = set()

def is_section_external(bv, sect):
    """Returns `True` if the given section contains only external references

    Args:
        bv (binja.BinaryView)
        sect (binja.binaryview.Section)
    """
    if sect.start in _EXT_SECTIONS:
        return True

    if sect.start in _INT_SECTIONS:
        return False

    if is_ELF(bv):
        if re.search(r'\.(got|plt)', sect.name):
            _EXT_SECTIONS.add(sect.start)
            return True

    if is_PE(bv):
        if '.idata' in sect.name:
            _EXT_SECTIONS.add(sect.start)
            return True

    _INT_SECTIONS.add(sect.start)
    return False

def is_tls_section(bv, addr):
    sect_names = (sect.name for sect in bv.get_sections_at(addr))
    return any(sect in ['.tbss', '.tdata', '.tls'] for sect in sect_names)
<|endoftext|>


<commit_before>import logging

logger = logging.getLogger(None if __name__ == '__main__' else __name__)

class ConstituencyTree:
    """
    Group together information from a constituency tree and expose convenient
    lookups.
    """

    def __init__(self, head2deps):
        """
        Initialise `ConstituencyTree`
        """
        self.head2deps = head2deps

        # Create the reverse too
        self.dep2heads = self.reverse_headdep_dict(head2deps)

        logger.debug("head2deps: {}".format(self.head2deps))
        logger.debug("dep2heads: {}".format(self.dep2heads))
        self._head2constituent = {}

    def __repr__(self):
        return self.__class__.__name__ + '({!r})'.format(self.head2deps)

    def __eq__(self, other):
        if isinstance(other, type(self)):
            return self.head2deps == other.head2deps
        return False

    @classmethod
    def from_naf(cls, nafobj, term_filter=lambda naf, t: True):
        """
        Initialise this ConstituencyTree from a NAF object

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :param term_filter:     filter for terms
        """
        unfiltered = cls.create_headdep_dict(nafobj)
        return cls(cls.filter_headdep_dict(
            unfiltered,
            lambda t: term_filter(nafobj, t)
        ))

    @staticmethod
    def reverse_headdep_dict(head2deps):
        """
        Create dep2heads from head2deps (or visa versa?)
        """
        dep2heads = {}
        for headID, deps in head2deps.items():
            for toID, relation in deps:
                dep2heads.setdefault(toID, set()).add((headID, relation))
        return dep2heads

    def get_direct_dependents(self, ID):
        """
        Get the term IDs of the terms directly dependent on `ID`.

        :param headID:  term ID to get direct dependents of
        :return:        a set of IDs that are direct dependents of `ID`
        """
        if ID in self.head2deps:
            return {ID for ID, _ in self.head2deps[ID]}

    def get_direct_parents(self, ID):
        """
        Get the term IDs of the terms of which `ID` is a direct dependent.

        :param headID:  term ID to get direct parents of
        :return:        a set of IDs that are direct parents of `ID`
        """
        if ID in self.dep2heads:
            return {ID for ID, _ in self.dep2heads[ID]}

    def get_constituent(self, ID):
        """
        Get the term IDs of the terms dependent on `ID`.

        Always contains `ID`, even if it is unknown.

        :param headID:  term ID to get constituent of
        :return:        a set of IDs of terms that are dependents of `ID`
        """
        try:
            return self._head2constituent[ID]
        except KeyError:
            return self._get_constituent(ID, set())

    def _get_constituent(self, ID, parents):
        """
        Calculate the constituent if it is not in the cache

        :param head:      head to find constituent of
        :param parents:   term IDs that are already being calculated, thus
                          should not be recursed into
        :return:          set of term IDs that are in the constituent of `head`
        """
        if ID in parents:
            # Prevent loop
            return set()
        elif ID in self._head2constituent:
            # Use cache
            return self._head2constituent[ID]
        # Base case: direct dependents
        deps = {term_id for term_id, _ in self.head2deps.get(ID, [])}
        # Recursive step: call _get_constituent for every direct dependent
        parents.add(ID)
        deps.update(*(
            self._get_constituent(term_id, parents)
            for term_id in deps
        ))
        # Make sure ID itself is in it
        deps.add(ID)
        # Cache
        self._head2constituent[ID] = deps
        return deps

    @staticmethod
    def create_headdep_dict(nafobj):
        """
        Create dictionary of head to direct dependents

        :param nafobj:          NAF object from input
        :return:                head2deps
        """

        head2deps = {}
        for dep in nafobj.get_dependencies():
            head2deps.setdefault(dep.get_from(), set()).add(
                (dep.get_to(), dep.get_function())
            )
        return head2deps

    @staticmethod
    def filter_headdep_dict(head2deps, term_filter):
        """
        Only keep terms where `term_filter(term)` evaluates True.

        If a term is removed, the link (or function) information is handled as
        follows:

            from  --parent info-->  removed  --child info-->  to

        becomes

            from  --child info-->  to

        !! NB !! If a circular reference exists in the dependency tree, this
                 will stay in, even if some of the terms in the circular
                 reference are filtered out.

        :param head2deps:       {head: {(dep, function), ...}}
        :param term_filter:     filter for terms
        """
        dep2headIDs = {}
        for headID, deps in head2deps.items():
            for toID, _ in deps:
                dep2headIDs.setdefault(toID, []).append(headID)
        logger.debug("dep2headIDs: {}".format(dep2headIDs))
        filtered = {}
        for headID, deps in head2deps.items():
            # I don't have to do something with the deps that are filtered out,
            # because if they are leaves they can be left out and if they
            # aren't leaves they will also appear as headID and handled there.
            logger.debug("headID: {}".format(headID))
            deps = {
                (toID, relation)
                for toID, relation in deps
                if term_filter(toID)
            }
            logger.debug("deps: {}".format(deps))
            if term_filter(headID):
                if deps:
                    filtered.setdefault(headID, set()).update(deps)
            elif deps:
                # Delete the head by adding its dependents to the most shallow
                # head that isn't filtered out.
                stack = dep2headIDs.get(headID, [])
                been_in_stack = set(stack)
                super_heads = []
                while stack:
                    logger.debug("stack: {}".format(stack))
                    super_head = stack.pop()
                    if term_filter(super_head):
                        super_heads.append(super_head)
                    else:
                        add_to_stack = [
                            h for h in dep2headIDs.get(super_head, [])
                            if h not in been_in_stack
                        ]
                        been_in_stack.update(add_to_stack)
                        logger.debug("add_to_stack: {}".format(add_to_stack))
                        stack.extend(
                            add_to_stack
                        )
                for super_head in super_heads:
                    filtered.setdefault(super_head, set()).update(deps)

        return filtered
<commit_msg>Add type annotation to `ConstituencyTree.from_naf`<commit_after>import logging

logger = logging.getLogger(None if __name__ == '__main__' else __name__)

class ConstituencyTree:
    """
    Group together information from a constituency tree and expose convenient
    lookups.
    """

    def __init__(self, head2deps):
        """
        Initialise `ConstituencyTree`
        """
        self.head2deps = head2deps

        # Create the reverse too
        self.dep2heads = self.reverse_headdep_dict(head2deps)

        logger.debug("head2deps: {}".format(self.head2deps))
        logger.debug("dep2heads: {}".format(self.dep2heads))
        self._head2constituent = {}

    def __repr__(self):
        return self.__class__.__name__ + '({!r})'.format(self.head2deps)

    def __eq__(self, other):
        if isinstance(other, type(self)):
            return self.head2deps == other.head2deps
        return False

    @classmethod
    def from_naf(cls, nafobj, term_filter=lambda naf, t: True):
        """
        Initialise this ConstituencyTree from a NAF object

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :type  nafobj:          KafNafParser
        :param term_filter:     filter for terms
        :type  term_filter:     Callable[[KafNafParser, str], bool]
        """
        unfiltered = cls.create_headdep_dict(nafobj)
        return cls(cls.filter_headdep_dict(
            unfiltered,
            lambda t: term_filter(nafobj, t)
        ))

    @staticmethod
    def reverse_headdep_dict(head2deps):
        """
        Create dep2heads from head2deps (or visa versa?)
        """
        dep2heads = {}
        for headID, deps in head2deps.items():
            for toID, relation in deps:
                dep2heads.setdefault(toID, set()).add((headID, relation))
        return dep2heads

    def get_direct_dependents(self, ID):
        """
        Get the term IDs of the terms directly dependent on `ID`.

        :param headID:  term ID to get direct dependents of
        :return:        a set of IDs that are direct dependents of `ID`
        """
        if ID in self.head2deps:
            return {ID for ID, _ in self.head2deps[ID]}

    def get_direct_parents(self, ID):
        """
        Get the term IDs of the terms of which `ID` is a direct dependent.

        :param headID:  term ID to get direct parents of
        :return:        a set of IDs that are direct parents of `ID`
        """
        if ID in self.dep2heads:
            return {ID for ID, _ in self.dep2heads[ID]}

    def get_constituent(self, ID):
        """
        Get the term IDs of the terms dependent on `ID`.

        Always contains `ID`, even if it is unknown.

        :param headID:  term ID to get constituent of
        :return:        a set of IDs of terms that are dependents of `ID`
        """
        try:
            return self._head2constituent[ID]
        except KeyError:
            return self._get_constituent(ID, set())

    def _get_constituent(self, ID, parents):
        """
        Calculate the constituent if it is not in the cache

        :param head:      head to find constituent of
        :param parents:   term IDs that are already being calculated, thus
                          should not be recursed into
        :return:          set of term IDs that are in the constituent of `head`
        """
        if ID in parents:
            # Prevent loop
            return set()
        elif ID in self._head2constituent:
            # Use cache
            return self._head2constituent[ID]
        # Base case: direct dependents
        deps = {term_id for term_id, _ in self.head2deps.get(ID, [])}
        # Recursive step: call _get_constituent for every direct dependent
        parents.add(ID)
        deps.update(*(
            self._get_constituent(term_id, parents)
            for term_id in deps
        ))
        # Make sure ID itself is in it
        deps.add(ID)
        # Cache
        self._head2constituent[ID] = deps
        return deps

    @staticmethod
    def create_headdep_dict(nafobj):
        """
        Create dictionary of head to direct dependents

        :param nafobj:          NAF object from input
        :return:                head2deps
        """

        head2deps = {}
        for dep in nafobj.get_dependencies():
            head2deps.setdefault(dep.get_from(), set()).add(
                (dep.get_to(), dep.get_function())
            )
        return head2deps

    @staticmethod
    def filter_headdep_dict(head2deps, term_filter):
        """
        Only keep terms where `term_filter(term)` evaluates True.

        If a term is removed, the link (or function) information is handled as
        follows:

            from  --parent info-->  removed  --child info-->  to

        becomes

            from  --child info-->  to

        !! NB !! If a circular reference exists in the dependency tree, this
                 will stay in, even if some of the terms in the circular
                 reference are filtered out.

        :param head2deps:       {head: {(dep, function), ...}}
        :param term_filter:     filter for terms
        """
        dep2headIDs = {}
        for headID, deps in head2deps.items():
            for toID, _ in deps:
                dep2headIDs.setdefault(toID, []).append(headID)
        logger.debug("dep2headIDs: {}".format(dep2headIDs))
        filtered = {}
        for headID, deps in head2deps.items():
            # I don't have to do something with the deps that are filtered out,
            # because if they are leaves they can be left out and if they
            # aren't leaves they will also appear as headID and handled there.
            logger.debug("headID: {}".format(headID))
            deps = {
                (toID, relation)
                for toID, relation in deps
                if term_filter(toID)
            }
            logger.debug("deps: {}".format(deps))
            if term_filter(headID):
                if deps:
                    filtered.setdefault(headID, set()).update(deps)
            elif deps:
                # Delete the head by adding its dependents to the most shallow
                # head that isn't filtered out.
                stack = dep2headIDs.get(headID, [])
                been_in_stack = set(stack)
                super_heads = []
                while stack:
                    logger.debug("stack: {}".format(stack))
                    super_head = stack.pop()
                    if term_filter(super_head):
                        super_heads.append(super_head)
                    else:
                        add_to_stack = [
                            h for h in dep2headIDs.get(super_head, [])
                            if h not in been_in_stack
                        ]
                        been_in_stack.update(add_to_stack)
                        logger.debug("add_to_stack: {}".format(add_to_stack))
                        stack.extend(
                            add_to_stack
                        )
                for super_head in super_heads:
                    filtered.setdefault(super_head, set()).update(deps)

        return filtered
<|endoftext|>


<commit_before>import logging

logger = logging.getLogger(None if __name__ == '__main__' else __name__)

class ConstituencyTree:
    """
    Group together information from a constituency tree and expose convenient
    lookups.
    """

    def __init__(self, head2deps):
        """
        Initialise `ConstituencyTree`
        """
        self.head2deps = head2deps

        # Create the reverse too
        self.dep2heads = {}
        for headID, deps in head2deps.items():
            for toID, relation in deps:
                self.dep2heads.setdefault(toID, set()).add((headID, relation))

        logger.debug("head2deps: {}".format(self.head2deps))
        logger.debug("dep2heads: {}".format(self.dep2heads))
        self._head2constituent = {}

    def __repr__(self):
        return self.__class__.__name__ + '({!r})'.format(self.head2deps)

    @classmethod
    def from_naf(cls, nafobj, term_filter=lambda naf, t: True):
        """
        Initialise this ConstituencyTree from a NAF object

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :param term_filter:     filter for terms
        """
        return cls(cls.create_headdep_dict(
            nafobj,
            term_filter
        ))

    def get_direct_dependents(self, ID):
        """
        Get the term IDs of the terms directly dependent on `ID`.

        :param headID:  term ID to get direct dependents of
        :return:        a set of IDs that are direct dependents of `ID`
        """
        if ID in self.head2deps:
            return {ID for ID, _ in self.head2deps[ID]}

    def get_direct_parents(self, ID):
        """
        Get the term IDs of the terms of which `ID` is a direct dependent.

        :param headID:  term ID to get direct parents of
        :return:        a set of IDs that are direct parents of `ID`
        """
        if ID in self.dep2heads:
            return {ID for ID, _ in self.dep2heads[ID]}

    def get_constituent(self, ID):
        """
        Get the term IDs of the terms dependent on `ID`.

        Always contains `ID`, even if it is an unknown `ID`.

        :param headID:  term ID to get constituent of
        :return:        a set of IDs of terms that are dependents of `ID`
        """
        try:
            return self._head2constituent[ID]
        except KeyError:
            return self._get_constituent(ID, set())

    def _get_constituent(self, ID, parents):
        """
        Calculate the constituent if it is not in the cache

        :param head:      head to find constituent of
        :param parents:   term IDs that are already being calculated, thus
                          should not be recursed into
        :return:          set of term IDs that are in the constituent of `head`
        """
        if ID in parents:
            # Prevent loop
            return set()
        elif ID in self._head2constituent:
            # Use cache
            return self._head2constituent[ID]
        # Base case: direct dependents
        deps = {term_id for term_id, _ in self.head2deps.get(ID, [])}
        # Recursive step: call _get_constituent for every direct dependent
        parents.add(ID)
        deps.update(*(
            self._get_constituent(term_id, parents)
            for term_id in deps
        ))
        # Make sure ID itself is in it
        deps.add(ID)
        # Cache
        self._head2constituent[ID] = deps
        return deps

    @staticmethod
    def create_headdep_dict(nafobj, term_filter):
        """
        Create dictionaries of dependent to heads and head to direct dependent

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :param term_filter:     filter for terms
        :return:        head2deps, dep2heads
        """

        allhead2deps = {}
        dep2headIDs = {}
        for dep in nafobj.get_dependencies():
            headID = dep.get_from()
            toID = dep.get_to()
            allhead2deps.setdefault(headID, []).append(
                (toID, dep.get_function())
            )
            dep2headIDs.setdefault(toID, []).append(headID)

        head2deps = {}
        for headID, deps in allhead2deps.items():
            # I don't have to do something with the deps that are filtered out,
            # because if they are leaves they can be left out and if they
            # aren't leaves they will also appear as headID and handled there.
            deps = {
                (toID, relation)
                for toID, relation in deps
                if term_filter(nafobj, toID)
            }
            if term_filter(nafobj, headID):
                head2deps.setdefault(headID, set()).update(deps)
            else:
                # Delete the head by adding its dependents to the heads of the
                # head.
                for super_headID in dep2headIDs[headID]:
                    if term_filter(nafobj, super_headID):
                        head2deps.setdefault(super_headID, set()).update(deps)

        return head2deps
<commit_msg>Fix docstring<commit_after>import logging

logger = logging.getLogger(None if __name__ == '__main__' else __name__)

class ConstituencyTree:
    """
    Group together information from a constituency tree and expose convenient
    lookups.
    """

    def __init__(self, head2deps):
        """
        Initialise `ConstituencyTree`
        """
        self.head2deps = head2deps

        # Create the reverse too
        self.dep2heads = {}
        for headID, deps in head2deps.items():
            for toID, relation in deps:
                self.dep2heads.setdefault(toID, set()).add((headID, relation))

        logger.debug("head2deps: {}".format(self.head2deps))
        logger.debug("dep2heads: {}".format(self.dep2heads))
        self._head2constituent = {}

    def __repr__(self):
        return self.__class__.__name__ + '({!r})'.format(self.head2deps)

    @classmethod
    def from_naf(cls, nafobj, term_filter=lambda naf, t: True):
        """
        Initialise this ConstituencyTree from a NAF object

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :param term_filter:     filter for terms
        """
        return cls(cls.create_headdep_dict(
            nafobj,
            term_filter
        ))

    def get_direct_dependents(self, ID):
        """
        Get the term IDs of the terms directly dependent on `ID`.

        :param headID:  term ID to get direct dependents of
        :return:        a set of IDs that are direct dependents of `ID`
        """
        if ID in self.head2deps:
            return {ID for ID, _ in self.head2deps[ID]}

    def get_direct_parents(self, ID):
        """
        Get the term IDs of the terms of which `ID` is a direct dependent.

        :param headID:  term ID to get direct parents of
        :return:        a set of IDs that are direct parents of `ID`
        """
        if ID in self.dep2heads:
            return {ID for ID, _ in self.dep2heads[ID]}

    def get_constituent(self, ID):
        """
        Get the term IDs of the terms dependent on `ID`.

        Always contains `ID`, even if it is an unknown `ID`.

        :param headID:  term ID to get constituent of
        :return:        a set of IDs of terms that are dependents of `ID`
        """
        try:
            return self._head2constituent[ID]
        except KeyError:
            return self._get_constituent(ID, set())

    def _get_constituent(self, ID, parents):
        """
        Calculate the constituent if it is not in the cache

        :param head:      head to find constituent of
        :param parents:   term IDs that are already being calculated, thus
                          should not be recursed into
        :return:          set of term IDs that are in the constituent of `head`
        """
        if ID in parents:
            # Prevent loop
            return set()
        elif ID in self._head2constituent:
            # Use cache
            return self._head2constituent[ID]
        # Base case: direct dependents
        deps = {term_id for term_id, _ in self.head2deps.get(ID, [])}
        # Recursive step: call _get_constituent for every direct dependent
        parents.add(ID)
        deps.update(*(
            self._get_constituent(term_id, parents)
            for term_id in deps
        ))
        # Make sure ID itself is in it
        deps.add(ID)
        # Cache
        self._head2constituent[ID] = deps
        return deps

    @staticmethod
    def create_headdep_dict(nafobj, term_filter):
        """
        Create dictionary of head to direct dependents

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :param term_filter:     filter for terms
        :return:                head2deps
        """

        allhead2deps = {}
        dep2headIDs = {}
        for dep in nafobj.get_dependencies():
            headID = dep.get_from()
            toID = dep.get_to()
            allhead2deps.setdefault(headID, []).append(
                (toID, dep.get_function())
            )
            dep2headIDs.setdefault(toID, []).append(headID)

        head2deps = {}
        for headID, deps in allhead2deps.items():
            # I don't have to do something with the deps that are filtered out,
            # because if they are leaves they can be left out and if they
            # aren't leaves they will also appear as headID and handled there.
            deps = {
                (toID, relation)
                for toID, relation in deps
                if term_filter(nafobj, toID)
            }
            if term_filter(nafobj, headID):
                head2deps.setdefault(headID, set()).update(deps)
            else:
                # Delete the head by adding its dependents to the heads of the
                # head.
                for super_headID in dep2headIDs[headID]:
                    if term_filter(nafobj, super_headID):
                        head2deps.setdefault(super_headID, set()).update(deps)

        return head2deps
<|endoftext|>