<commit_before>#!/usr/bin/env python
# Copyright (c) 2013, Eduard Broecker
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that
# the following conditions are met:
#
#    Redistributions of source code must retain the above copyright notice, this list of conditions and the
#    following disclaimer.
#    Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
#    following disclaimer in the documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
# DAMAGE.

#
# this script exports json-files from a canmatrix-object
# json-files are the can-matrix-definitions of the CANard-project
# (https://github.com/ericevenchick/CANard)

from __future__ import absolute_import

import json
import sys
import typing
from builtins import *

import canmatrix

extension = 'json'

def dump(db, f, **options):
    # type: (canmatrix.CanMatrix, typing.BinaryIO, **str) -> None

    exportCanard = options.get('jsonCanard', False)
    motorolaBitFormat = options.get('jsonMotorolaBitFormat', "lsb")
    exportAll = options.get('jsonAll', False)

    if (sys.version_info > (3, 0)):
        mode = 'w'
    else:
        mode = 'wb'

    additionalFrameColums = []  # type: typing.List[str]
    if "additionalFrameAttributes" in options and options["additionalFrameAttributes"]:
        additionalFrameColums = options["additionalFrameAttributes"].split(",")

    exportArray = []  # type: typing.List[typing.Union[str, int, list, dict]]

    if exportCanard:
        for frame in db.frames:
            signals = {}
            for signal in frame.signals:
                signals[
                    signal.get_startbit(
                        bit_numbering=1,
                        start_little=True)] = {
                    "name": signal.name,
                    "bit_length": signal.size,
                    "factor": signal.factor,
                    "offset": signal.offset}
            exportArray.append(
                {"name": frame.name, "id": hex(frame.arbitration_id.id), "signals": signals})

    elif exportAll is False:
        for frame in db.frames:
            symbolic_signals = []
            for signal in frame.signals:
                if not signal.is_little_endian:
                    if motorolaBitFormat == "msb":
                        startBit = signal.get_startbit(bit_numbering=1)
                    elif motorolaBitFormat == "msbreverse":
                        startBit = signal.get_startbit()
                    else:  # motorolaBitFormat == "lsb"
                        startBit = signal.get_startbit(bit_numbering=1, start_little=True)
                else:
                    startBit = signal.get_startbit(bit_numbering=1, start_little=True)

                symbolic_signals.append({
                    "name": signal.name,
                    "start_bit": startBit,
                    "bit_length": signal.size,
                    "factor": str(signal.factor),
                    "offset": str(signal.offset),
                    "is_big_endian": signal.is_little_endian == 0,
                    "is_signed": signal.is_signed,
                    "is_float": signal.is_float
                })
            symbolic_frame = {"name": frame.name,
                              "id": int(frame.arbitration_id.id),
                              "is_extended_frame": frame.arbitration_id.extended,
                              "signals": symbolic_signals}
            frame_attributes = {}
            for frame_info in additionalFrameColums:  # Look for additional Frame Attributes
                if frame.attribute(frame_info) is not None:  # does the attribute exist? None value isn't exported.
                    frame_attributes[frame_info] = frame.attribute(frame_info)
            if frame_attributes:  # only add Attributes if there are any
                symbolic_frame["attributes"] = frame_attributes
            exportArray.append(symbolic_frame)
    else:  # exportAll
        for frame in db.frames:
            frameattribs = {}
            for attribute in db.frame_defines:
                frameattribs[attribute] = frame.attribute(attribute, db=db)
            symbolic_signals = []
            for signal in frame.signals:
                attribs = {}
                for attribute in db.signal_defines:
                    attribs[attribute] = signal.attribute(attribute, db=db)

                values = {}
                for key in signal.values:
                    values[key] = signal.values[key]
                if not signal.is_little_endian:
                    if motorolaBitFormat == "msb":
                        startBit = signal.get_startbit(bit_numbering=1)
                    elif motorolaBitFormat == "msbreverse":
                        startBit = signal.get_startbit()
                    else:  # motorolaBitFormat == "lsb"
                        startBit = signal.get_startbit(bit_numbering=1, start_little=True)
                else:  # motorolaBitFormat == "lsb"
                    startBit = signal.get_startbit(bit_numbering=1, start_little=True)

                signalDict = {
                    "name": signal.name,
                    "start_bit": startBit,
                    "bit_length": signal.size,
                    "factor": str(signal.factor),
                    "offset": str(signal.offset),
                    "min": str(signal.min),
                    "max": str(signal.max),
                    "is_big_endian": signal.is_little_endian == 0,
                    "is_signed": signal.is_signed,
                    "is_float": signal.is_float,
                    "comment": signal.comment,
                    "attributes": attribs,
                    "values": values
                }
                if signal.multiplex is not None:
                    signalDict["multiplex"] = signal.multiplex
                if signal.unit:
                    signalDict["unit"] = signal.unit
                symbolic_signals.append(signalDict)

            exportArray.append(
                {"name": frame.name,
                 "id": int(frame.arbitration_id.id),
                 "is_extended_frame": frame.arbitration_id.extended,
                 "signals": symbolic_signals,
                 "attributes": frameattribs,
                 "comment": frame.comment,
                 "length": frame.size})
    if sys.version_info > (3, 0):
        import io
        temp = io.TextIOWrapper(f, encoding='UTF-8')
    else:
        temp = f

    try:
        json.dump({"messages": exportArray}, temp, sort_keys=True,
                  indent=4, separators=(',', ': '))
    finally:
        if sys.version_info > (3, 0):
            # When TextIOWrapper is garbage collected, it closes the raw stream
            # unless the raw stream is detached first
            temp.detach()

def load(f, **options):
    # type: (typing.BinaryIO, **str) -> canmatrix.CanMatrix

    db = canmatrix.CanMatrix()

    if (sys.version_info > (3, 0)):
        import io
        jsonData = json.load(io.TextIOWrapper(f, encoding='UTF-8'))
    else:
        jsonData = json.load(f)

    if "messages" in jsonData:
        for frame in jsonData["messages"]:
            # newframe = Frame(frame["id"],frame["name"],8,None)
            newframe = canmatrix.Frame(frame["name"],
                                   arbitration_id=frame["id"],
                             size=8)
            if "length" in frame:
                newframe.size = frame["length"]

            if "is_extended_frame" in frame and frame["is_extended_frame"]:
                newframe.arbitration_id.extended = 1
            else:
                newframe.arbitration_id.extended = 0

            for signal in frame["signals"]:
                is_little_endian = not signal.get("is_big_endian", False)

                if signal.get("is_float", False):
                    is_float = True
                else:
                    is_float = False

                if signal.get("is_signed", False):
                    is_signed = True
                else:
                    is_signed = False
                newsignal = canmatrix.Signal(signal["name"],
                                   start_bit=signal["start_bit"],
                                   size=signal["bit_length"],
                                   is_little_endian=is_little_endian,
                                   is_signed=is_signed,
                                   is_float = is_float,
                                   factor=signal["factor"],
                                   offset=signal["offset"])

                if signal.get("min") is not None:
                    newsignal.min = newsignal.float_factory(signal["min"])

                if signal.get("max", False):
                    newsignal.max = newsignal.float_factory(signal["max"])

                if signal.get("unit", False):
                    newsignal.unit = signal["unit"]

                if signal.get("multiplex", False):
                    newsignal.unit = signal["multiplex"]

                if signal.get("values", False):
                    for key in signal["values"]:
                        newsignal.add_values(key, signal["values"][key])
                if newsignal.is_little_endian is False:
                    newsignal.set_startbit(
                        newsignal.start_bit, bitNumbering=1, startLittle=True)
                newframe.add_signal(newsignal)
            db.add_frame(newframe)
    f.close()
    return db
<commit_msg>annotation and cleanup, pep8 (#341)<commit_after>#!/usr/bin/env python
# Copyright (c) 2013, Eduard Broecker
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that
# the following conditions are met:
#
#    Redistributions of source code must retain the above copyright notice, this list of conditions and the
#    following disclaimer.
#    Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
#    following disclaimer in the documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
# DAMAGE.

#
# this script exports json-files from a canmatrix-object
# json-files are the can-matrix-definitions of the CANard-project
# (https://github.com/ericevenchick/CANard)

from __future__ import absolute_import

import json
import sys
import typing
from builtins import *

import canmatrix

extension = 'json'

def dump(db, f, **options):
    # type: (canmatrix.CanMatrix, typing.BinaryIO, **str) -> None

    export_canard = options.get('jsonCanard', False)
    motorola_bit_format = options.get('jsonMotorolaBitFormat', "lsb")
    export_all = options.get('jsonAll', False)
    additional_frame_columns = [x for x in options.get("additionalFrameAttributes", "").split(",") if x]

    export_array = []  # type: typing.List[typing.Union[str, float, list, dict]]

    if export_canard:
        for frame in db.frames:
            signals = {}
            for signal in frame.signals:
                signals[
                    signal.get_startbit(
                        bit_numbering=1,
                        start_little=True)] = {
                    "name": signal.name,
                    "bit_length": signal.size,
                    "factor": signal.factor,
                    "offset": signal.offset}
            export_array.append(
                {"name": frame.name, "id": hex(frame.arbitration_id.id), "signals": signals})

    elif export_all is False:
        for frame in db.frames:
            symbolic_signals = []
            for signal in frame.signals:
                if not signal.is_little_endian:
                    if motorola_bit_format == "msb":
                        start_bit = signal.get_startbit(bit_numbering=1)
                    elif motorola_bit_format == "msbreverse":
                        start_bit = signal.get_startbit()
                    else:  # motorola_bit_format == "lsb"
                        start_bit = signal.get_startbit(bit_numbering=1, start_little=True)
                else:
                    start_bit = signal.get_startbit(bit_numbering=1, start_little=True)

                symbolic_signals.append({
                    "name": signal.name,
                    "start_bit": start_bit,
                    "bit_length": signal.size,
                    "factor": str(signal.factor),
                    "offset": str(signal.offset),
                    "is_big_endian": signal.is_little_endian == 0,
                    "is_signed": signal.is_signed,
                    "is_float": signal.is_float
                })
            symbolic_frame = {"name": frame.name,
                              "id": int(frame.arbitration_id.id),
                              "is_extended_frame": frame.arbitration_id.extended,
                              "signals": symbolic_signals}
            frame_attributes = {
                attr: frame.attribute(attr)
                for attr in additional_frame_columns
                if frame.attribute(attr) is not None  # don't export None parameters
            }
            if frame_attributes:  # only add attributes if there are any
                symbolic_frame["attributes"] = frame_attributes
            export_array.append(symbolic_frame)
    else:  # export_all
        for frame in db.frames:
            frame_attributes = {attribute: frame.attribute(attribute, db=db) for attribute in db.frame_defines}
            symbolic_signals = []
            for signal in frame.signals:
                attributes = {attribute: signal.attribute(attribute, db=db) for attribute in db.signal_defines}
                values = {key: signal.values[key] for key in signal.values}
                if not signal.is_little_endian:
                    if motorola_bit_format == "msb":
                        start_bit = signal.get_startbit(bit_numbering=1)
                    elif motorola_bit_format == "msbreverse":
                        start_bit = signal.get_startbit()
                    else:  # motorola_bit_format == "lsb"
                        start_bit = signal.get_startbit(bit_numbering=1, start_little=True)
                else:  # motorola_bit_format == "lsb"
                    start_bit = signal.get_startbit(bit_numbering=1, start_little=True)

                symbolic_signal = {
                    "name": signal.name,
                    "start_bit": start_bit,
                    "bit_length": signal.size,
                    "factor": str(signal.factor),
                    "offset": str(signal.offset),
                    "min": str(signal.min),
                    "max": str(signal.max),
                    "is_big_endian": signal.is_little_endian == 0,
                    "is_signed": signal.is_signed,
                    "is_float": signal.is_float,
                    "comment": signal.comment,
                    "attributes": attributes,
                    "values": values
                }
                if signal.multiplex is not None:
                    symbolic_signal["multiplex"] = signal.multiplex
                if signal.unit:
                    symbolic_signal["unit"] = signal.unit
                symbolic_signals.append(symbolic_signal)

            export_array.append(
                {"name": frame.name,
                 "id": int(frame.arbitration_id.id),
                 "is_extended_frame": frame.arbitration_id.extended,
                 "signals": symbolic_signals,
                 "attributes": frame_attributes,
                 "comment": frame.comment,
                 "length": frame.size})
    if sys.version_info > (3, 0):
        import io
        temp = io.TextIOWrapper(f, encoding='UTF-8')
    else:
        temp = f

    try:
        json.dump({"messages": export_array}, temp, sort_keys=True,
                  indent=4, separators=(',', ': '))
    finally:
        if sys.version_info > (3, 0):
            # When TextIOWrapper is garbage collected, it closes the raw stream
            # unless the raw stream is detached first
            temp.detach()

def load(f, **_options):
    # type: (typing.BinaryIO, **str) -> canmatrix.CanMatrix

    db = canmatrix.CanMatrix()

    if sys.version_info > (3, 0):
        import io
        json_data = json.load(io.TextIOWrapper(f, encoding='UTF-8'))
    else:
        json_data = json.load(f)

    if "messages" in json_data:
        for frame in json_data["messages"]:
            # new_frame = Frame(frame["id"],frame["name"],8,None)
            new_frame = canmatrix.Frame(frame["name"], arbitration_id=frame["id"], size=8)
            if "length" in frame:
                new_frame.size = frame["length"]

            new_frame.arbitration_id.extended = frame.get("is_extended_frame", False)

            for signal in frame["signals"]:
                is_little_endian = not signal.get("is_big_endian", False)
                is_float = signal.get("is_float", False)
                is_signed = signal.get("is_signed", False)

                new_signal = canmatrix.Signal(
                    signal["name"],
                    start_bit=signal["start_bit"],
                    size=signal["bit_length"],
                    is_little_endian=is_little_endian,
                    is_signed=is_signed,
                    is_float=is_float,
                    factor=signal["factor"],
                    offset=signal["offset"]
                )

                if signal.get("min") is not None:
                    new_signal.min = new_signal.float_factory(signal["min"])

                if signal.get("max", False):
                    new_signal.max = new_signal.float_factory(signal["max"])

                if signal.get("unit", False):
                    new_signal.unit = signal["unit"]

                if signal.get("multiplex", False):
                    new_signal.unit = signal["multiplex"]

                if signal.get("values", False):
                    for key in signal["values"]:
                        new_signal.add_values(key, signal["values"][key])
                if new_signal.is_little_endian is False:
                    new_signal.set_startbit(
                        new_signal.start_bit, bitNumbering=1, startLittle=True)
                new_frame.add_signal(new_signal)
            db.add_frame(new_frame)
    f.close()
    return db
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals
import json
from copy import copy

from django.conf import settings
from kafka import KafkaConsumer
from kafka.common import TopicPartition

from corehq.apps.change_feed.data_sources import get_document_store
from corehq.apps.change_feed.exceptions import UnknownDocumentStore
from corehq.apps.change_feed.topics import get_multi_topic_offset, validate_offsets
from memoized import memoized
from dimagi.utils.logging import notify_error
from pillowtop.checkpoints.manager import PillowCheckpointEventHandler
from pillowtop.models import kafka_seq_to_str
from pillowtop.feed.interface import ChangeFeed, Change, ChangeMeta
from six.moves import range

MIN_TIMEOUT = 500

class KafkaChangeFeed(ChangeFeed):
    """
    Kafka-based implementation of a ChangeFeed
    """
    sequence_format = 'json'

    def __init__(self, topics, client_id, strict=False, num_processes=1, process_num=0):
        """
        Create a change feed listener for a list of kafka topics, a client ID, and partition.

        See http://kafka.apache.org/documentation.html#introduction for a description of what these are.
        """
        self._topics = topics
        self._client_id = client_id
        self._processed_topic_offsets = {}
        self.strict = strict
        self.num_processes = num_processes
        self.process_num = process_num

    def __unicode__(self):
        return 'KafkaChangeFeed: topics: {}, client: {}'.format(self._topics, self._client_id)

    @property
    def topics(self):
        return self._topics

    @property
    @memoized
    def topic_and_partitions(self):
        return list(self._get_partitioned_offsets(get_multi_topic_offset(self.topics)))

    def _get_single_topic_or_fail(self):
        if len(self._topics) != 1:
            raise ValueError("This function requires a single topic but found {}!".format(self._topics))
        return self._topics[0]

    def iter_changes(self, since, forever):
        """
        Since must be a dictionary of topic partition offsets.
        """
        since = self._filter_offsets(since)
        # a special value of since=None will start from the end of the change stream
        if since is not None and (not isinstance(since, dict) or not since):
            raise ValueError("'since' must be None or a topic offset dictionary")

        timeout = float('inf') if forever else MIN_TIMEOUT

        start_from_latest = since is None

        reset = 'largest' if start_from_latest else 'smallest'
        consumer = self._get_consumer(timeout, auto_offset_reset=reset)
        consumer.assign([TopicPartition(a, b) for a, b in self.topic_and_partitions])

        if not start_from_latest:
            if self.strict:
                validate_offsets(since)

            checkpoint_topics = {tp[0] for tp in since}
            extra_topics = checkpoint_topics - set(self._topics)
            if extra_topics:
                raise ValueError("'since' contains extra topics: {}".format(list(extra_topics)))

            self._processed_topic_offsets = copy(since)

            # this is how you tell the consumer to start from a certain point in the sequence
            for topic_partition, offset in since.items():
                consumer.seek(TopicPartition(topic_partition[0], topic_partition[1]), int(offset))

        try:
            for message in consumer:
                self._processed_topic_offsets[(message.topic, message.partition)] = message.offset
                yield change_from_kafka_message(message)
        except StopIteration:
            assert not forever, 'Kafka pillow should not timeout when waiting forever!'
            # no need to do anything since this is just telling us we've reached the end of the feed

    def get_current_checkpoint_offsets(self):
        # the way kafka works, the checkpoint should increment by 1 because
        # querying the feed is inclusive of the value passed in.
        latest_offsets = self.get_latest_offsets()
        ret = {}
        for topic_partition, sequence in self.get_processed_offsets().items():
            if sequence == latest_offsets[topic_partition]:
                # this topic and partition is totally up to date and if we add 1
                # then kafka will give us an offset out of range error.
                # not adding 1 to the partition means that we may process this
                # change again later, but that should be OK
                sequence = latest_offsets[topic_partition]
            else:
                sequence += 1
            ret[topic_partition] = sequence
        return self._filter_offsets(ret)

    def get_processed_offsets(self):
        return copy(self._processed_topic_offsets)

    def get_latest_offsets(self):
        return self._filter_offsets(get_multi_topic_offset(self.topics))

    def get_latest_offsets_json(self):
        return json.loads(kafka_seq_to_str(self.get_latest_offsets()))

    def get_latest_offsets_as_checkpoint_value(self):
        return self.get_latest_offsets()

    def _get_consumer(self, timeout, auto_offset_reset='smallest'):
        config = {
            'client_id': self._client_id,
            'bootstrap_servers': settings.KAFKA_BROKERS,
            'consumer_timeout_ms': timeout,
            'auto_offset_reset': auto_offset_reset,
            'enable_auto_commit': False,
            'api_version': settings.KAFKA_API_VERSION,
        }
        return KafkaConsumer(
            **config
        )

    def _filter_offsets(self, offsets):
        if offsets is None:
            return offsets

        return {
            tp: offsets[tp]
            for tp in self.topic_and_partitions
            if tp in offsets
        }

    def _get_partitioned_offsets(self, offsets):
        topic_partitions = sorted(list(offsets))

        partitioned_topic_partitions = [
            topic_partitions[num::self.num_processes]
            for num in range(self.num_processes)
        ][self.process_num]

        return {
            tp: offset
            for tp, offset in offsets.items()
            if tp in partitioned_topic_partitions
        }

class KafkaCheckpointEventHandler(PillowCheckpointEventHandler):
    """
    Event handler that supports checkpoints when subscribing to multiple topics.
    """

    def __init__(self, checkpoint, checkpoint_frequency, change_feed, checkpoint_callback=None):
        super(KafkaCheckpointEventHandler, self).__init__(checkpoint, checkpoint_frequency, checkpoint_callback)
        assert isinstance(change_feed, KafkaChangeFeed)
        self.change_feed = change_feed

    def fire_change_processed(self, change, context):
        if self.should_update_checkpoint(context):
            updated_to = self.change_feed.get_current_checkpoint_offsets()
            self.update_checkpoint(updated_to)
            return True

        return False

def change_from_kafka_message(message):
    change_meta = change_meta_from_kafka_message(message.value)
    try:
        document_store = get_document_store(
            data_source_type=change_meta.data_source_type,
            data_source_name=change_meta.data_source_name,
            domain=change_meta.domain
        )
    except UnknownDocumentStore:
        document_store = None
        notify_error("Unknown document store: {}".format(change_meta.data_source_type))
    return Change(
        id=change_meta.document_id,
        sequence_id=message.offset,
        document=None,
        deleted=change_meta.is_deletion,
        metadata=change_meta,
        document_store=document_store,
        topic=message.topic,
    )

def change_meta_from_kafka_message(message):
    return ChangeMeta.wrap(json.loads(message))
<commit_msg>Get latest offsets from the consumer instead of utility<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals
import json
from copy import copy

from django.conf import settings
from kafka import KafkaConsumer
from kafka.common import TopicPartition

from corehq.apps.change_feed.data_sources import get_document_store
from corehq.apps.change_feed.exceptions import UnknownDocumentStore
from corehq.apps.change_feed.topics import validate_offsets
from dimagi.utils.logging import notify_error
from pillowtop.checkpoints.manager import PillowCheckpointEventHandler
from pillowtop.models import kafka_seq_to_str
from pillowtop.feed.interface import ChangeFeed, Change, ChangeMeta
from six.moves import range

MIN_TIMEOUT = 500

class KafkaChangeFeed(ChangeFeed):
    """
    Kafka-based implementation of a ChangeFeed
    """
    sequence_format = 'json'

    def __init__(self, topics, client_id, strict=False, num_processes=1, process_num=0):
        """
        Create a change feed listener for a list of kafka topics, a client ID, and partition.

        See http://kafka.apache.org/documentation.html#introduction for a description of what these are.
        """
        self._topics = topics
        self._client_id = client_id
        self._processed_topic_offsets = {}
        self.strict = strict
        self.num_processes = num_processes
        self.process_num = process_num
        self._consumer = None

    def __unicode__(self):
        return 'KafkaChangeFeed: topics: {}, client: {}'.format(self._topics, self._client_id)

    @property
    def topics(self):
        return self._topics

    def _get_single_topic_or_fail(self):
        if len(self._topics) != 1:
            raise ValueError("This function requires a single topic but found {}!".format(self._topics))
        return self._topics[0]

    def iter_changes(self, since, forever):
        """
        Since must be a dictionary of topic partition offsets.
        """
        timeout = float('inf') if forever else MIN_TIMEOUT
        start_from_latest = since is None
        reset = 'largest' if start_from_latest else 'smallest'
        self._init_consumer(timeout, auto_offset_reset=reset)

        since = self._filter_offsets(since)
        # a special value of since=None will start from the end of the change stream
        if since is not None and (not isinstance(since, dict) or not since):
            raise ValueError("'since' must be None or a topic offset dictionary")

        if not start_from_latest:
            if self.strict:
                validate_offsets(since)

            checkpoint_topics = {tp[0] for tp in since}
            extra_topics = checkpoint_topics - set(self._topics)
            if extra_topics:
                raise ValueError("'since' contains extra topics: {}".format(list(extra_topics)))

            self._processed_topic_offsets = copy(since)

            # Tell the consumer to start from offsets that were passed in
            for topic_partition, offset in since.items():
                self.consumer.seek(TopicPartition(topic_partition[0], topic_partition[1]), int(offset))

        try:
            for message in self.consumer:
                self._processed_topic_offsets[(message.topic, message.partition)] = message.offset
                yield change_from_kafka_message(message)
        except StopIteration:
            assert not forever, 'Kafka pillow should not timeout when waiting forever!'
            # no need to do anything since this is just telling us we've reached the end of the feed

    def get_current_checkpoint_offsets(self):
        # the way kafka works, the checkpoint should increment by 1 because
        # querying the feed is inclusive of the value passed in.
        latest_offsets = self.get_latest_offsets()
        ret = {}
        for topic_partition, sequence in self.get_processed_offsets().items():
            if sequence == latest_offsets[topic_partition]:
                # this topic and partition is totally up to date and if we add 1
                # then kafka will give us an offset out of range error.
                # not adding 1 to the partition means that we may process this
                # change again later, but that should be OK
                sequence = latest_offsets[topic_partition]
            else:
                sequence += 1
            ret[topic_partition] = sequence
        return self._filter_offsets(ret)

    def get_processed_offsets(self):
        return copy(self._processed_topic_offsets)

    def get_latest_offsets(self):
        return self.consumer.end_offsets(self.consumer.assignment())

    def get_latest_offsets_json(self):
        return json.loads(kafka_seq_to_str(self.get_latest_offsets()))

    def get_latest_offsets_as_checkpoint_value(self):
        return self.get_latest_offsets()

    @property
    def consumer(self):
        if self._consumer is None:
            return self._init_consumer()
        return self._consumer

    def _init_consumer(self, timeout=MIN_TIMEOUT, auto_offset_reset='smallest'):
        """Allow re-initing the consumer if necessary
        """
        config = {
            'client_id': self._client_id,
            'bootstrap_servers': settings.KAFKA_BROKERS,
            'consumer_timeout_ms': timeout,
            'auto_offset_reset': auto_offset_reset,
            'enable_auto_commit': False,
            'api_version': settings.KAFKA_API_VERSION,
        }
        self._consumer = KafkaConsumer(**config)

        topic_partitions = []
        for topic in self.topics:
            for partition in self._consumer.partitions_for_topic(topic):
                topic_partitions.append(TopicPartition(topic, partition))

        self._consumer.assign(self._filter_partitions(topic_partitions))
        return self._consumer

    def _filter_offsets(self, offsets):
        if offsets is None:
            return offsets

        return {
            tp: offsets[tp]
            for tp in self.consumer.assignment()
            if tp in offsets
        }

    def _filter_partitions(self, topic_partitions):
        topic_partitions.sort()

        return [
            topic_partitions[num::self.num_processes]
            for num in range(self.num_processes)
        ][self.process_num]

class KafkaCheckpointEventHandler(PillowCheckpointEventHandler):
    """
    Event handler that supports checkpoints when subscribing to multiple topics.
    """

    def __init__(self, checkpoint, checkpoint_frequency, change_feed, checkpoint_callback=None):
        super(KafkaCheckpointEventHandler, self).__init__(checkpoint, checkpoint_frequency, checkpoint_callback)
        assert isinstance(change_feed, KafkaChangeFeed)
        self.change_feed = change_feed

    def fire_change_processed(self, change, context):
        if self.should_update_checkpoint(context):
            updated_to = self.change_feed.get_current_checkpoint_offsets()
            self.update_checkpoint(updated_to)
            return True

        return False

def change_from_kafka_message(message):
    change_meta = change_meta_from_kafka_message(message.value)
    try:
        document_store = get_document_store(
            data_source_type=change_meta.data_source_type,
            data_source_name=change_meta.data_source_name,
            domain=change_meta.domain
        )
    except UnknownDocumentStore:
        document_store = None
        notify_error("Unknown document store: {}".format(change_meta.data_source_type))
    return Change(
        id=change_meta.document_id,
        sequence_id=message.offset,
        document=None,
        deleted=change_meta.is_deletion,
        metadata=change_meta,
        document_store=document_store,
        topic=message.topic,
    )

def change_meta_from_kafka_message(message):
    return ChangeMeta.wrap(json.loads(message))
<|endoftext|>


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

import sys, re, json

filepath = sys.argv[1]
with open(filepath, 'r') as xml_file:
    xml = xml_file.read()

drawMap = False
if len(sys.argv) > 2:
    drawMap = True

minl = 100
mint = 130
maxt = 730
l1 = 200
parls_type = "senateurs"
parls_first = True
if "senateurs_collaborateurs" in filepath:
    mint = 150
elif "collaborateurs_senateurs" in filepath:
    parls_first = False
elif "deputes_collaborateurs" in filepath:
    parls_type = "deputes"
    minl = 50
    mint = 85
    maxt = 800
    l1 = 75
    l2 = 150
    l3 = 250
    l4 = 275
    l5 = 350
parl_type = parls_type.rstrip("s")

with open("data/%s.json" % parls_type, 'r') as f:
    parls = [p[parl_type] for p in json.load(f)[parls_type]]

Mme = r"^M[.mle]+\s+"
clean_Mme = re.compile(Mme)

re_clean_bal = re.compile(r'<[^>]+>')
re_clean_dash = re.compile(r'\s*-\s*')
re_clean_spaces = re.compile(r'\s+')
clean = lambda x: re_clean_spaces.sub(' ', re_clean_dash.sub('-', re_clean_bal.sub('', x))).strip()

regexps = [(re.compile(r), s) for r, s in [
    (u'[àÀâÂ]', 'a'),
    (u'[éÉèÈêÊëË]', 'e'),
    (u'[îÎïÏ]', 'i'),
    (u'[ôÔöÔ]', 'o'),
    (u'[ùÙûÛüÜ]', 'u'),
    (u'[çÇ]', 'c'),
]]
def clean_accents(t):
    if not isinstance(t, unicode):
        t = t.decode('utf-8')
    for r, s in regexps:
        t = r.sub(s, t)
    return t
checker = lambda x: clean(clean_accents(x)).lower().strip()

reorder = lambda p: checker("%s %s" % (p['nom'].replace("%s " % p['prenom'], ""), p['prenom']))

maj = ur"A-ZÀÂÉÈÊËÎÏÔÖÙÛÜÇ"
particule = r"d(?:(?:e l)?'|[iu] |e(?:[ls]| la)? )"
re_name = re.compile(ur"(%s)((?:(?:%s)?[%s'\-]+\s+)+)([%s][a-zàâéèêëîïôöùûüç].*)$" % (Mme, particule, maj, maj))

def split_name(name):
    match = re_name.search(name.decode('utf-8'))
    if not match:
        sys.stderr.write("WARNING: could not split name %s\n" % name)
        return name, "", ""
    sexe = "H" if "." in match.group(1) else "F"
    nom = match.group(2).strip()
    prenom = match.group(3).strip()
    return nom, prenom, sexe

def split_collab(record):
    record[5], record[6], record[7] = split_name(record[4])
    if record[6].endswith("*"):
        record[6].rstrip(" *")
        record[10] = "Congé sans solde"
    else: record[10] = ""

exceptions = [
    ("deromedi jacqueline", "deromedi jacky"),
    ("yonnet-salvator evelyne", "yonnet evelyne"),
    ("laufoaulu lopeleto", "laufoaulu robert"),
    ("azerot bruno", "azerot bruno nestor"),
    ("de la verpillere charles", "de la verpilliere charles"),
    ("debre bernard andre", "debre bernard"),
    ("destans jean", "destans jean-louis"),
    ("le borgn pierre-yves", "le borgn' pierre-yves"),
    ("vlody jean-jacques", "vlody jean jacques"),
    ("zimmermann marie jo", "zimmermann marie-jo")
]

def find_parl(record, splitted=False):
    nom = checker(clean_Mme.sub('', record[0]))
    for bad, gd in exceptions:
        nom = nom.replace(bad, gd)
    if not splitted:
        record[1], record[2], record[3] = split_name(record[0])
    for parl in parls:
        if reorder(parl) == nom:
            record[0] = parl["nom"]
            #record[1] = parl["nom_de_famille"]
            #record[2] = parl["prenom"]
            #record[3] = parl["sexe"]
            record[8] = parl["url_nos%s_api" % parls_type].replace('json', 'xml')
            record[9] = parl.get("url_institution", parl.get("url_an", ""))
            return
    sys.stderr.write("Could not find %s -> %s\n" % (record[0], nom.encode('utf-8')))

# Reorder xml lines
xml_ordered = ""
page_lines = []
re_ordline = re.compile(r'<text top="(\d+)" left="(\d+)"[^>]', re.I)
ordline = lambda l: [int(v) for v in re_ordline.search(l).groups()]
for line in (xml).split("\n"):
    if line.startswith('<text'):
        page_lines.append(line)
    elif line.startswith('</page'):
        xml_ordered += "\n".join(sorted(page_lines, key=ordline))
        page_lines = []

topvals = {}
leftvals = {}
maxtop = 0
maxleft = 0
results = []
headers = ['parlementaire', 'nom_parlementaire', 'prénom_parlementaire', 'sexe_parlementaire', 'collaborateur', 'nom_collaborateur', 'prénom_collaborateur', 'sexe_collaborateur', 'url_api_RC', 'url_institution', 'information complémentaire']
record = ["", "", "", "", "", "", "", "", "", "", ""]
re_line = re.compile(r'<page number|text top="(\d+)" left="(\d+)"[^>]*font="(\d+)">(.*)</text>', re.I)
re_tosplit = re.compile(r'^(.*) ((?:M.|Mme) .*)$')
re_collabtosplit = re.compile(r'^\s*(M\.|Mme) (.+)$')
sexize = lambda val: "H" if val == "M." else "F"
for line in (xml_ordered).split("\n"):
    #print >> sys.stderr, "DEBUG %s" % line
    attrs = re_line.search(line)
    if not attrs or not attrs.groups():
        raise Exception("WARNING : line detected with good font but wrong format %s" % line)
    font = int(attrs.group(3))
    top = int(attrs.group(1))
    if top > maxtop:
        maxtop = top
    if not font in topvals:
        topvals[font] = []
    topvals[font].append(top)
    left = int(attrs.group(2))
    if left > maxleft:
        maxleft = left
    if not font in leftvals:
        leftvals[font] = []
    leftvals[font].append(left)
    if drawMap:
        continue
    #print "DEBUG %s %s %s %s" % (font, left, top, text)
    if top < mint or top > maxt:
        continue
    if left < minl:
        continue
    text = attrs.group(4).replace("&", "&")
    # Handle députés
    if parls_type == "deputes":
        val = clean(text)
        if left < l1:
            record[3] = val
        elif left < l2:
            record[2] = val
        elif left < l3:
            record[1] = val
            record[0] = record[3] + " " + record[1] + " " + record[2]
            find_parl(record, splitted=True)
            record[3] = sexize(record[3])
        elif left < l4:
            splitted = re_collabtosplit.search(val)
            if splitted:
                record[7] = splitted.group(1)
                record[6] = splitted.group(2)
            else:
                record[7] = val
        elif left < l5:
            record[6] = val
        else:
            record[5] = val
            record[4] = record[7] + " " + record[5] + " " + record[6]
            record[7] = sexize(record[7])
            results.append(list(record))
        continue
    # Handle sénateurs
    if left < l1:
        val = clean(text)
        idx = 0 if parls_first else 4
        tosplit = re_tosplit.search(val)
        if tosplit:
            sys.stderr.write("WARNING: splitting %s\n" % val)
            idx2 = 4 if parls_first else 0
            record[idx], record[idx2] = tosplit.groups()
            find_parl(record)
            split_collab(record)
            results.append(list(record))
            continue
        record[idx] = val
        if parls_first:
            find_parl(record)
        else:
            split_collab(record)
    else:
        idx = 4 if parls_first else 0
        record[idx] = clean(text)
        if not parls_first:
            find_parl(record)
        else:
            split_collab(record)
        results.append(list(record))

if not drawMap:
    print ",".join(['"%s"' % h for h in headers])
    for i in sorted(results, key=lambda x: ("%s %s - %s %s" % (x[1], x[2], x[5], x[6])).lower()):
        for j in range(len(i)):
            i[j] = clean(i[j])
            try:    i[j] = i[j].encode('utf-8')
            except: pass
        print ",".join([str(i[a]) if isinstance(i[a], int) else "\"%s\"" % i[a].replace('"', '""') for a,_ in enumerate(i)])

else:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import numpy as np
    from matplotlib import cm

    fig = plt.figure(figsize=(8.5, 12))
    ax = fig.add_subplot(111)
    ax.grid(True, fillstyle='left')
    nf = len(leftvals)
    for font in leftvals:
        color = cm.jet(1.5*font/nf)
        ax.plot(leftvals[font], topvals[font], 'ro', color=color, marker=".")
        plt.figtext((font+1.)/(nf+1), 0.95, "font %d" % font, color=color)
    plt.xticks(np.arange(0, maxleft + 50, 50))
    plt.yticks(np.arange(0, maxtop + 50, 50))
    plt.xlim(0, maxleft + 50)
    plt.ylim(0, maxtop + 50)
    plt.gca().invert_yaxis()
    mappath = filepath.replace(".xml", ".png").replace("pdfs/", "pdfmaps/")
    fig.savefig(mappath)
    fig.clf()
    plt.close(fig)

<commit_msg>fix parsing AN since AN added CNIL warning<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, re, json

filepath = sys.argv[1]
with open(filepath, 'r') as xml_file:
    xml = xml_file.read()

drawMap = False
if len(sys.argv) > 2:
    drawMap = True

minl = 100
mint = 130
maxt = 730
l1 = 200
parls_type = "senateurs"
parls_first = True
if "senateurs_collaborateurs" in filepath:
    mint = 150
elif "collaborateurs_senateurs" in filepath:
    parls_first = False
elif "deputes_collaborateurs" in filepath:
    parls_type = "deputes"
    minl = 50
    mint = 220
    maxt = 800
    l1 = 75
    l2 = 150
    l3 = 250
    l4 = 275
    l5 = 350
parl_type = parls_type.rstrip("s")

with open("data/%s.json" % parls_type, 'r') as f:
    parls = [p[parl_type] for p in json.load(f)[parls_type]]

Mme = r"^M[.mle]+\s+"
clean_Mme = re.compile(Mme)

re_clean_bal = re.compile(r'<[^>]+>')
re_clean_dash = re.compile(r'\s*-\s*')
re_clean_spaces = re.compile(r'\s+')
clean = lambda x: re_clean_spaces.sub(' ', re_clean_dash.sub('-', re_clean_bal.sub('', x))).strip()

regexps = [(re.compile(r), s) for r, s in [
    (u'[àÀâÂ]', 'a'),
    (u'[éÉèÈêÊëË]', 'e'),
    (u'[îÎïÏ]', 'i'),
    (u'[ôÔöÔ]', 'o'),
    (u'[ùÙûÛüÜ]', 'u'),
    (u'[çÇ]', 'c'),
]]
def clean_accents(t):
    if not isinstance(t, unicode):
        t = t.decode('utf-8')
    for r, s in regexps:
        t = r.sub(s, t)
    return t
checker = lambda x: clean(clean_accents(x)).lower().strip()

reorder = lambda p: checker("%s %s" % (p['nom'].replace("%s " % p['prenom'], ""), p['prenom']))

maj = ur"A-ZÀÂÉÈÊËÎÏÔÖÙÛÜÇ"
particule = r"d(?:(?:e l)?'|[iu] |e(?:[ls]| la)? )"
re_name = re.compile(ur"(%s)((?:(?:%s)?[%s'\-]+\s+)+)([%s][a-zàâéèêëîïôöùûüç].*)$" % (Mme, particule, maj, maj))

def split_name(name):
    match = re_name.search(name.decode('utf-8'))
    if not match:
        sys.stderr.write("WARNING: could not split name %s\n" % name)
        return name, "", ""
    sexe = "H" if "." in match.group(1) else "F"
    nom = match.group(2).strip()
    prenom = match.group(3).strip()
    return nom, prenom, sexe

def split_collab(record):
    record[5], record[6], record[7] = split_name(record[4])
    if record[6].endswith("*"):
        record[6].rstrip(" *")
        record[10] = "Congé sans solde"
    else: record[10] = ""

exceptions = [
    ("deromedi jacqueline", "deromedi jacky"),
    ("yonnet-salvator evelyne", "yonnet evelyne"),
    ("laufoaulu lopeleto", "laufoaulu robert"),
    ("azerot bruno", "azerot bruno nestor"),
    ("de la verpillere charles", "de la verpilliere charles"),
    ("debre bernard andre", "debre bernard"),
    ("destans jean", "destans jean-louis"),
    ("le borgn pierre-yves", "le borgn' pierre-yves"),
    ("vlody jean-jacques", "vlody jean jacques"),
    ("zimmermann marie jo", "zimmermann marie-jo")
]

def find_parl(record, splitted=False):
    nom = checker(clean_Mme.sub('', record[0]))
    for bad, gd in exceptions:
        nom = nom.replace(bad, gd)
    if not splitted:
        record[1], record[2], record[3] = split_name(record[0])
    for parl in parls:
        if reorder(parl) == nom:
            record[0] = parl["nom"]
            #record[1] = parl["nom_de_famille"]
            #record[2] = parl["prenom"]
            #record[3] = parl["sexe"]
            record[8] = parl["url_nos%s_api" % parls_type].replace('json', 'xml')
            record[9] = parl.get("url_institution", parl.get("url_an", ""))
            return
    sys.stderr.write("Could not find %s -> %s\n" % (record[0], nom.encode('utf-8')))

# Reorder xml lines
xml_ordered = ""
page_lines = []
re_ordline = re.compile(r'<text top="(\d+)" left="(\d+)"[^>]', re.I)
ordline = lambda l: [int(v) for v in re_ordline.search(l).groups()]
for line in (xml).split("\n"):
    if line.startswith('<text'):
        page_lines.append(line)
    elif line.startswith('</page'):
        xml_ordered += "\n".join(sorted(page_lines, key=ordline))
        page_lines = []

topvals = {}
leftvals = {}
maxtop = 0
maxleft = 0
results = []
headers = ['parlementaire', 'nom_parlementaire', 'prénom_parlementaire', 'sexe_parlementaire', 'collaborateur', 'nom_collaborateur', 'prénom_collaborateur', 'sexe_collaborateur', 'url_api_RC', 'url_institution', 'information complémentaire']
record = ["", "", "", "", "", "", "", "", "", "", ""]
re_line = re.compile(r'<page number|text top="(\d+)" left="(\d+)"[^>]*font="(\d+)">(.*)</text>', re.I)
re_tosplit = re.compile(r'^(.*) ((?:M.|Mme) .*)$')
re_collabtosplit = re.compile(r'^\s*(M\.|Mme) (.+)$')
sexize = lambda val: "H" if val == "M." else "F"
for line in (xml_ordered).split("\n"):
    #print >> sys.stderr, "DEBUG %s" % line
    attrs = re_line.search(line)
    if not attrs or not attrs.groups():
        raise Exception("WARNING : line detected with good font but wrong format %s" % line)
    font = int(attrs.group(3))
    top = int(attrs.group(1))
    if top > maxtop:
        maxtop = top
    if not font in topvals:
        topvals[font] = []
    topvals[font].append(top)
    left = int(attrs.group(2))
    if left > maxleft:
        maxleft = left
    if not font in leftvals:
        leftvals[font] = []
    leftvals[font].append(left)
    if drawMap:
        continue
    #print "DEBUG %s %s %s %s" % (font, left, top, text)
    if top < mint or top > maxt:
        continue
    if left < minl:
        continue
    text = attrs.group(4).replace("&", "&")
    # Handle députés
    if parls_type == "deputes":
        val = clean(text)
        if left < l1:
            record[3] = val
        elif left < l2:
            record[2] = val
        elif left < l3:
            record[1] = val
            record[0] = record[3] + " " + record[1] + " " + record[2]
            find_parl(record, splitted=True)
            record[3] = sexize(record[3])
        elif left < l4:
            splitted = re_collabtosplit.search(val)
            if splitted:
                record[7] = splitted.group(1)
                record[6] = splitted.group(2)
            else:
                record[7] = val
        elif left < l5:
            record[6] = val
        else:
            record[5] = val
            record[4] = record[7] + " " + record[5] + " " + record[6]
            record[7] = sexize(record[7])
            results.append(list(record))
        continue
    # Handle sénateurs
    if left < l1:
        val = clean(text)
        idx = 0 if parls_first else 4
        tosplit = re_tosplit.search(val)
        if tosplit:
            sys.stderr.write("WARNING: splitting %s\n" % val)
            idx2 = 4 if parls_first else 0
            record[idx], record[idx2] = tosplit.groups()
            find_parl(record)
            split_collab(record)
            results.append(list(record))
            continue
        record[idx] = val
        if parls_first:
            find_parl(record)
        else:
            split_collab(record)
    else:
        idx = 4 if parls_first else 0
        record[idx] = clean(text)
        if not parls_first:
            find_parl(record)
        else:
            split_collab(record)
        results.append(list(record))

if not drawMap:
    print ",".join(['"%s"' % h for h in headers])
    for i in sorted(results, key=lambda x: ("%s %s - %s %s" % (x[1], x[2], x[5], x[6])).lower()):
        for j in range(len(i)):
            i[j] = clean(i[j])
            try:    i[j] = i[j].encode('utf-8')
            except: pass
        print ",".join([str(i[a]) if isinstance(i[a], int) else "\"%s\"" % i[a].replace('"', '""') for a,_ in enumerate(i)])

else:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import numpy as np
    from matplotlib import cm

    fig = plt.figure(figsize=(8.5, 12))
    ax = fig.add_subplot(111)
    ax.grid(True, fillstyle='left')
    nf = len(leftvals)
    for font in leftvals:
        color = cm.jet(1.5*font/nf)
        ax.plot(leftvals[font], topvals[font], 'ro', color=color, marker=".")
        plt.figtext((font+1.)/(nf+1), 0.95, "font %d" % font, color=color)
    plt.xticks(np.arange(0, maxleft + 50, 50))
    plt.yticks(np.arange(0, maxtop + 50, 50))
    plt.xlim(0, maxleft + 50)
    plt.ylim(0, maxtop + 50)
    plt.gca().invert_yaxis()
    mappath = filepath.replace(".xml", ".png").replace("pdfs/", "pdfmaps/")
    fig.savefig(mappath)
    fig.clf()
    plt.close(fig)

<|endoftext|>


<commit_before>#! /usr/bin/env python
#Copyright (c) 2016, Buti Al Delail
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#* Redistributions of source code must retain the above copyright notice, this
#  list of conditions and the following disclaimer.
#
#* Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.
#
#* Neither the name of kuri_mbzirc_challenge_3 nor the names of its
#  contributors may be used to endorse or promote products derived from
#  this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,

import rospy
import thread
import threading
import time
import mavros
import actionlib
from math import *
from mavros.utils import *
from mavros import setpoint as SP
from tf.transformations import quaternion_from_euler
from kuri_msgs.msg import *

class dropzone_landing:

    def __init__(self):
        self.done = False
        self.done_evt = threading.Event()
        self.isExploring = False
        self.progress = 0.0
        self.x = 0.0
        self.y = 0.0
        self.z = 0.0
        self.currentPoseX = 0
        self.currentPoseY = 0
        self.currentPoseZ = 0
        self.navigating = False

        mavros.set_namespace('/uav_1/mavros')
        # publisher for mavros/setpoint_position/local
        self.pub = SP.get_pub_position_local(queue_size=10)
        # subscriber for mavros/local_position/local
        self.sub = rospy.Subscriber(mavros.get_topic('local_position', 'pose'),
                                    SP.PoseStamped, self.reached)

        self.objects_map = ObjectsMap()

        self.client = actionlib.SimpleActionClient('TrackingAction', TrackingAction)
        #client = self.client
        #client = self.actionServer.client

        print "Waiting for tracking server"
        self.client.wait_for_server()
        self.goal = TrackingGoal()
        self.goal.uav_id = 1
        self.client.send_goal(self.goal)
        print "Waiting for result"
        self.client.wait_for_result()
        print "Result:"
        self.objects =self.client.get_result().tracked_objects.objects
        print self.objects

        try:
            thread.start_new_thread(self.navigate, ())
        except:
            fault("Error: Unable to start thread")

    def navigate(self):
        rate = rospy.Rate(10)   # 10hz

        msg = SP.PoseStamped(
            header=SP.Header(
                frame_id="base_footprint",  # no matter, plugin don't use TF
                stamp=rospy.Time.now()),    # stamp should update
        )

        while not rospy.is_shutdown():
            msg.pose.position.x = self.x
            msg.pose.position.y = self.y
            msg.pose.position.z = self.z

            # For demo purposes we will lock yaw/heading to north.
            yaw_degrees = 0  # North
            yaw = radians(yaw_degrees)
            quaternion = quaternion_from_euler(0, 0, yaw)
            msg.pose.orientation = SP.Quaternion(*quaternion)

            self.pub.publish(msg)
            rate.sleep()

    def setPose(self, x, y, z, delay=0, wait=True):
        self.done = False
        self.x = x
        self.y = y
        self.z = z
        self.navigating = True
        if wait:
            rate = rospy.Rate(5)
            while not self.done and not rospy.is_shutdown():
                rate.sleep()
        time.sleep(delay)

    def takeoff(self, z, delay=0, wait=True):
        diff = z - self.currentPoseZ
        while not abs(diff)<0.2:
            diff = z - self.currentPoseZ
            if diff>0:
                self.setPose(self.currentPoseX,self.currentPoseY,self.currentPoseZ + 0.5 ,2)
            else:
                self.setPose(self.currentPoseX,self.currentPoseY,self.currentPoseZ - 0.5 ,2)

    def land(self, delay=0, wait=True):
        altitude = self.currentPoseZ
        while altitude > 0:
            altitude = self.currentPoseZ
            self.setPose(self.currentPoseX,self.currentPoseY,self.currentPoseZ - 0.5 ,2)

    def reached(self, topic):
        def is_near(msg, x, y, d):
            rospy.logdebug("Position %s: local: %d, target: %d, abs diff: %d",
                           msg, x, y, abs(x - y))
            return abs(x - y) < d
        self.currentPoseX = topic.pose.position.x
        self.currentPoseY = topic.pose.position.y
        self.currentPoseZ = topic.pose.position.z

        if is_near('X', topic.pose.position.x, self.x, 0.2) and \
           is_near('Y', topic.pose.position.y, self.y, 0.2) and \
           is_near('Z', topic.pose.position.z, self.z, 0.5):
            if  self.navigating:
                self.done = True
                self.navigating = False
                self.done_evt.set()

    def explore(self):
        print 'explore started '
        rate = rospy.Rate(30)
        self.newGoal = True
        if self.isExploring == False:
            #Change this later when we have a better exploration
            #self.isExploring = True
            while self.done == False:
                    time.sleep(1)
                    rospy.loginfo("Climb")
                    self.progress += 0.1
                    self.takeoff(7)
                    self.progress += 0.1
                    rospy.loginfo("Moving to Red_Object")
                    self.reached_object = False
                    red_object_id = -1
                    xspeed = 3
                    while self.reached_object == False:
                        self.client.send_goal(self.goal)
                        self.client.wait_for_result()
                        self.objects = self.client.get_result().tracked_objects.objects
                        islost = True
                        for obj in self.objects:
                            if red_object_id == -1 and (obj.color == 'RED' or obj.color == 'BLUE' or obj.color == 'GREEN'): #pick any nearby object
                                red_object_id = obj.object_id
                            if obj.object_id == red_object_id:
                                islost = False
                                print 'Moving to Drop zone', self.currentPoseX-obj.pose.pose.position.x, self.currentPoseY-obj.pose.pose.position.y, obj.pose.pose.position.x, obj.pose.pose.position.y
                                if fabs(obj.pose.pose.position.x) < 0.1 and fabs(obj.pose.pose.position.y) > 0.1:
                                    print 'Moving Y'
                                    self.setPose(self.x, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , False)
                                elif fabs(obj.pose.pose.position.y) < 0.1 and fabs(obj.pose.pose.position.x) > 0.1:
                                    print 'Moving X'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.y, self.z, 0 , False)
                                else:
                                    print 'Moving XY'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , True)
                                if fabs(obj.pose.pose.position.x) < 0.3 and fabs(obj.pose.pose.position.y) < 0.3 and self.z > 0.2:
                                    print 'Moving Z'
                                    land = 0.5
                                    if self.z <= 3:
                                        land = 0.2
                                    if self.z <= 1:
                                        xspeed = 0.5
                                    self.setPose(self.x, self.y, self.z - land, 1, False)
                                    if self.z < 0.3:
                                        self.reached_object = True
                        if islost == True:
                            red_object_id = -1
                        if red_object_id == -1:
                            rospy.loginfo("No object in sight, exploring")
                            self.setPose(self.x, self.y - 5, self.z, 1, True)
                        rate.sleep()
                    rospy.loginfo("Picked Object, climb")
                    self.takeoff(5)
                    rospy.loginfo("Moving to DropZone")
                    self.setPose(1, -21, 5) ##Go near dropzone
                    self.progress += 0.1
                    self.reached_dropzone = False
                    xspeed = 3
                    while self.reached_dropzone == False:
                        self.client.send_goal(self.goal)
                        self.client.wait_for_result()
                        self.objects = self.client.get_result().tracked_objects.objects

                        for obj in self.objects:
                            if obj.color == 'DROP_ZONE':
                                print 'Moving to Drop zone', self.currentPoseX-obj.pose.pose.position.x, self.currentPoseY-obj.pose.pose.position.y, obj.pose.pose.position.x, obj.pose.pose.position.y
                                if fabs(obj.pose.pose.position.x) < 0.1 and fabs(obj.pose.pose.position.y) > 0.1:
                                    print 'Moving Y'
                                    self.setPose(self.x, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , False)
                                elif fabs(obj.pose.pose.position.y) < 0.1 and fabs(obj.pose.pose.position.x) > 0.1:
                                    print 'Moving X'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.y, self.z, 0 , False)
                                else:
                                    print 'Moving XY'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , True)
                                if fabs(obj.pose.pose.position.x) < 0.3 and fabs(obj.pose.pose.position.y) < 0.3 and self.z > 1:
                                    print 'Moving Z'
                                    land = 0.5
                                    if self.z <= 3:
                                        land = 0.2
                                        xspeed = 0.5
                                    self.setPose(self.x, self.y, self.z - land, 5, True)
                                    if self.z < 1.5:
                                        self.reached_dropzone = True
                        rate.sleep()
                    self.progress += 0.1
                    rospy.loginfo("Landed Object, climb")
                    self.takeoff(7)
                    rospy.loginfo("Bye!")

def main(args):
    '''Initializes and cleanup ros node'''
    rospy.init_node('dropzone_landing', anonymous=True)
    d = dropzone_landing()
    d.explore()
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print "Shutting down ROS Image feature detector module"
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main(sys.argv)
<commit_msg>sped up pick and drop<commit_after>#! /usr/bin/env python
#Copyright (c) 2016, Buti Al Delail
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#* Redistributions of source code must retain the above copyright notice, this
#  list of conditions and the following disclaimer.
#
#* Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.
#
#* Neither the name of kuri_mbzirc_challenge_3 nor the names of its
#  contributors may be used to endorse or promote products derived from
#  this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,

import rospy
import thread
import threading
import time
import mavros
import actionlib
from math import *
from mavros.utils import *
from mavros import setpoint as SP
from tf.transformations import quaternion_from_euler
from kuri_msgs.msg import *

class dropzone_landing:

    def __init__(self):
        self.done = False
        self.done_evt = threading.Event()
        self.isExploring = False
        self.progress = 0.0
        self.x = 0.0
        self.y = 0.0
        self.z = 0.0
        self.currentPoseX = 0
        self.currentPoseY = 0
        self.currentPoseZ = 0
        self.navigating = False

        mavros.set_namespace('/uav_1/mavros')
        # publisher for mavros/setpoint_position/local
        self.pub = SP.get_pub_position_local(queue_size=10)
        # subscriber for mavros/local_position/local
        self.sub = rospy.Subscriber(mavros.get_topic('local_position', 'pose'),
                                    SP.PoseStamped, self.reached)

        self.objects_map = ObjectsMap()

        self.client = actionlib.SimpleActionClient('TrackingAction', TrackingAction)
        #client = self.client
        #client = self.actionServer.client

        print "Waiting for tracking server"
        self.client.wait_for_server()
        self.goal = TrackingGoal()
        self.goal.uav_id = 1
        self.client.send_goal(self.goal)
        print "Waiting for result"
        self.client.wait_for_result()
        print "Result:"
        self.objects =self.client.get_result().tracked_objects.objects
        print self.objects

        try:
            thread.start_new_thread(self.navigate, ())
        except:
            fault("Error: Unable to start thread")

    def navigate(self):
        rate = rospy.Rate(40)   # 10hz

        msg = SP.PoseStamped(
            header=SP.Header(
                frame_id="base_footprint",  # no matter, plugin don't use TF
                stamp=rospy.Time.now()),    # stamp should update
        )

        while not rospy.is_shutdown():
            msg.pose.position.x = self.x
            msg.pose.position.y = self.y
            msg.pose.position.z = self.z

            # For demo purposes we will lock yaw/heading to north.
            yaw_degrees = 0  # North
            yaw = radians(yaw_degrees)
            quaternion = quaternion_from_euler(0, 0, yaw)
            msg.pose.orientation = SP.Quaternion(*quaternion)

            self.pub.publish(msg)
            rate.sleep()

    def setPose(self, x, y, z, delay=0, wait=True):
        self.done = False
        self.x = x
        self.y = y
        self.z = z
        self.navigating = True
        if wait:
            rate = rospy.Rate(5)
            while not self.done and not rospy.is_shutdown():
                rate.sleep()
        time.sleep(delay)

    def takeoff(self, z, delay=0, wait=True):
        diff = z - self.currentPoseZ
        while not abs(diff)<0.2:
            diff = z - self.currentPoseZ
            if diff>0:
                self.setPose(self.currentPoseX,self.currentPoseY,self.currentPoseZ + 1, 0, False)
            else:
                self.setPose(self.currentPoseX,self.currentPoseY,self.currentPoseZ - 0.1, 0, False)

    def land(self, delay=0, wait=True):
        altitude = self.currentPoseZ
        while altitude > 0:
            altitude = self.currentPoseZ
            self.setPose(self.currentPoseX,self.currentPoseY,self.currentPoseZ - 0.5 ,2)

    def reached(self, topic):
        def is_near(msg, x, y, d):
            rospy.logdebug("Position %s: local: %d, target: %d, abs diff: %d",
                           msg, x, y, abs(x - y))
            return abs(x - y) < d
        self.currentPoseX = topic.pose.position.x
        self.currentPoseY = topic.pose.position.y
        self.currentPoseZ = topic.pose.position.z

        if is_near('X', topic.pose.position.x, self.x, 0.2) and \
           is_near('Y', topic.pose.position.y, self.y, 0.2) and \
           is_near('Z', topic.pose.position.z, self.z, 0.5):
            if  self.navigating:
                self.done = True
                self.navigating = False
                self.done_evt.set()

    def explore(self):
        print 'explore started '
        rate = rospy.Rate(30)
        self.newGoal = True
        if self.isExploring == False:
            #Change this later when we have a better exploration
            #self.isExploring = True
            while self.done == False:
                    time.sleep(1)
                    rospy.loginfo("Climb")
                    self.progress += 0.1
                    self.takeoff(5)
                    self.progress += 0.1
                    rospy.loginfo("Moving to Red_Object")
                    self.reached_object = False
                    red_object_id = -1
                    xspeed = 5
                    while self.reached_object == False:
                        self.client.send_goal(self.goal)
                        self.client.wait_for_result()
                        self.objects = self.client.get_result().tracked_objects.objects
                        islost = True
                        for obj in self.objects:
                            if red_object_id == -1 and (obj.color == 'RED' or obj.color == 'BLUE' or obj.color == 'GREEN'): #pick any nearby object
                                red_object_id = obj.object_id
                            if obj.object_id == red_object_id:
                                islost = False
                                print 'Moving to Drop zone', self.currentPoseX-obj.pose.pose.position.x, self.currentPoseY-obj.pose.pose.position.y, obj.pose.pose.position.x, obj.pose.pose.position.y
                                if fabs(obj.pose.pose.position.x) < 0.01 and fabs(obj.pose.pose.position.y) > 0.01:
                                    print 'Moving Y'
                                    self.setPose(self.x, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , False)
                                elif fabs(obj.pose.pose.position.y) < 0.01 and fabs(obj.pose.pose.position.x) > 0.01:
                                    print 'Moving X'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.y, self.z, 0 , False)
                                else:
                                    print 'Moving XY'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , True)
                                if fabs(obj.pose.pose.position.x) < 0.3 and fabs(obj.pose.pose.position.y) < 0.3 and self.z > 0.0:
                                    print 'Moving Z'
                                    land = 0.2
                                    if self.z <= 3:
                                        xspeed = 1
                                    if self.z <= 1.5:
                                        xspeed = 0.5
                                    if self.z < 0.5:
                                        land = 0.05
                                    self.setPose(self.x, self.y, self.z - land * xspeed, 1, False)
                                    if self.z <= 0.4:
                                        self.reached_object = True
                        if islost == True:
                            red_object_id = -1
                        if red_object_id == -1:
                            rospy.loginfo("No object in sight, exploring")
                            #self.setPose(self.x, self.y - 5, self.z, 1, True)
                        rate.sleep()
                    time.sleep(10)
                    rospy.loginfo("Picked Object, climb")
                    self.takeoff(1)
                    self.takeoff(2)
                    self.takeoff(3)
                    self.takeoff(4)
                    self.takeoff(5)
                    #self.setPose(self.x, self.y, self.z)
                    time.sleep(10)
                    rospy.loginfo("Moving to DropZone")
                    self.setPose(1, -21, 5) ##Go near dropzone
                    self.progress += 0.1
                    self.reached_dropzone = False
                    xspeed = 3
                    while self.reached_dropzone == False:
                        self.client.send_goal(self.goal)
                        self.client.wait_for_result()
                        self.objects = self.client.get_result().tracked_objects.objects

                        for obj in self.objects:
                            if obj.color == 'DROP_ZONE':
                                print 'Moving to Drop zone', self.currentPoseX-obj.pose.pose.position.x, self.currentPoseY-obj.pose.pose.position.y, obj.pose.pose.position.x, obj.pose.pose.position.y
                                if fabs(obj.pose.pose.position.x) < 0.1 and fabs(obj.pose.pose.position.y) > 0.1:
                                    print 'Moving Y'
                                    self.setPose(self.x, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , False)
                                elif fabs(obj.pose.pose.position.y) < 0.1 and fabs(obj.pose.pose.position.x) > 0.1:
                                    print 'Moving X'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.y, self.z, 0 , False)
                                else:
                                    print 'Moving XY'
                                    self.setPose(self.currentPoseX-obj.pose.pose.position.x*xspeed, self.currentPoseY+obj.pose.pose.position.y*xspeed, self.z, 0 , True)
                                if fabs(obj.pose.pose.position.x) < 0.3 and fabs(obj.pose.pose.position.y) < 0.3 and self.z > 1:
                                    print 'Moving Z'
                                    land = 0.5
                                    if self.z <= 3:
                                        land = 0.2
                                        xspeed = 0.5
                                    self.setPose(self.x, self.y, self.z - land, 1, False)
                                    if self.z < 1.5:
                                        self.reached_dropzone = True
                        rate.sleep()
                    self.progress += 0.1
                    rospy.loginfo("Landed Object, climb")
                    self.takeoff(7)
                    rospy.loginfo("Bye!")

def main(args):
    '''Initializes and cleanup ros node'''
    rospy.init_node('dropzone_landing', anonymous=True)
    d = dropzone_landing()
    d.explore()
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print "Shutting down ROS Image feature detector module"
    cv2.destroyAllWindows()

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


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

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import tensorflow as tf

from tensorforce import util
from tensorforce.core.memories import Memory

class Queue(Memory):
    """
    Base class for memories organized as a queue (FIFO).
    """

    def __init__(self, states, internals, actions, include_next_states, capacity, scope='queue', summary_labels=None):
        """
        Queue memory.

        Args:
            capacity: Memory capacity.
        """
        self.capacity = capacity
        self.scope = scope

        # Pieces of the records are stored in different tensors:
        self.states_memory = dict()  # keys=state space components
        self.internals_memory = dict()  # keys=internal state components
        self.actions_memory = dict()  # keys=action space components
        self.terminal_memory = None  # 1D tensor
        self.reward_memory = None  # 1D tensor
        self.memory_index = None  # 0D (int) tensor (points to the next record to be overwritten)
        self.episode_indices = None  # 1D tensor of indexes where episodes start.
        self.episode_count = None  # 0D (int) tensor: How many episodes do we have stored?

        self.retrieve_indices = None

        super(Queue, self).__init__(
            states=states,
            internals=internals,
            actions=actions,
            include_next_states=include_next_states,
            scope=scope,
            summary_labels=summary_labels
        )

    def setup_template_funcs(self, custom_getter=None):
        custom_getter = super(Queue, self).setup_template_funcs(custom_getter=custom_getter)

        self.retrieve_indices = tf.make_template(
            name_=(self.scope + '/retrieve_indices'),
            func_=self.tf_retrieve_indices,
            custom_getter_=custom_getter
        )

    def tf_initialize(self):
        # States
        for name, state in self.states_spec.items():
            self.states_memory[name] = tf.get_variable(
                name=('state-' + name),
                shape=(self.capacity,) + tuple(state['shape']),
                dtype=util.tf_dtype(state['type']),
                trainable=False
            )

        # Internals
        for name, internal in self.internals_spec.items():
            self.internals_memory[name] = tf.get_variable(
                name=('internal-' + name),
                shape=(self.capacity,) + tuple(internal['shape']),
                dtype=util.tf_dtype(internal['type']),
                trainable=False
            )

        # Actions
        for name, action in self.actions_spec.items():
            self.actions_memory[name] = tf.get_variable(
                name=('action-' + name),
                shape=(self.capacity,) + tuple(action['shape']),
                dtype=util.tf_dtype(action['type']),
                trainable=False
            )

        # Terminal
        self.terminal_memory = tf.get_variable(
            name='terminal',
            shape=(self.capacity,),
            dtype=util.tf_dtype('bool'),
            initializer=tf.constant_initializer(
                value=tuple(n == self.capacity - 1 for n in range(self.capacity)),
                dtype=util.tf_dtype('bool')
            ),
            trainable=False
        )

        # Reward
        self.reward_memory = tf.get_variable(
            name='reward',
            shape=(self.capacity,),
            dtype=util.tf_dtype('float'),
            trainable=False
        )

        # Memory index
        self.memory_index = tf.get_variable(
            name='memory-index',
            dtype=util.tf_dtype('int'),
            initializer=0,
            trainable=False
        )

        # Episode indices
        self.episode_indices = tf.get_variable(
            name='episode-indices',
            shape=(self.capacity + 1,),
            dtype=util.tf_dtype('int'),
            initializer=tf.constant_initializer(value=(self.capacity - 1), dtype=util.tf_dtype('int')),
            trainable=False
        )

        # Episodes index
        self.episode_count = tf.get_variable(
            name='episode-count',
            dtype=util.tf_dtype('int'),
            initializer=0,
            trainable=False
        )

    def tf_store(self, states, internals, actions, terminal, reward):
        # Memory indices to overwrite.
        num_instances = tf.shape(input=terminal)[0]
        indices = tf.range(start=self.memory_index, limit=(self.memory_index + num_instances)) % self.capacity

        # Remove episode indices.
        num_episodes = tf.count_nonzero(
            input_tensor=tf.gather(params=self.terminal_memory, indices=indices),
            axis=0,
            dtype=util.tf_dtype('int')
        )
        num_episodes = tf.minimum(x=num_episodes, y=self.episode_count)
        assignment = tf.assign(
            ref=self.episode_indices[:self.episode_count + 1 - num_episodes],
            value=self.episode_indices[num_episodes: self.episode_count + 1]
        )

        # Decrement episode count.
        with tf.control_dependencies(control_inputs=(assignment,)):
            assignment = tf.assign_sub(ref=self.episode_count, value=num_episodes)

        # Assign new observations.
        with tf.control_dependencies(control_inputs=(assignment,)):
            assignments = list()
            for name, state in states.items():
                assignments.append(tf.scatter_update(
                    ref=self.states_memory[name],
                    indices=indices,
                    updates=state
                ))
            for name, internal in internals.items():
                assignments.append(tf.scatter_update(
                    ref=self.internals_memory[name],
                    indices=indices,
                    updates=internal
                ))
            for name, action in actions.items():
                assignments.append(tf.scatter_update(
                    ref=self.actions_memory[name],
                    indices=indices,
                    updates=action
                ))
            assignments.append(tf.scatter_update(ref=self.terminal_memory, indices=indices, updates=terminal))
            assignments.append(tf.scatter_update(ref=self.reward_memory, indices=indices, updates=reward))

        # Increment memory index.
        with tf.control_dependencies(control_inputs=assignments):
            assignment = tf.assign(ref=self.memory_index, value=((self.memory_index + num_instances) % self.capacity))

        # Add episode indices.
        with tf.control_dependencies(control_inputs=(assignment,)):
            num_episodes = tf.count_nonzero(input_tensor=terminal, axis=0, dtype=util.tf_dtype('int'))
            assignment = tf.assign(
                ref=self.episode_indices[self.episode_count + 1: self.episode_count + 1 + num_episodes],
                value=tf.boolean_mask(tensor=indices, mask=terminal)
            )

        # Increment episode count.
        with tf.control_dependencies(control_inputs=(assignment,)):
            assignment = tf.assign_add(ref=self.episode_count, value=num_episodes)

        with tf.control_dependencies(control_inputs=(assignment,)):
            return tf.no_op()

    def tf_retrieve_indices(self, indices):
        """
        Fetches experiences for given indices.

        Args:
            indices: Index tensor

        Returns: Batch of experiences
        """
        states = dict()
        for name, state_memory in self.states_memory.items():
            states[name] = tf.gather(params=state_memory, indices=indices)

        internals = dict()
        for name, internal_memory in self.internals_memory.items():
            internals[name] = tf.gather(params=internal_memory, indices=indices)

        actions = dict()
        for name, action_memory in self.actions_memory.items():
            actions[name] = tf.gather(params=action_memory, indices=indices)

        terminal = tf.gather(params=self.terminal_memory, indices=indices)
        reward = tf.gather(params=self.reward_memory, indices=indices)

        if self.include_next_states:
            assert util.rank(indices) == 1
            next_indices = (indices + 1) % self.capacity

            next_states = dict()
            for name, state_memory in self.states_memory.items():
                next_states[name] = tf.gather(params=state_memory, indices=next_indices)

            next_internals = dict()
            for name, internal_memory in self.internals_memory.items():
                next_internals[name] = tf.gather(params=internal_memory, indices=next_indices)

            return dict(
                states=states,
                internals=internals,
                actions=actions,
                terminal=terminal,
                reward=reward,
                next_states=next_states,
                next_internals=next_internals
            )
        else:
            return dict(
                states=states,
                internals=internals,
                actions=actions,
                terminal=terminal,
                reward=reward
            )
<commit_msg>fixing np.bool_ vs tf.bool type issue on memory initialization<commit_after># Copyright 2017 reinforce.io. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import tensorflow as tf

from tensorforce import util
from tensorforce.core.memories import Memory

class Queue(Memory):
    """
    Base class for memories organized as a queue (FIFO).
    """

    def __init__(self, states, internals, actions, include_next_states, capacity, scope='queue', summary_labels=None):
        """
        Queue memory.

        Args:
            capacity: Memory capacity.
        """
        self.capacity = capacity
        self.scope = scope

        # Pieces of the records are stored in different tensors:
        self.states_memory = dict()  # keys=state space components
        self.internals_memory = dict()  # keys=internal state components
        self.actions_memory = dict()  # keys=action space components
        self.terminal_memory = None  # 1D tensor
        self.reward_memory = None  # 1D tensor
        self.memory_index = None  # 0D (int) tensor (points to the next record to be overwritten)
        self.episode_indices = None  # 1D tensor of indexes where episodes start.
        self.episode_count = None  # 0D (int) tensor: How many episodes do we have stored?

        self.retrieve_indices = None

        super(Queue, self).__init__(
            states=states,
            internals=internals,
            actions=actions,
            include_next_states=include_next_states,
            scope=scope,
            summary_labels=summary_labels
        )

    def setup_template_funcs(self, custom_getter=None):
        custom_getter = super(Queue, self).setup_template_funcs(custom_getter=custom_getter)

        self.retrieve_indices = tf.make_template(
            name_=(self.scope + '/retrieve_indices'),
            func_=self.tf_retrieve_indices,
            custom_getter_=custom_getter
        )

    def tf_initialize(self):
        # States
        for name, state in self.states_spec.items():
            self.states_memory[name] = tf.get_variable(
                name=('state-' + name),
                shape=(self.capacity,) + tuple(state['shape']),
                dtype=util.tf_dtype(state['type']),
                trainable=False
            )

        # Internals
        for name, internal in self.internals_spec.items():
            self.internals_memory[name] = tf.get_variable(
                name=('internal-' + name),
                shape=(self.capacity,) + tuple(internal['shape']),
                dtype=util.tf_dtype(internal['type']),
                trainable=False
            )

        # Actions
        for name, action in self.actions_spec.items():
            self.actions_memory[name] = tf.get_variable(
                name=('action-' + name),
                shape=(self.capacity,) + tuple(action['shape']),
                dtype=util.tf_dtype(action['type']),
                trainable=False
            )

        # Terminal
        self.terminal_memory = tf.get_variable(
            name='terminal',
            shape=(self.capacity,),
            dtype=util.tf_dtype('bool'),
            initializer=tf.constant_initializer(
                value=[n == self.capacity - 1 for n in range(self.capacity)],
                dtype=util.tf_dtype('bool')
            ),
            trainable=False
        )

        # Reward
        self.reward_memory = tf.get_variable(
            name='reward',
            shape=(self.capacity,),
            dtype=util.tf_dtype('float'),
            trainable=False
        )

        # Memory index
        self.memory_index = tf.get_variable(
            name='memory-index',
            dtype=util.tf_dtype('int'),
            initializer=0,
            trainable=False
        )

        # Episode indices
        self.episode_indices = tf.get_variable(
            name='episode-indices',
            shape=(self.capacity + 1,),
            dtype=util.tf_dtype('int'),
            initializer=tf.constant_initializer(value=(self.capacity - 1), dtype=util.tf_dtype('int')),
            trainable=False
        )

        # Episodes index
        self.episode_count = tf.get_variable(
            name='episode-count',
            dtype=util.tf_dtype('int'),
            initializer=0,
            trainable=False
        )

    def tf_store(self, states, internals, actions, terminal, reward):
        # Memory indices to overwrite.
        num_instances = tf.shape(input=terminal)[0]
        indices = tf.range(start=self.memory_index, limit=(self.memory_index + num_instances)) % self.capacity

        # Remove episode indices.
        num_episodes = tf.count_nonzero(
            input_tensor=tf.gather(params=self.terminal_memory, indices=indices),
            axis=0,
            dtype=util.tf_dtype('int')
        )
        num_episodes = tf.minimum(x=num_episodes, y=self.episode_count)
        assignment = tf.assign(
            ref=self.episode_indices[:self.episode_count + 1 - num_episodes],
            value=self.episode_indices[num_episodes: self.episode_count + 1]
        )

        # Decrement episode count.
        with tf.control_dependencies(control_inputs=(assignment,)):
            assignment = tf.assign_sub(ref=self.episode_count, value=num_episodes)

        # Assign new observations.
        with tf.control_dependencies(control_inputs=(assignment,)):
            assignments = list()
            for name, state in states.items():
                assignments.append(tf.scatter_update(
                    ref=self.states_memory[name],
                    indices=indices,
                    updates=state
                ))
            for name, internal in internals.items():
                assignments.append(tf.scatter_update(
                    ref=self.internals_memory[name],
                    indices=indices,
                    updates=internal
                ))
            for name, action in actions.items():
                assignments.append(tf.scatter_update(
                    ref=self.actions_memory[name],
                    indices=indices,
                    updates=action
                ))
            assignments.append(tf.scatter_update(ref=self.terminal_memory, indices=indices, updates=terminal))
            assignments.append(tf.scatter_update(ref=self.reward_memory, indices=indices, updates=reward))

        # Increment memory index.
        with tf.control_dependencies(control_inputs=assignments):
            assignment = tf.assign(ref=self.memory_index, value=((self.memory_index + num_instances) % self.capacity))

        # Add episode indices.
        with tf.control_dependencies(control_inputs=(assignment,)):
            num_episodes = tf.count_nonzero(input_tensor=terminal, axis=0, dtype=util.tf_dtype('int'))
            assignment = tf.assign(
                ref=self.episode_indices[self.episode_count + 1: self.episode_count + 1 + num_episodes],
                value=tf.boolean_mask(tensor=indices, mask=terminal)
            )

        # Increment episode count.
        with tf.control_dependencies(control_inputs=(assignment,)):
            assignment = tf.assign_add(ref=self.episode_count, value=num_episodes)

        with tf.control_dependencies(control_inputs=(assignment,)):
            return tf.no_op()

    def tf_retrieve_indices(self, indices):
        """
        Fetches experiences for given indices.

        Args:
            indices: Index tensor

        Returns: Batch of experiences
        """
        states = dict()
        for name, state_memory in self.states_memory.items():
            states[name] = tf.gather(params=state_memory, indices=indices)

        internals = dict()
        for name, internal_memory in self.internals_memory.items():
            internals[name] = tf.gather(params=internal_memory, indices=indices)

        actions = dict()
        for name, action_memory in self.actions_memory.items():
            actions[name] = tf.gather(params=action_memory, indices=indices)

        terminal = tf.gather(params=self.terminal_memory, indices=indices)
        reward = tf.gather(params=self.reward_memory, indices=indices)

        if self.include_next_states:
            assert util.rank(indices) == 1
            next_indices = (indices + 1) % self.capacity

            next_states = dict()
            for name, state_memory in self.states_memory.items():
                next_states[name] = tf.gather(params=state_memory, indices=next_indices)

            next_internals = dict()
            for name, internal_memory in self.internals_memory.items():
                next_internals[name] = tf.gather(params=internal_memory, indices=next_indices)

            return dict(
                states=states,
                internals=internals,
                actions=actions,
                terminal=terminal,
                reward=reward,
                next_states=next_states,
                next_internals=next_internals
            )
        else:
            return dict(
                states=states,
                internals=internals,
                actions=actions,
                terminal=terminal,
                reward=reward
            )
<|endoftext|>