<commit_before>import glob
import os

from fabric.api import task, local
from fabric.context_managers import lcd

class Test(object):
    def __init__(self, func, deps=[], needs_rpython=True, needs_rubyspec=False,
                 builds_release=False):
        self.func = func
        self.deps = deps
        self.needs_rpython = needs_rpython
        self.needs_rubyspec = needs_rubyspec
        self.builds_release = builds_release

    def install_deps(self):
        local("pip install --use-mirrors {}".format(" ".join(self.deps)))

    def download_rpython(self):
        local("wget https://bitbucket.org/pypy/pypy/get/default.tar.bz2 -O `pwd`/../pypy.tar.bz2")
        local("tar -xf `pwd`/../pypy.tar.bz2 -C `pwd`/../")
        [path_name] = glob.glob("../pypy-pypy*")
        path_name = os.path.abspath(path_name)
        with open("rpython_marker", "w") as f:
            f.write(path_name)

    def download_mspec(self):
        with lcd(".."):
            local("git clone --depth=100 --quiet https://github.com/rubyspec/mspec")

    def download_rubyspec(self):
        with lcd(".."):
            local("git clone --depth=100 --quiet https://github.com/rubyspec/rubyspec")

    def run_tests(self):
        env = {}
        if self.needs_rpython:
            with open("rpython_marker") as f:
                env["rpython_path"] = f.read()
        self.func(env)

    def build_release(self):
        local("python topaz/tool/make_release.py topaz.tar")
        # TODO: the part where we upload it somewhere.

@task
def install_requirements():
    t = TEST_TYPES[os.environ["TEST_TYPE"]]
    if t.deps:
        t.install_deps()
    if t.needs_rpython:
        t.download_rpython()
    if t.needs_rubyspec:
        t.download_mspec()
        t.download_rubyspec()

@task
def run_tests():
    t = TEST_TYPES[os.environ["TEST_TYPE"]]
    t.run_tests()

@task
def build_release():
    t = TEST_TYPES[os.environ["TEST_TYPE"]]
    if t.builds_release:
        t.build_release()

def run_own_tests(env):
    local("PYTHONPATH=$PYTHONPATH:{rpython_path} py.test".format(**env))

def run_rubyspec_untranslated(env):
    run_specs("bin/topaz_untranslated.py", prefix="PYTHONPATH=$PYTHONPATH:{rpython_path} ".format(**env))

def run_translate_tests(env):
    local("PYTHONPATH={rpython_path}:$PYTHONPATH python {rpython_path}/rpython/bin/rpython --batch -Ojit targettopaz.py".format(**env))
    run_specs("`pwd`/bin/topaz")
    local("PYTHONPATH={rpython_path}:$PYTHONPATH py.test --topaz=bin/topaz tests/jit/".format(**env))

def run_specs(binary, prefix=""):
    # TODO: this list is temporary until we have all the machinery necessary to
    # run the full rubyspec directory (including the tagging feature)
    rubyspec_tests = [
        "language/and_spec.rb",
        "language/array_spec.rb",
        "language/class_variable_spec.rb",
        "language/match_spec.rb",
        "language/metaclass_spec.rb",
        "language/module_spec.rb",
        "language/next_spec.rb",
        "language/not_spec.rb",
        "language/numbers_spec.rb",
        "language/order_spec.rb",
        "language/splat_spec.rb",
        "language/undef_spec.rb",
        "language/unless_spec.rb",
        "language/yield_spec.rb",

        "language/regexp/grouping_spec.rb",
        "language/regexp/repetition_spec.rb",

        "core/array/allocate_spec.rb",
        "core/array/array_spec.rb",
        "core/array/at_spec.rb",
        "core/array/empty_spec.rb",
        "core/array/frozen_spec.rb",
        "core/array/length_spec.rb",
        "core/array/size_spec.rb",

        "core/basicobject/ancestors_spec.rb",
        "core/basicobject/class_spec.rb",
        "core/basicobject/new_spec.rb",
        "core/basicobject/superclass_spec.rb",

        "core/class/superclass_spec.rb",

        "core/comparable/between_spec.rb",

        "core/exception/message_spec.rb",
        "core/exception/standard_error_spec.rb",

        "core/false/and_spec.rb",
        "core/false/inspect_spec.rb",
        "core/false/or_spec.rb",
        "core/false/to_s_spec.rb",
        "core/false/xor_spec.rb",

        "core/fixnum/comparison_spec.rb",
        "core/fixnum/even_spec.rb",
        "core/fixnum/hash_spec.rb",
        "core/fixnum/odd_spec.rb",
        "core/fixnum/to_f_spec.rb",
        "core/fixnum/zero_spec.rb",

        "core/hash/empty_spec.rb",
        "core/hash/length_spec.rb",
        "core/hash/size_spec.rb",

        "core/nil/and_spec.rb",
        "core/nil/inspect_spec.rb",
        "core/nil/nil_spec.rb",
        "core/nil/or_spec.rb",
        "core/nil/to_a_spec.rb",
        "core/nil/to_i_spec.rb",
        "core/nil/to_s_spec.rb",
        "core/nil/xor_spec.rb",

        "core/object/clone_spec.rb",
        "core/object/dup_spec.rb",
        "core/object/match_spec.rb",

        "core/process/pid_spec.rb",

        "core/regexp/casefold_spec.rb",
        "core/regexp/source_spec.rb",

        "core/true/and_spec.rb",
        "core/true/inspect_spec.rb",
        "core/true/or_spec.rb",
        "core/true/to_s_spec.rb",
        "core/true/xor_spec.rb",
    ]
    local("{prefix}../mspec/bin/mspec -t {binary} --format=dotted {spec_files}".format(
        prefix=prefix,
        binary=binary,
        spec_files=" ".join(os.path.join("../rubyspec", p) for p in rubyspec_tests),
    ))

def run_docs_tests(env):
    local("sphinx-build -W -b html docs/ docs/_build/")

TEST_TYPES = {
    "own": Test(run_own_tests, deps=["pytest", "-r requirements.txt"]),
    "rubyspec_untranslated": Test(run_rubyspec_untranslated, deps=["-r requirements.txt"], needs_rubyspec=True),
    "translate": Test(run_translate_tests, deps=["-r requirements.txt"], needs_rubyspec=True),
    "docs": Test(run_docs_tests, deps=["sphinx"], needs_rpython=False),
}
<commit_msg>actually enable this<commit_after>import glob
import os

from fabric.api import task, local
from fabric.context_managers import lcd

class Test(object):
    def __init__(self, func, deps=[], needs_rpython=True, needs_rubyspec=False,
                 builds_release=False):
        self.func = func
        self.deps = deps
        self.needs_rpython = needs_rpython
        self.needs_rubyspec = needs_rubyspec
        self.builds_release = builds_release

    def install_deps(self):
        local("pip install --use-mirrors {}".format(" ".join(self.deps)))

    def download_rpython(self):
        local("wget https://bitbucket.org/pypy/pypy/get/default.tar.bz2 -O `pwd`/../pypy.tar.bz2")
        local("tar -xf `pwd`/../pypy.tar.bz2 -C `pwd`/../")
        [path_name] = glob.glob("../pypy-pypy*")
        path_name = os.path.abspath(path_name)
        with open("rpython_marker", "w") as f:
            f.write(path_name)

    def download_mspec(self):
        with lcd(".."):
            local("git clone --depth=100 --quiet https://github.com/rubyspec/mspec")

    def download_rubyspec(self):
        with lcd(".."):
            local("git clone --depth=100 --quiet https://github.com/rubyspec/rubyspec")

    def run_tests(self):
        env = {}
        if self.needs_rpython:
            with open("rpython_marker") as f:
                env["rpython_path"] = f.read()
        self.func(env)

    def build_release(self):
        local("python topaz/tool/make_release.py topaz.tar")
        # TODO: the part where we upload it somewhere.

@task
def install_requirements():
    t = TEST_TYPES[os.environ["TEST_TYPE"]]
    if t.deps:
        t.install_deps()
    if t.needs_rpython:
        t.download_rpython()
    if t.needs_rubyspec:
        t.download_mspec()
        t.download_rubyspec()

@task
def run_tests():
    t = TEST_TYPES[os.environ["TEST_TYPE"]]
    t.run_tests()

@task
def build_release():
    t = TEST_TYPES[os.environ["TEST_TYPE"]]
    if t.builds_release:
        t.build_release()

def run_own_tests(env):
    local("PYTHONPATH=$PYTHONPATH:{rpython_path} py.test".format(**env))

def run_rubyspec_untranslated(env):
    run_specs("bin/topaz_untranslated.py", prefix="PYTHONPATH=$PYTHONPATH:{rpython_path} ".format(**env))

def run_translate_tests(env):
    local("PYTHONPATH={rpython_path}:$PYTHONPATH python {rpython_path}/rpython/bin/rpython --batch -Ojit targettopaz.py".format(**env))
    run_specs("`pwd`/bin/topaz")
    local("PYTHONPATH={rpython_path}:$PYTHONPATH py.test --topaz=bin/topaz tests/jit/".format(**env))

def run_specs(binary, prefix=""):
    # TODO: this list is temporary until we have all the machinery necessary to
    # run the full rubyspec directory (including the tagging feature)
    rubyspec_tests = [
        "language/and_spec.rb",
        "language/array_spec.rb",
        "language/class_variable_spec.rb",
        "language/match_spec.rb",
        "language/metaclass_spec.rb",
        "language/module_spec.rb",
        "language/next_spec.rb",
        "language/not_spec.rb",
        "language/numbers_spec.rb",
        "language/order_spec.rb",
        "language/splat_spec.rb",
        "language/undef_spec.rb",
        "language/unless_spec.rb",
        "language/yield_spec.rb",

        "language/regexp/grouping_spec.rb",
        "language/regexp/repetition_spec.rb",

        "core/array/allocate_spec.rb",
        "core/array/array_spec.rb",
        "core/array/at_spec.rb",
        "core/array/empty_spec.rb",
        "core/array/frozen_spec.rb",
        "core/array/length_spec.rb",
        "core/array/size_spec.rb",

        "core/basicobject/ancestors_spec.rb",
        "core/basicobject/class_spec.rb",
        "core/basicobject/new_spec.rb",
        "core/basicobject/superclass_spec.rb",

        "core/class/superclass_spec.rb",

        "core/comparable/between_spec.rb",

        "core/exception/message_spec.rb",
        "core/exception/standard_error_spec.rb",

        "core/false/and_spec.rb",
        "core/false/inspect_spec.rb",
        "core/false/or_spec.rb",
        "core/false/to_s_spec.rb",
        "core/false/xor_spec.rb",

        "core/fixnum/comparison_spec.rb",
        "core/fixnum/even_spec.rb",
        "core/fixnum/hash_spec.rb",
        "core/fixnum/odd_spec.rb",
        "core/fixnum/to_f_spec.rb",
        "core/fixnum/zero_spec.rb",

        "core/hash/empty_spec.rb",
        "core/hash/length_spec.rb",
        "core/hash/size_spec.rb",

        "core/nil/and_spec.rb",
        "core/nil/inspect_spec.rb",
        "core/nil/nil_spec.rb",
        "core/nil/or_spec.rb",
        "core/nil/to_a_spec.rb",
        "core/nil/to_i_spec.rb",
        "core/nil/to_s_spec.rb",
        "core/nil/xor_spec.rb",

        "core/object/clone_spec.rb",
        "core/object/dup_spec.rb",
        "core/object/match_spec.rb",

        "core/process/pid_spec.rb",

        "core/regexp/casefold_spec.rb",
        "core/regexp/source_spec.rb",

        "core/true/and_spec.rb",
        "core/true/inspect_spec.rb",
        "core/true/or_spec.rb",
        "core/true/to_s_spec.rb",
        "core/true/xor_spec.rb",
    ]
    local("{prefix}../mspec/bin/mspec -t {binary} --format=dotted {spec_files}".format(
        prefix=prefix,
        binary=binary,
        spec_files=" ".join(os.path.join("../rubyspec", p) for p in rubyspec_tests),
    ))

def run_docs_tests(env):
    local("sphinx-build -W -b html docs/ docs/_build/")

TEST_TYPES = {
    "own": Test(run_own_tests, deps=["pytest", "-r requirements.txt"]),
    "rubyspec_untranslated": Test(run_rubyspec_untranslated, deps=["-r requirements.txt"], needs_rubyspec=True),
    "translate": Test(run_translate_tests, deps=["-r requirements.txt"], needs_rubyspec=True, builds_release=True),
    "docs": Test(run_docs_tests, deps=["sphinx"], needs_rpython=False),
}
<|endoftext|>


<commit_before>from threading import _Event
from urlparse import urlparse
from functools import partial
from socketIO_client import SocketIO, SocketIONamespace, EngineIONamespace

from exceptions import HeimdallrClientException
from utils import timestamp, for_own_methods, on_ready

__all__ = ['Client', 'Provider', 'Consumer']

def _init(self, io):
    self._io = io
    self._callback_by_event = {}
    self._log_name = Client._url
    self.initialize()

EngineIONamespace.__init__ = _init

class _SocketIO(SocketIO):
    def _should_stop_waiting(self, **kwargs):
        event = kwargs.pop('event', None)
        event_set = False
        if isinstance(event, _Event):
            event_set = event.is_set()
        return super(_SocketIO, self)._should_stop_waiting(**kwargs) or \
            event_set

class Client():
    """
    The ``Client`` class provides most of the behavior for
    Heimdallr clients. However, it is not intended to be
    used directly. The ``Client`` constructor creates the
    basic connection which in this case is a
    ``SocketIONamespace``. It sets up callbacks for
    connection and authentication as well as a default
    callback for errors. The default error handler can be
    removed by ``client.remove_listener('err')``.

    Args:
        token (str): Authentication token
    """

    _url = 'https://heimdallr.co'
    _auth_source = 'heimdallr'
    _namespace = '/'

    def __init__(self, token):
        self.ready = False
        self.ready_callbacks = []
        self.callbacks = {}
        self.token = token
        self.connection = SocketIONamespace(None, self._namespace)

        @self.on('err')
        def fn(err):
            if 'message' in err:
                raise HeimdallrClientException(err['message'])
            else:
                raise HeimdallrClientException(err)

        @self.on('auth-success')
        def fn(*args):
            self.ready = True
            while self.ready_callbacks:
                self.ready_callbacks.pop(0)()

        @self.on('connect')
        def fn(*args):
            self.connection.emit(
                'authorize',
                {'token': self.token, 'authSource': self._auth_source}
            )

    def run(self, seconds=None, **kwargs):
        """ Main loop for a client.

        The ``run`` method is the main loop for a client and is where
        all communication between the Heimdallr server and client
        takes place. The ``run`` method is just a proxy for the
        ``SocketIO.wait`` method so you can call it with the
        same arguments. However, an additional ``event`` option has
        been added. If a :py:class:`threading.Event` object is passed in for
        ``event``, the wait loop will terminate once the flag is set.

        Args:
            seconds (float): Number of seconds to loop for
            event (:py:class:`threading.Event`): Triggers the exit of the run
                loop when the flag is set
            for_connect (bool): Run until the SocketIO connect event
            for_callback (bool): Run until the server has acknowledged all
                emits

        :returns: :class:`Client <Client>`

        **Usage:**

        .. code-block:: python

            client.run(1)  # Loops for 1 second
            from threading import Event
            event = Event()
            client.run(event=event)  # Loops until event.is_set() is True
            client.run()  # Loops forever
        """

        kwargs['seconds'] = seconds
        self.connection._io.wait(**kwargs)

        return self

    def connect(self):
        """ Connect to the Heimdallr server.

        The ``connect`` method blocks until the the socket connection
        to the server has been established.

        :returns: :class:`Client <Client>`
        """

        parsed = urlparse(self._url)
        io = _SocketIO(parsed.hostname, parsed.port)
        io._namespace = self.connection
        io._namespace_by_path[self._namespace] = self.connection
        io.connect(self._namespace)
        io.wait(for_connect=True)
        self.connection._io = io

        return self

    def __trigger_callbacks(self, message_name, *args):
        """ Call all of the callbacks for a socket.io message.

        A version of this method curried with ``message_name``
        is given to the underlying ``SocketIONamespace``. When the
        ``SocketIONamespace`` calls it each of the callbacks that
        have been attached to ``message_name`` will be called.

        Args:
            message_name (str): Name of the socket.io message to listen for
            args: Data sent with message
        """

        callbacks = self.callbacks.get(message_name, [])
        for callback in callbacks:
            callback(*args)

    def __on(self, message_name, callback):
        """ Store ``callback`` and register a placeholder callback.

        Appends ``callback`` to the list of callbacks for the
        given ``message_name``. Also assigns a placeholder
        callback to the underlying ``SocketIONamespace`` so that
        the placeholder can call all of the callbacks in
        the list.

        Args:
            message_name (str): Name of the socket.io message to listen for
            callback (function): Callback to be run when the socket.io message
                is heard
        """

        self.callbacks.setdefault(message_name, [])
        self.callbacks[message_name].append(callback)
        self.connection.on(
            message_name,
            partial(self.__trigger_callbacks, message_name)
        )

    def on(self, message_name, callback=None):
        """ Add a socket.io message listener.

        The ``on`` method will add a callback for socket.io messages
        of the specified message name. Multiple callbacks can be
        added for the same message name. They will be triggered
        in the order in which they were added. This method can be
        called outright or it can be used as a decorator.

        Args:
            message_name (str): Name of the socket.io message to listen for
            callback (function): Callback to run when the socket.io
                message is heard

        :returns: :class:`Client <Client>`

        **Usage:**

        .. code-block:: python

            def first(*args):
                print 'FIRST'

            client.on('myMessage', first)

            @client.on('myMessage')
            def second(*args):
                print 'SECOND'
        """

        # Decorator syntax
        if callback is None:
            def decorator(fn):
                self.__on(message_name, fn)
            return decorator

        # SocketIO-Client syntax
        self.__on(message_name, callback)

        return self

    def remove_listener(self, message_name, callback=None):
        """ Remove listener for socket.io message.

        If ``callback`` is specified, only the callbacks registered
        for ``message_name`` that match ``callback`` will be removed.
        If only ``message_name`` is specified, all of the callbacks
        will be removed.

        Args:
            message_name (str): Name of the socket.io message to remove
            callback (function): Specific callback to remove

        :returns: :class:`Client <Client>`
        """

        if callback:
            while callback in self.callbacks.get(message_name, []):
                self.callbacks[message_name].remove(callback)
        else:
            self.callbacks.pop(message_name, None)
            self.connection._callback_by_event.pop(message_name, None)

        return self

@for_own_methods(on_ready)
class Provider(Client):
    """
    This class should be used to create a Heimdallr provider.
    It inherits most of its functionality but it also
    automatically connects to the provider namespace and
    provides some convenience functions.
    """

    _namespace = '/provider'

    def send_event(self, subtype, data=None):
        """ Emit a Heimdallr event packet.

        This will send a Heimdallr event packet to the
        Heimdallr server where it will be rebroadcast.
        ``data`` must adhere to the provider's schema for
        the given ``subtype``.

        Args:
            subtype (str): The event packet subtype
            data: The event packet data

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'event',
            {'subtype': subtype, 'data': data, 't': timestamp()}
        )

    def send_sensor(self, subtype, data=None):
        """ Emit a Heimdallr sensor packet.

        This will send a Heimdallr sensor packet to the
        Heimdallr server where it will be rebroadcast.
        ``data`` must adhere to the provider's schema for
        the given ``subtype``.

        Args:
            subtype (str): The sensor packet subtype
            data: The sensor packet data

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'sensor',
            {'subtype': subtype, 'data': data, 't': timestamp()}
        )

    def send_stream(self, data):
        """ Send binary data to the Heimdallr server.

        This should only be used when the Heimdallr server
        has issued a ``{'stream': 'start'}`` control packet
        and should stop being used when the Heimdallr
        server issues a ``{'stream': 'start'}`` control
        packet.

        Args:
            data: The binary data to be sent.

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'stream',
            data
        )

    def completed(self, uuid):
        """ Signal the Heimdallr server that a control has been completed.

        This should be used when a control that has a persistent
        field set to ``uuid`` has been completed.

        Args:
            uuid (str): UUID of the persistent control packet that has been
                completed

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'event',
            {'subtype': 'completed', 'data': uuid, 't': timestamp()}
        )

@for_own_methods(on_ready)
class Consumer(Client):
    """
    This class should be used to create a Heimdallr consumer.
    It inherits most of its functionality but it also
    automatically connects to the consumer namespace and
    provides some convenience functions.
    """

    _namespace = '/consumer'

    def send_control(self, uuid, subtype, data=None, persistent=False):
        """ Emit a Heimdallr control packet.

        This will send a control to the provider specified by
        ``uuid``. ``data`` must adhere to the provider's schema
        for the given ``subtype``. If `persistent` is ``True``,
        the control packet will be sent immediately and then
        again every time the provider connects until the
        provider signals the Heimdallr server that it has
        completed the control.

        Args:
            uuid (str): UUID of the provider to send the control packet to
            subtype (str): The control packet subtype
            data: The control packet data
            persistent (bool): Whether or not the control should persist

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'control',
            {
                'provider': uuid,
                'subtype': subtype,
                'data': data,
                'persistent': persistent
            }
        )

    def subscribe(self, uuid):
        """ Subscribe to a provider.

        A consumer must subscribe to a provider before it
        receives event or sensor packets from the provider
        or can send control packets to the provider.

        Args:
            uuid (str): UUID of the provider to subscribe to

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'subscribe',
            {'provider': uuid}
        )

    def unsubscribe(self, uuid):
        """ Unsubscribe from a provider.

        The consumer will no longer receive packets from the
        provider or be able to send it controls. This will
        be done automatically by the Heimdallr server on
        disconnect.

        Args:
            uuid (str): UUID of the provider to subscribe to

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'unsubscribe',
            {'provider': uuid}
        )

    def set_filter(self, uuid, filter_):
        """ Control which event and sensor subtypes to hear from provider.

        Set which packet subtypes you want to hear from the provider.
        `filter` should be a dictionary with the keys `event` and/or
        `sensor`. The value of those fields should be an array of
        strings of the subtypes that you want to hear for the
        provider given by `uuid`.

        Args:
            uuid (str): UUID of the provider to filter packets from
            filter (dict): Dictionary containing event and/or sensor packet
                subtypes that you want to receive

        :returns: :class:`Consumer <Consumer>`
        """

        filter_['provider'] = uuid
        self.connection.emit(
            'setFilter',
            filter_
        )

    def get_state(self, uuid, subtypes):
        """ Get the current state of a provider.

        For each event packet subtype in `subtypes`, the most recent
        event packet of that subtype will be sent to the consumer by
        the Heimdallr server.

        Args:
            uuid (str): UUID of the provider to get the state of
            subtypes (list): Event subtypes to get the state of

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'getState',
            {'provider': uuid, 'subtypes': subtypes}
        )

    def join_stream(self, uuid):
        """ Join binary data stream from a provider.

        If this is the first consumer to join the stream of
        a provider, the Heimdallr server will send a
        ``{'stream': 'start'}`` control packet to the provider.

        Args:
            uuid (str): UUID of the provider to join the stream of

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'joinStream',
            {'provider': uuid}
        )

    def leave_stream(self, uuid):
        """ Leave binary data stream for a provider.

        If this is the last consumer to leave the stream for a
        provider the Heimdallr server will send a
        ``{'stream': 'stop'}`` control packet to the provider.
        This will be done automatically by the Heimdallr server
        on disconnect.

        Args:
            uuid (str): UUID of the provider to leave the stream of

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'leaveStream',
            {'provider': uuid}
        )<commit_msg>Need to share the models before calling connect<commit_after>from threading import _Event
from urlparse import urlparse
from functools import partial
from socketIO_client import SocketIO, SocketIONamespace, EngineIONamespace

from exceptions import HeimdallrClientException
from utils import timestamp, for_own_methods, on_ready

__all__ = ['Client', 'Provider', 'Consumer']

def _init(self, io):
    self._io = io
    self._callback_by_event = {}
    self._log_name = Client._url
    self.initialize()

EngineIONamespace.__init__ = _init

class _SocketIO(SocketIO):
    def _should_stop_waiting(self, **kwargs):
        event = kwargs.pop('event', None)
        event_set = False
        if isinstance(event, _Event):
            event_set = event.is_set()
        return super(_SocketIO, self)._should_stop_waiting(**kwargs) or \
            event_set

class Client():
    """
    The ``Client`` class provides most of the behavior for
    Heimdallr clients. However, it is not intended to be
    used directly. The ``Client`` constructor creates the
    basic connection which in this case is a
    ``SocketIONamespace``. It sets up callbacks for
    connection and authentication as well as a default
    callback for errors. The default error handler can be
    removed by ``client.remove_listener('err')``.

    Args:
        token (str): Authentication token
    """

    _url = 'https://heimdallr.co'
    _auth_source = 'heimdallr'
    _namespace = '/'

    def __init__(self, token):
        self.ready = False
        self.ready_callbacks = []
        self.callbacks = {}
        self.token = token
        self.connection = SocketIONamespace(None, self._namespace)

        @self.on('err')
        def fn(err):
            if 'message' in err:
                raise HeimdallrClientException(err['message'])
            else:
                raise HeimdallrClientException(err)

        @self.on('auth-success')
        def fn(*args):
            self.ready = True
            while self.ready_callbacks:
                self.ready_callbacks.pop(0)()

        @self.on('connect')
        def fn(*args):
            self.connection.emit(
                'authorize',
                {'token': self.token, 'authSource': self._auth_source}
            )

    def run(self, seconds=None, **kwargs):
        """ Main loop for a client.

        The ``run`` method is the main loop for a client and is where
        all communication between the Heimdallr server and client
        takes place. The ``run`` method is just a proxy for the
        ``SocketIO.wait`` method so you can call it with the
        same arguments. However, an additional ``event`` option has
        been added. If a :py:class:`threading.Event` object is passed in for
        ``event``, the wait loop will terminate once the flag is set.

        Args:
            seconds (float): Number of seconds to loop for
            event (:py:class:`threading.Event`): Triggers the exit of the run
                loop when the flag is set
            for_connect (bool): Run until the SocketIO connect event
            for_callback (bool): Run until the server has acknowledged all
                emits

        :returns: :class:`Client <Client>`

        **Usage:**

        .. code-block:: python

            client.run(1)  # Loops for 1 second
            from threading import Event
            event = Event()
            client.run(event=event)  # Loops until event.is_set() is True
            client.run()  # Loops forever
        """

        kwargs['seconds'] = seconds
        self.connection._io.wait(**kwargs)

        return self

    def connect(self):
        """ Connect to the Heimdallr server.

        The ``connect`` method blocks until the the socket connection
        to the server has been established.

        :returns: :class:`Client <Client>`
        """

        parsed = urlparse(self._url)
        self.connection._io = _SocketIO(parsed.hostname, parsed.port)
        io = self.connection._io
        io._namespace = self.connection
        io._namespace_by_path[self._namespace] = self.connection
        io.connect(self._namespace)
        io.wait(for_connect=True)

        return self

    def __trigger_callbacks(self, message_name, *args):
        """ Call all of the callbacks for a socket.io message.

        A version of this method curried with ``message_name``
        is given to the underlying ``SocketIONamespace``. When the
        ``SocketIONamespace`` calls it each of the callbacks that
        have been attached to ``message_name`` will be called.

        Args:
            message_name (str): Name of the socket.io message to listen for
            args: Data sent with message
        """

        callbacks = self.callbacks.get(message_name, [])
        for callback in callbacks:
            callback(*args)

    def __on(self, message_name, callback):
        """ Store ``callback`` and register a placeholder callback.

        Appends ``callback`` to the list of callbacks for the
        given ``message_name``. Also assigns a placeholder
        callback to the underlying ``SocketIONamespace`` so that
        the placeholder can call all of the callbacks in
        the list.

        Args:
            message_name (str): Name of the socket.io message to listen for
            callback (function): Callback to be run when the socket.io message
                is heard
        """

        self.callbacks.setdefault(message_name, [])
        self.callbacks[message_name].append(callback)
        self.connection.on(
            message_name,
            partial(self.__trigger_callbacks, message_name)
        )

    def on(self, message_name, callback=None):
        """ Add a socket.io message listener.

        The ``on`` method will add a callback for socket.io messages
        of the specified message name. Multiple callbacks can be
        added for the same message name. They will be triggered
        in the order in which they were added. This method can be
        called outright or it can be used as a decorator.

        Args:
            message_name (str): Name of the socket.io message to listen for
            callback (function): Callback to run when the socket.io
                message is heard

        :returns: :class:`Client <Client>`

        **Usage:**

        .. code-block:: python

            def first(*args):
                print 'FIRST'

            client.on('myMessage', first)

            @client.on('myMessage')
            def second(*args):
                print 'SECOND'
        """

        # Decorator syntax
        if callback is None:
            def decorator(fn):
                self.__on(message_name, fn)
            return decorator

        # SocketIO-Client syntax
        self.__on(message_name, callback)

        return self

    def remove_listener(self, message_name, callback=None):
        """ Remove listener for socket.io message.

        If ``callback`` is specified, only the callbacks registered
        for ``message_name`` that match ``callback`` will be removed.
        If only ``message_name`` is specified, all of the callbacks
        will be removed.

        Args:
            message_name (str): Name of the socket.io message to remove
            callback (function): Specific callback to remove

        :returns: :class:`Client <Client>`
        """

        if callback:
            while callback in self.callbacks.get(message_name, []):
                self.callbacks[message_name].remove(callback)
        else:
            self.callbacks.pop(message_name, None)
            self.connection._callback_by_event.pop(message_name, None)

        return self

@for_own_methods(on_ready)
class Provider(Client):
    """
    This class should be used to create a Heimdallr provider.
    It inherits most of its functionality but it also
    automatically connects to the provider namespace and
    provides some convenience functions.
    """

    _namespace = '/provider'

    def send_event(self, subtype, data=None):
        """ Emit a Heimdallr event packet.

        This will send a Heimdallr event packet to the
        Heimdallr server where it will be rebroadcast.
        ``data`` must adhere to the provider's schema for
        the given ``subtype``.

        Args:
            subtype (str): The event packet subtype
            data: The event packet data

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'event',
            {'subtype': subtype, 'data': data, 't': timestamp()}
        )

    def send_sensor(self, subtype, data=None):
        """ Emit a Heimdallr sensor packet.

        This will send a Heimdallr sensor packet to the
        Heimdallr server where it will be rebroadcast.
        ``data`` must adhere to the provider's schema for
        the given ``subtype``.

        Args:
            subtype (str): The sensor packet subtype
            data: The sensor packet data

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'sensor',
            {'subtype': subtype, 'data': data, 't': timestamp()}
        )

    def send_stream(self, data):
        """ Send binary data to the Heimdallr server.

        This should only be used when the Heimdallr server
        has issued a ``{'stream': 'start'}`` control packet
        and should stop being used when the Heimdallr
        server issues a ``{'stream': 'start'}`` control
        packet.

        Args:
            data: The binary data to be sent.

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'stream',
            data
        )

    def completed(self, uuid):
        """ Signal the Heimdallr server that a control has been completed.

        This should be used when a control that has a persistent
        field set to ``uuid`` has been completed.

        Args:
            uuid (str): UUID of the persistent control packet that has been
                completed

        :returns: :class:`Provider <Provider>`
        """

        self.connection.emit(
            'event',
            {'subtype': 'completed', 'data': uuid, 't': timestamp()}
        )

@for_own_methods(on_ready)
class Consumer(Client):
    """
    This class should be used to create a Heimdallr consumer.
    It inherits most of its functionality but it also
    automatically connects to the consumer namespace and
    provides some convenience functions.
    """

    _namespace = '/consumer'

    def send_control(self, uuid, subtype, data=None, persistent=False):
        """ Emit a Heimdallr control packet.

        This will send a control to the provider specified by
        ``uuid``. ``data`` must adhere to the provider's schema
        for the given ``subtype``. If `persistent` is ``True``,
        the control packet will be sent immediately and then
        again every time the provider connects until the
        provider signals the Heimdallr server that it has
        completed the control.

        Args:
            uuid (str): UUID of the provider to send the control packet to
            subtype (str): The control packet subtype
            data: The control packet data
            persistent (bool): Whether or not the control should persist

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'control',
            {
                'provider': uuid,
                'subtype': subtype,
                'data': data,
                'persistent': persistent
            }
        )

    def subscribe(self, uuid):
        """ Subscribe to a provider.

        A consumer must subscribe to a provider before it
        receives event or sensor packets from the provider
        or can send control packets to the provider.

        Args:
            uuid (str): UUID of the provider to subscribe to

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'subscribe',
            {'provider': uuid}
        )

    def unsubscribe(self, uuid):
        """ Unsubscribe from a provider.

        The consumer will no longer receive packets from the
        provider or be able to send it controls. This will
        be done automatically by the Heimdallr server on
        disconnect.

        Args:
            uuid (str): UUID of the provider to subscribe to

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'unsubscribe',
            {'provider': uuid}
        )

    def set_filter(self, uuid, filter_):
        """ Control which event and sensor subtypes to hear from provider.

        Set which packet subtypes you want to hear from the provider.
        `filter` should be a dictionary with the keys `event` and/or
        `sensor`. The value of those fields should be an array of
        strings of the subtypes that you want to hear for the
        provider given by `uuid`.

        Args:
            uuid (str): UUID of the provider to filter packets from
            filter (dict): Dictionary containing event and/or sensor packet
                subtypes that you want to receive

        :returns: :class:`Consumer <Consumer>`
        """

        filter_['provider'] = uuid
        self.connection.emit(
            'setFilter',
            filter_
        )

    def get_state(self, uuid, subtypes):
        """ Get the current state of a provider.

        For each event packet subtype in `subtypes`, the most recent
        event packet of that subtype will be sent to the consumer by
        the Heimdallr server.

        Args:
            uuid (str): UUID of the provider to get the state of
            subtypes (list): Event subtypes to get the state of

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'getState',
            {'provider': uuid, 'subtypes': subtypes}
        )

    def join_stream(self, uuid):
        """ Join binary data stream from a provider.

        If this is the first consumer to join the stream of
        a provider, the Heimdallr server will send a
        ``{'stream': 'start'}`` control packet to the provider.

        Args:
            uuid (str): UUID of the provider to join the stream of

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'joinStream',
            {'provider': uuid}
        )

    def leave_stream(self, uuid):
        """ Leave binary data stream for a provider.

        If this is the last consumer to leave the stream for a
        provider the Heimdallr server will send a
        ``{'stream': 'stop'}`` control packet to the provider.
        This will be done automatically by the Heimdallr server
        on disconnect.

        Args:
            uuid (str): UUID of the provider to leave the stream of

        :returns: :class:`Consumer <Consumer>`
        """

        self.connection.emit(
            'leaveStream',
            {'provider': uuid}
        )<|endoftext|>


<commit_before>import datetime
import email
import re
import time
import os
from email.header import decode_header, make_header
from imaplib import ParseFlags

class Message():

    def __init__(self, mailbox, uid):
        self.uid = uid
        self.mailbox = mailbox
        self.gmail = mailbox.gmail if mailbox else None

        self.message = None
        self.headers = {}

        self.subject = None
        self.body = None
        self.html = None

        self.to = None
        self.fr = None
        self.cc = None
        self.delivered_to = None

        self.sent_at = None

        self.flags = []
        self.labels = []

        self.thread_id = None
        self.thread = []
        self.message_id = None

        self.attachments = None

    def is_read(self):
        return ('\\Seen' in self.flags)

    def read(self):
        flag = '\\Seen'
        self.gmail.imap.uid('STORE', self.uid, '+FLAGS', flag)
        if flag not in self.flags: self.flags.append(flag)

    def unread(self):
        flag = '\\Seen'
        self.gmail.imap.uid('STORE', self.uid, '-FLAGS', flag)
        if flag in self.flags: self.flags.remove(flag)

    def is_starred(self):
        return ('\\Flagged' in self.flags)

    def star(self):
        flag = '\\Flagged'
        self.gmail.imap.uid('STORE', self.uid, '+FLAGS', flag)
        if flag not in self.flags: self.flags.append(flag)

    def unstar(self):
        flag = '\\Flagged'
        self.gmail.imap.uid('STORE', self.uid, '-FLAGS', flag)
        if flag in self.flags: self.flags.remove(flag)

    def is_draft(self):
        return ('\\Draft' in self.flags)

    def has_label(self, label):
        full_label = '%s' % label
        return (full_label in self.labels)

    def add_label(self, label):
        full_label = '%s' % label
        self.gmail.imap.uid('STORE', self.uid, '+X-GM-LABELS', full_label)
        if full_label not in self.labels: self.labels.append(full_label)

    def remove_label(self, label):
        full_label = '%s' % label
        self.gmail.imap.uid('STORE', self.uid, '-X-GM-LABELS', full_label)
        if full_label in self.labels: self.labels.remove(full_label)

    def is_deleted(self):
        return ('\\Deleted' in self.flags)

    def delete(self):
        flag = '\\Deleted'
        self.gmail.imap.uid('STORE', self.uid, '+FLAGS', flag)
        if flag not in self.flags: self.flags.append(flag)

        trash = '[Gmail]/Trash' if '[Gmail]/Trash' in self.gmail.labels() else '[Gmail]/Bin'
        if self.mailbox.name not in ['[Gmail]/Bin', '[Gmail]/Trash']:
            self.move_to(trash)

    # def undelete(self):
    #     flag = '\\Deleted'
    #     self.gmail.imap.uid('STORE', self.uid, '-FLAGS', flag)
    #     if flag in self.flags: self.flags.remove(flag)

    def move_to(self, name):
        self.gmail.copy(self.uid, name, self.mailbox.name)
        if name not in ['[Gmail]/Bin', '[Gmail]/Trash']:
            self.delete()

    def archive(self):
        self.move_to('[Gmail]/All Mail')

    def parse_headers(self, message):
        hdrs = {}
        for hdr in message.keys():
            hdrs[hdr] = message[hdr]
        return hdrs

    def parse_flags(self, headers):
        return list(ParseFlags(headers))
        # flags = re.search(r'FLAGS \(([^\)]*)\)', headers).groups(1)[0].split(' ')

    def parse_labels(self, headers):
        if re.search(r'X-GM-LABELS \(([^\)]+)\)', headers):
            labels = re.search(r'X-GM-LABELS \(([^\)]+)\)', headers).groups(1)[0].split(' ')
            return map(lambda l: l.replace('"', '').decode("string_escape"), labels)
        else:
            return list()

    def parse_subject(self, encoded_subject):
        dh = decode_header(encoded_subject)
        default_charset = 'ASCII'
        return ''.join([ unicode(t[0], t[1] or default_charset) for t in dh ])

    def parse(self, raw_message):
        raw_headers = raw_message[0]
        raw_email = raw_message[1]

        self.message = email.message_from_string(raw_email)
        self.headers = self.parse_headers(self.message)

        self.to = self.message['to']
        self.fr = self.message['from']
        self.delivered_to = self.message['delivered_to']

        self.subject = self.parse_subject(self.message['subject'])

        if self.message.get_content_maintype() == "multipart":
            for content in self.message.walk():
                if content.get_content_type() == "text/plain":
                    self.body = content.get_payload(decode=True)
                elif content.get_content_type() == "text/html":
                    self.html = content.get_payload(decode=True)
        elif self.message.get_content_maintype() == "text":
            self.body = self.message.get_payload()

        self.sent_at = datetime.datetime.fromtimestamp(time.mktime(email.utils.parsedate_tz(self.message['date'])[:9]))

        self.flags = self.parse_flags(raw_headers)

        self.labels = self.parse_labels(raw_headers)

        if re.search(r'X-GM-THRID (\d+)', raw_headers):
            self.thread_id = re.search(r'X-GM-THRID (\d+)', raw_headers).groups(1)[0]
        if re.search(r'X-GM-MSGID (\d+)', raw_headers):
            self.message_id = re.search(r'X-GM-MSGID (\d+)', raw_headers).groups(1)[0]

        # Parse attachments into attachment objects array for this message
        self.attachments = [
            Attachment(attachment) for attachment in self.message._payload
                if not isinstance(attachment, basestring) and attachment.get('Content-Disposition') is not None
        ]

    def fetch(self):
        if not self.message:
            response, results = self.gmail.imap.uid('FETCH', self.uid, '(BODY.PEEK[] FLAGS X-GM-THRID X-GM-MSGID X-GM-LABELS)')

            self.parse(results[0])

        return self.message

    # returns a list of fetched messages (both sent and received) in chronological order
    def fetch_thread(self):
        self.fetch()
        original_mailbox = self.mailbox
        self.gmail.use_mailbox(original_mailbox.name)

        # fetch and cache messages from inbox or other received mailbox
        response, results = self.gmail.imap.uid('SEARCH', None, '(X-GM-THRID ' + self.thread_id + ')')
        received_messages = {}
        uids = results[0].split(' ')
        if response == 'OK':
            for uid in uids: received_messages[uid] = Message(original_mailbox, uid)
            self.gmail.fetch_multiple_messages(received_messages)
            self.mailbox.messages.update(received_messages)

        # fetch and cache messages from 'sent'
        self.gmail.use_mailbox('[Gmail]/Sent Mail')
        response, results = self.gmail.imap.uid('SEARCH', None, '(X-GM-THRID ' + self.thread_id + ')')
        sent_messages = {}
        uids = results[0].split(' ')
        if response == 'OK':
            for uid in uids: sent_messages[uid] = Message(self.gmail.mailboxes['[Gmail]/Sent Mail'], uid)
            self.gmail.fetch_multiple_messages(sent_messages)
            self.gmail.mailboxes['[Gmail]/Sent Mail'].messages.update(sent_messages)

        self.gmail.use_mailbox(original_mailbox.name)

        # combine and sort sent and received messages
        return sorted(dict(received_messages.items() + sent_messages.items()).values(), key=lambda m: m.sent_at)

class Attachment:

    def __init__(self, attachment):
        self.name = decode_header(attachment.get_filename())[0][0]
        # Raw file data
        self.payload = attachment.get_payload(decode=True)
        # Filesize in kilobytes
        self.size = int(round(len(self.payload)/1000.0))

    def save(self, path=None):
        if path is None:
            # Save as name of attachment if there is no path specified
            path = self.name
        elif os.path.isdir(path):
            # If the path is a directory, save as name of attachment in that directory
            path = os.path.join(path, self.name)

        with open(path, 'wb') as f:
            f.write(self.payload)
<commit_msg>re-add changes to find apple attachments, with old typos fixed<commit_after>import datetime
import email
import re
import time
import os
from email.header import decode_header, make_header
from imaplib import ParseFlags

class Message():

    def __init__(self, mailbox, uid):
        self.uid = uid
        self.mailbox = mailbox
        self.gmail = mailbox.gmail if mailbox else None

        self.message = None
        self.headers = {}

        self.subject = None
        self.body = None
        self.html = None

        self.partitioned_body = []
        self.partitioned_html = []

        self.to = None
        self.fr = None
        self.cc = None
        self.delivered_to = None

        self.sent_at = None

        self.flags = []
        self.labels = []

        self.thread_id = None
        self.thread = []
        self.message_id = None

        self.attachments = None

    def is_read(self):
        return ('\\Seen' in self.flags)

    def read(self):
        flag = '\\Seen'
        self.gmail.imap.uid('STORE', self.uid, '+FLAGS', flag)
        if flag not in self.flags: self.flags.append(flag)

    def unread(self):
        flag = '\\Seen'
        self.gmail.imap.uid('STORE', self.uid, '-FLAGS', flag)
        if flag in self.flags: self.flags.remove(flag)

    def is_starred(self):
        return ('\\Flagged' in self.flags)

    def star(self):
        flag = '\\Flagged'
        self.gmail.imap.uid('STORE', self.uid, '+FLAGS', flag)
        if flag not in self.flags: self.flags.append(flag)

    def unstar(self):
        flag = '\\Flagged'
        self.gmail.imap.uid('STORE', self.uid, '-FLAGS', flag)
        if flag in self.flags: self.flags.remove(flag)

    def is_draft(self):
        return ('\\Draft' in self.flags)

    def has_label(self, label):
        full_label = '%s' % label
        return (full_label in self.labels)

    def add_label(self, label):
        full_label = '%s' % label
        self.gmail.imap.uid('STORE', self.uid, '+X-GM-LABELS', full_label)
        if full_label not in self.labels: self.labels.append(full_label)

    def remove_label(self, label):
        full_label = '%s' % label
        self.gmail.imap.uid('STORE', self.uid, '-X-GM-LABELS', full_label)
        if full_label in self.labels: self.labels.remove(full_label)

    def is_deleted(self):
        return ('\\Deleted' in self.flags)

    def delete(self):
        flag = '\\Deleted'
        self.gmail.imap.uid('STORE', self.uid, '+FLAGS', flag)
        if flag not in self.flags: self.flags.append(flag)

        trash = '[Gmail]/Trash' if '[Gmail]/Trash' in self.gmail.labels() else '[Gmail]/Bin'
        if self.mailbox.name not in ['[Gmail]/Bin', '[Gmail]/Trash']:
            self.move_to(trash)

    # def undelete(self):
    #     flag = '\\Deleted'
    #     self.gmail.imap.uid('STORE', self.uid, '-FLAGS', flag)
    #     if flag in self.flags: self.flags.remove(flag)

    def move_to(self, name):
        self.gmail.copy(self.uid, name, self.mailbox.name)
        if name not in ['[Gmail]/Bin', '[Gmail]/Trash']:
            self.delete()

    def archive(self):
        self.move_to('[Gmail]/All Mail')

    def parse_headers(self, message):
        hdrs = {}
        for hdr in message.keys():
            hdrs[hdr] = message[hdr]
        return hdrs

    def parse_flags(self, headers):
        return list(ParseFlags(headers))
        # flags = re.search(r'FLAGS \(([^\)]*)\)', headers).groups(1)[0].split(' ')

    def parse_labels(self, headers):
        if re.search(r'X-GM-LABELS \(([^\)]+)\)', headers):
            labels = re.search(r'X-GM-LABELS \(([^\)]+)\)', headers).groups(1)[0].split(' ')
            return map(lambda l: l.replace('"', '').decode("string_escape"), labels)
        else:
            return list()

    def parse_subject(self, encoded_subject):
        dh = decode_header(encoded_subject)
        default_charset = 'ASCII'
        return ''.join([ unicode(t[0], t[1] or default_charset) for t in dh ])

    def parse(self, raw_message):
        raw_headers = raw_message[0]
        raw_email = raw_message[1]

        self.message = email.message_from_string(raw_email)
        self.headers = self.parse_headers(self.message)

        self.to = self.message['to']
        self.fr = self.message['from']
        self.delivered_to = self.message['delivered_to']

        self.subject = self.parse_subject(self.message['subject'])

        if self.message.is_multipart():

            for part in self.message.walk():

                # since we're already walking the tree, we only
                # care about leaf nodes; ignore everything else
                if not part.is_multipart():

                    content_disposition = part.get('Content-Disposition', None)
                    if content_disposition:
                        # if it has a content disposition, it should
                        # be an attachment of some kind
                        self.attachments.append(Attachment(part))
                    else:
                        content = part.get_payload(decode=True)
                        content_type = part.get_content_type()
                        if content_type == "text/plain":
                            self.partitioned_body.append(content)
                        elif content_type == "text/html":
                            self.partitioned_html.append(content)

            # This determines what the body and html of the
            # email appears to be to our client. Because of
            # multipart content, there can be multiple values
            # for each; here we take the longest value, which
            # serves to drop the empty HTML that gets created
            # with attachments
            self.body = max(self.partitioned_body)
            self.html = max(self.partitioned_html)

        elif self.message.get_content_maintype() == "text":
            self.body = self.message.get_payload()

        self.sent_at = datetime.datetime.fromtimestamp(time.mktime(email.utils.parsedate_tz(self.message['date'])[:9]))

        self.flags = self.parse_flags(raw_headers)

        self.labels = self.parse_labels(raw_headers)

        if re.search(r'X-GM-THRID (\d+)', raw_headers):
            self.thread_id = re.search(r'X-GM-THRID (\d+)', raw_headers).groups(1)[0]
        if re.search(r'X-GM-MSGID (\d+)', raw_headers):
            self.message_id = re.search(r'X-GM-MSGID (\d+)', raw_headers).groups(1)[0]

        # Parse attachments into attachment objects array for this message
        self.attachments = [
            Attachment(attachment) for attachment in self.message._payload
                if not isinstance(attachment, basestring) and attachment.get('Content-Disposition') is not None
        ]

    def fetch(self):
        if not self.message:
            response, results = self.gmail.imap.uid('FETCH', self.uid, '(BODY.PEEK[] FLAGS X-GM-THRID X-GM-MSGID X-GM-LABELS)')

            self.parse(results[0])

        return self.message

    # returns a list of fetched messages (both sent and received) in chronological order
    def fetch_thread(self):
        self.fetch()
        original_mailbox = self.mailbox
        self.gmail.use_mailbox(original_mailbox.name)

        # fetch and cache messages from inbox or other received mailbox
        response, results = self.gmail.imap.uid('SEARCH', None, '(X-GM-THRID ' + self.thread_id + ')')
        received_messages = {}
        uids = results[0].split(' ')
        if response == 'OK':
            for uid in uids: received_messages[uid] = Message(original_mailbox, uid)
            self.gmail.fetch_multiple_messages(received_messages)
            self.mailbox.messages.update(received_messages)

        # fetch and cache messages from 'sent'
        self.gmail.use_mailbox('[Gmail]/Sent Mail')
        response, results = self.gmail.imap.uid('SEARCH', None, '(X-GM-THRID ' + self.thread_id + ')')
        sent_messages = {}
        uids = results[0].split(' ')
        if response == 'OK':
            for uid in uids: sent_messages[uid] = Message(self.gmail.mailboxes['[Gmail]/Sent Mail'], uid)
            self.gmail.fetch_multiple_messages(sent_messages)
            self.gmail.mailboxes['[Gmail]/Sent Mail'].messages.update(sent_messages)

        self.gmail.use_mailbox(original_mailbox.name)

        # combine and sort sent and received messages
        return sorted(dict(received_messages.items() + sent_messages.items()).values(), key=lambda m: m.sent_at)

class Attachment:

    def __init__(self, attachment):
        self.name = decode_header(attachment.get_filename())[0][0]
        # Raw file data
        self.payload = attachment.get_payload(decode=True)
        # Filesize in kilobytes
        self.size = int(round(len(self.payload)/1000.0))

    def save(self, path=None):
        if path is None:
            # Save as name of attachment if there is no path specified
            path = self.name
        elif os.path.isdir(path):
            # If the path is a directory, save as name of attachment in that directory
            path = os.path.join(path, self.name)

        with open(path, 'wb') as f:
            f.write(self.payload)
<|endoftext|>


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

# Copyright (c) 2013, Felix Kolbe
# All rights reserved. BSD License
#
# 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 the {organization} 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 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.

import roslib; roslib.load_manifest('goap')

import thread

import rospy
import rostopic
import tf

from smach import Sequence, State, StateMachine
from smach_ros import ActionServerWrapper, IntrospectionServer

from std_msgs.msg import String
from geometry_msgs.msg import Pose, Point, Quaternion
from task_msgs.msg import TaskActivationAction, TaskActivationGoal

from uashh_smach.util import CheckSmachEnabledState, TopicToOutcomeState, UserDataToOutcomeState, SleepState, execute_smach_container
from uashh_smach.platform.move_base import WaitForGoalState, get_random_goal_smach
from uashh_smach.manipulator.look_around import get_lookaround_smach
from uashh_smach.tasks import task_go_and_return, task_move_around, task_patrol

from common import ActionBag, Condition, Goal, Precondition, WorldState
from inheriting import Memory
from planning import Planner, PlanExecutor
from introspection import Introspector
from smach_bridge import SmachStateAction, GOAPActionWrapperState

import config_scitos

def calc_Pose(x, y, yaw):
    quat = tf.transformations.quaternion_from_euler(0, 0, yaw)
    orientation = Quaternion(*quat)
    position = Point(x, y, 0)
    return Pose(position, orientation)

class Runner(object):
    """
    self.memory: memory to be used for conditions and actions
    self.worldstate: the default/start worldstate
    self.actionbag: the actions this runner uses
    self.planner: the planner this runner uses
    """

    def __init__(self, config_module=None):
        """
        param:config_module: a scenario/robot specific module to prepare setup,
                that has the following members:
                    get_all_conditions() -> return a list of conditions
                    get_all_actions() -> return a list of actions
        """
        self.memory = Memory()
        self.worldstate = WorldState()
        self.actionbag = ActionBag()

        if config_module is not None:
            for condition in config_module.get_all_conditions(self.memory):
                Condition.add(condition)
            for action in config_module.get_all_actions(self.memory):
                self.actionbag.add(action)

        self.planner = Planner(self.actionbag, self.worldstate, None)

        self._introspector = None

    def __repr__(self):
        return '<%s memory=%s worldstate=%s actions=%s planner=%s>' % (self.__class__.__name__,
                                self.memory, self.worldstate, self.actionbag, self.planner)

    def _setup_introspection(self):
        # init what could have been initialized externally
        if not rospy.core.is_initialized():
            rospy.init_node('goap_runner_introspector')
        # init everything else but only once
        if self._introspector is None:
            self._introspector = Introspector()
            thread.start_new_thread(rospy.spin, ())
            print "introspection spinner started"

    def update_and_plan(self, goal, tries=1, introspection=False):
        """introspection: introspect GOAP planning via smach.introspection"""
        # update to reality
        Condition.initialize_worldstate(self.worldstate)

        print "worldstate initialized/updated to: ", self.worldstate
        for (condition, value) in self.worldstate._condition_values.iteritems():
            if value is None:
                rospy.logwarn("Condition still 'None': %s", condition)

        if introspection:
            self._setup_introspection()

        while tries > 0:
            tries -= 1
            start_node = self.planner.plan(goal=goal)
            if start_node is not None:
                break

        if introspection:
            if start_node is not None:
                self._introspector.publish(start_node)
            self._introspector.publish_net(start_node,
                           self.planner.last_goal_node)

        return start_node

    def update_and_plan_and_execute(self, goal, tries=1, introspection=False):
        """introspection: introspect GOAP planning and SMACH execution via
        smach.introspection"""
        start_node = self.update_and_plan(goal, tries, introspection)
        if start_node is not None:
            #PlanExecutor().execute(start_node)
            return self.execute_as_smach(start_node, introspection)
        else:
            return 'aborted'

    def execute_as_smach(self, start_node, introspection=False):
        sm = self.path_to_smach(start_node)
        outcome = execute_smach_container(sm, introspection, name='/SM_GENERATED')
        return outcome

    def path_to_smach(self, start_node):
        sm = StateMachine(outcomes=['succeeded', 'aborted', 'preempted'])

        node = start_node
        with sm:
            while len(node.parent_nodes_path_list) > 0: # skipping the goal node at the end
                next_node = node.parent_nodes_path_list[-1]

                if isinstance(node.action, SmachStateAction):
                    StateMachine.add_auto('%s_%X' % (node.action.__class__.__name__, id(node)),
                                          node.action.state,
                                          ['succeeded'],
                                          remapping=node.action.get_remapping())
                    node.action.translate_worldstate_to_userdata(next_node.worldstate, sm.userdata)
                else:
                    StateMachine.add_auto('%s_%X' % (node.action.__class__.__name__, id(node)),
                                          GOAPActionWrapperState(node.action, next_node.worldstate),
                                          ['succeeded'])

                node = next_node

        return sm

    def print_worldstate_loop(self):
        rate = rospy.Rate(0.5)
        while not rospy.is_shutdown():
            # update to reality
            Condition.initialize_worldstate(self.worldstate)
            print self.worldstate
            rate.sleep()

#class MoveBaseGoalListener(object):
#
#    def __init__(self, topic, msg_type):
#        self._subscriber = rospy.Subscriber(topic, msg_type, self._callback, queue_size=1)

class MoveState(State):

    def __init__(self):
        State.__init__(self, outcomes=['succeeded', 'aborted', 'preempted'], input_keys=['x', 'y', 'yaw'], output_keys=['user_input'])
        self.runner = Runner(config_scitos)

    def execute(self, ud):
        pose = calc_Pose(ud.x, ud.y, ud.yaw)
        goal = Goal([Precondition(Condition.get('robot.pose'), pose)])
        # avoid duplicate introspection:
        outcome = self.runner.update_and_plan_and_execute(goal, introspection=False)
        print "GOAP returns: %s" % outcome
        return outcome

def test_runner():
    rospy.init_node('runner_test')

    sq = Sequence(outcomes=['succeeded', 'aborted', 'preempted'],
                  connector_outcome='succeeded')

    wfg = WaitForGoalState() # We don't want multiple subscribers so we need one WaitForGoal state

    with sq:

        Sequence.add('SLEEP', SleepState(5))

#        Sequence.add('CHECK', CheckSmachEnabledState(),
#                    transitions={'aborted':'SLEEP'})

        Sequence.add('WAIT_FOR_GOAL', wfg,
                     transitions={'aborted':'SLEEP'})

        Sequence.add('MOVE_GOAP', MoveState(),
                     transitions={'succeeded':'SLEEP'})

    execute_smach_container(sq, enable_introspection=True)

def test_tasker():

    rospy.init_node('tasker_test')

    wfg = WaitForGoalState() # We don't want multiple subscribers so we need one WaitForGoal state

    sq_move_to_new_goal = Sequence(outcomes=['succeeded', 'aborted', 'preempted'],
                                   connector_outcome='succeeded')
    with sq_move_to_new_goal:
        Sequence.add('WAIT_FOR_GOAL', wfg)
        Sequence.add('MOVE_GOAP', MoveState())

    sm_tasker = StateMachine(outcomes=['succeeded', 'aborted', 'preempted',
                                       'field_error', 'undefined_task'],
                             input_keys=['task_goal'])

    with sm_tasker:
        ## add all tasks to be available
        StateMachine.add('MOVE_TO_NEW_GOAL', sq_move_to_new_goal)

        StateMachine.add('LOOK_AROUND', get_lookaround_smach())

        StateMachine.add('GLIMPSE_AROUND', get_lookaround_smach(glimpse=True))

        StateMachine.add('MOVE_TO_RANDOM_GOAL', get_random_goal_smach())

        StateMachine.add('MOVE_TO_NEW_GOAL_AND_RETURN', task_go_and_return.get_go_and_return_smach())

        StateMachine.add('PATROL_TO_NEW_GOAL', task_patrol.get_patrol_smach())

        StateMachine.add('MOVE_AROUND', task_move_around.get_move_around_smach())

        StateMachine.add('SLEEP_FIVE_SEC', SleepState(5))

        ## now the task receiver is created and automatically links to
        ##   all task states added above
        task_states_labels = sm_tasker.get_children().keys()
        task_receiver_transitions = {'undefined_outcome':'undefined_task'}
        task_receiver_transitions.update({l:l for l in task_states_labels})

        StateMachine.add('TASK_RECEIVER',
                         UserDataToOutcomeState(task_states_labels,
                                                'task_goal',
                                                lambda ud: ud.task_id),
                         task_receiver_transitions)

    sm_tasker.set_initial_state(['TASK_RECEIVER'])

    rospy.loginfo('tasker starting, available tasks: %s', ', '.join(task_states_labels))
    pub = rospy.Publisher('/task/available_tasks', String, latch=True)
    thread.start_new_thread(rostopic.publish_message, (pub, String, [', '.join(task_states_labels)], 1))

    asw = ActionServerWrapper('activate_task', TaskActivationAction,
                              wrapped_container=sm_tasker,
                              succeeded_outcomes=['succeeded'],
                              aborted_outcomes=['aborted', 'undefined_task'],
                              preempted_outcomes=['preempted'],
                              goal_key='task_goal'
                              )

    # Create and start the introspection server
    sis = IntrospectionServer('smach_tasker_action', sm_tasker, '/SM_ROOT')
    sis.start()

    asw.run_server()

    rospy.spin()
    sis.stop()

if __name__ == '__main__':

    #test_runner()
    test_tasker()
<commit_msg>Add 'crazy' mode to lookaround smach<commit_after>#!/usr/bin/env python

# Copyright (c) 2013, Felix Kolbe
# All rights reserved. BSD License
#
# 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 the {organization} 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 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.

import roslib; roslib.load_manifest('goap')

import thread

import rospy
import rostopic
import tf

from smach import Sequence, State, StateMachine
from smach_ros import ActionServerWrapper, IntrospectionServer

from std_msgs.msg import String
from geometry_msgs.msg import Pose, Point, Quaternion
from task_msgs.msg import TaskActivationAction, TaskActivationGoal

from uashh_smach.util import CheckSmachEnabledState, TopicToOutcomeState, UserDataToOutcomeState, SleepState, execute_smach_container
from uashh_smach.platform.move_base import WaitForGoalState, get_random_goal_smach
from uashh_smach.manipulator.look_around import get_lookaround_smach
from uashh_smach.tasks import task_go_and_return, task_move_around, task_patrol

from common import ActionBag, Condition, Goal, Precondition, WorldState
from inheriting import Memory
from planning import Planner, PlanExecutor
from introspection import Introspector
from smach_bridge import SmachStateAction, GOAPActionWrapperState

import config_scitos

def calc_Pose(x, y, yaw):
    quat = tf.transformations.quaternion_from_euler(0, 0, yaw)
    orientation = Quaternion(*quat)
    position = Point(x, y, 0)
    return Pose(position, orientation)

class Runner(object):
    """
    self.memory: memory to be used for conditions and actions
    self.worldstate: the default/start worldstate
    self.actionbag: the actions this runner uses
    self.planner: the planner this runner uses
    """

    def __init__(self, config_module=None):
        """
        param:config_module: a scenario/robot specific module to prepare setup,
                that has the following members:
                    get_all_conditions() -> return a list of conditions
                    get_all_actions() -> return a list of actions
        """
        self.memory = Memory()
        self.worldstate = WorldState()
        self.actionbag = ActionBag()

        if config_module is not None:
            for condition in config_module.get_all_conditions(self.memory):
                Condition.add(condition)
            for action in config_module.get_all_actions(self.memory):
                self.actionbag.add(action)

        self.planner = Planner(self.actionbag, self.worldstate, None)

        self._introspector = None

    def __repr__(self):
        return '<%s memory=%s worldstate=%s actions=%s planner=%s>' % (self.__class__.__name__,
                                self.memory, self.worldstate, self.actionbag, self.planner)

    def _setup_introspection(self):
        # init what could have been initialized externally
        if not rospy.core.is_initialized():
            rospy.init_node('goap_runner_introspector')
        # init everything else but only once
        if self._introspector is None:
            self._introspector = Introspector()
            thread.start_new_thread(rospy.spin, ())
            print "introspection spinner started"

    def update_and_plan(self, goal, tries=1, introspection=False):
        """introspection: introspect GOAP planning via smach.introspection"""
        # update to reality
        Condition.initialize_worldstate(self.worldstate)

        print "worldstate initialized/updated to: ", self.worldstate
        for (condition, value) in self.worldstate._condition_values.iteritems():
            if value is None:
                rospy.logwarn("Condition still 'None': %s", condition)

        if introspection:
            self._setup_introspection()

        while tries > 0:
            tries -= 1
            start_node = self.planner.plan(goal=goal)
            if start_node is not None:
                break

        if introspection:
            if start_node is not None:
                self._introspector.publish(start_node)
            self._introspector.publish_net(start_node,
                           self.planner.last_goal_node)

        return start_node

    def update_and_plan_and_execute(self, goal, tries=1, introspection=False):
        """introspection: introspect GOAP planning and SMACH execution via
        smach.introspection"""
        start_node = self.update_and_plan(goal, tries, introspection)
        if start_node is not None:
            #PlanExecutor().execute(start_node)
            return self.execute_as_smach(start_node, introspection)
        else:
            return 'aborted'

    def execute_as_smach(self, start_node, introspection=False):
        sm = self.path_to_smach(start_node)
        outcome = execute_smach_container(sm, introspection, name='/SM_GENERATED')
        return outcome

    def path_to_smach(self, start_node):
        sm = StateMachine(outcomes=['succeeded', 'aborted', 'preempted'])

        node = start_node
        with sm:
            while len(node.parent_nodes_path_list) > 0: # skipping the goal node at the end
                next_node = node.parent_nodes_path_list[-1]

                if isinstance(node.action, SmachStateAction):
                    StateMachine.add_auto('%s_%X' % (node.action.__class__.__name__, id(node)),
                                          node.action.state,
                                          ['succeeded'],
                                          remapping=node.action.get_remapping())
                    node.action.translate_worldstate_to_userdata(next_node.worldstate, sm.userdata)
                else:
                    StateMachine.add_auto('%s_%X' % (node.action.__class__.__name__, id(node)),
                                          GOAPActionWrapperState(node.action, next_node.worldstate),
                                          ['succeeded'])

                node = next_node

        return sm

    def print_worldstate_loop(self):
        rate = rospy.Rate(0.5)
        while not rospy.is_shutdown():
            # update to reality
            Condition.initialize_worldstate(self.worldstate)
            print self.worldstate
            rate.sleep()

#class MoveBaseGoalListener(object):
#
#    def __init__(self, topic, msg_type):
#        self._subscriber = rospy.Subscriber(topic, msg_type, self._callback, queue_size=1)

class MoveState(State):

    def __init__(self):
        State.__init__(self, outcomes=['succeeded', 'aborted', 'preempted'], input_keys=['x', 'y', 'yaw'], output_keys=['user_input'])
        self.runner = Runner(config_scitos)

    def execute(self, ud):
        pose = calc_Pose(ud.x, ud.y, ud.yaw)
        goal = Goal([Precondition(Condition.get('robot.pose'), pose)])
        # avoid duplicate introspection:
        outcome = self.runner.update_and_plan_and_execute(goal, introspection=False)
        print "GOAP returns: %s" % outcome
        return outcome

def test_runner():
    rospy.init_node('runner_test')

    sq = Sequence(outcomes=['succeeded', 'aborted', 'preempted'],
                  connector_outcome='succeeded')

    wfg = WaitForGoalState() # We don't want multiple subscribers so we need one WaitForGoal state

    with sq:

        Sequence.add('SLEEP', SleepState(5))

#        Sequence.add('CHECK', CheckSmachEnabledState(),
#                    transitions={'aborted':'SLEEP'})

        Sequence.add('WAIT_FOR_GOAL', wfg,
                     transitions={'aborted':'SLEEP'})

        Sequence.add('MOVE_GOAP', MoveState(),
                     transitions={'succeeded':'SLEEP'})

    execute_smach_container(sq, enable_introspection=True)

def test_tasker():

    rospy.init_node('tasker_test')

    wfg = WaitForGoalState() # We don't want multiple subscribers so we need one WaitForGoal state

    sq_move_to_new_goal = Sequence(outcomes=['succeeded', 'aborted', 'preempted'],
                                   connector_outcome='succeeded')
    with sq_move_to_new_goal:
        Sequence.add('WAIT_FOR_GOAL', wfg)
        Sequence.add('MOVE_GOAP', MoveState())

    sm_tasker = StateMachine(outcomes=['succeeded', 'aborted', 'preempted',
                                       'field_error', 'undefined_task'],
                             input_keys=['task_goal'])

    with sm_tasker:
        ## add all tasks to be available
        StateMachine.add('MOVE_TO_NEW_GOAL', sq_move_to_new_goal)

        StateMachine.add('LOOK_AROUND', get_lookaround_smach())

        StateMachine.add('GLIMPSE_AROUND', get_lookaround_smach(glimpse=True))

        StateMachine.add('MOVE_ARM_CRAZY', get_lookaround_smach(crazy=True))

        StateMachine.add('MOVE_TO_RANDOM_GOAL', get_random_goal_smach())

        StateMachine.add('MOVE_TO_NEW_GOAL_AND_RETURN', task_go_and_return.get_go_and_return_smach())

        StateMachine.add('PATROL_TO_NEW_GOAL', task_patrol.get_patrol_smach())

        StateMachine.add('MOVE_AROUND', task_move_around.get_move_around_smach())

        StateMachine.add('SLEEP_FIVE_SEC', SleepState(5))

        ## now the task receiver is created and automatically links to
        ##   all task states added above
        task_states_labels = sm_tasker.get_children().keys()
        task_receiver_transitions = {'undefined_outcome':'undefined_task'}
        task_receiver_transitions.update({l:l for l in task_states_labels})

        StateMachine.add('TASK_RECEIVER',
                         UserDataToOutcomeState(task_states_labels,
                                                'task_goal',
                                                lambda ud: ud.task_id),
                         task_receiver_transitions)

    sm_tasker.set_initial_state(['TASK_RECEIVER'])

    rospy.loginfo('tasker starting, available tasks: %s', ', '.join(task_states_labels))
    pub = rospy.Publisher('/task/available_tasks', String, latch=True)
    thread.start_new_thread(rostopic.publish_message, (pub, String, [', '.join(task_states_labels)], 1))

    asw = ActionServerWrapper('activate_task', TaskActivationAction,
                              wrapped_container=sm_tasker,
                              succeeded_outcomes=['succeeded'],
                              aborted_outcomes=['aborted', 'undefined_task'],
                              preempted_outcomes=['preempted'],
                              goal_key='task_goal'
                              )

    # Create and start the introspection server
    sis = IntrospectionServer('smach_tasker_action', sm_tasker, '/SM_ROOT')
    sis.start()

    asw.run_server()

    rospy.spin()
    sis.stop()

if __name__ == '__main__':

    #test_runner()
    test_tasker()
<|endoftext|>


<commit_before>"""
Show transit information for normal user. (Not a contributor)
"""
from django.http import HttpResponse
from django.http import Http404
from django.template import Context, loader , RequestContext
from django.conf import settings
from django.utils.translation import ugettext_lazy as _
from django import forms
from django.forms import ModelForm

from ragendja.auth.decorators import staff_only
from ragendja.template import render_to_response
from ragendja.dbutils import get_object_or_404

from gogogo.models import *
from gogogo.geo.LatLng import LatLng
from widgets import Pathbar

def index(request):
	"""
	Show transit information
	"""

	pathbar = Pathbar()
	pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)

	query = Agency.all()

	agency_list = []
	for row in query:
		entity = create_entity(row,request)
		entity['key_name'] = row.key().name()
		agency_list.append(entity)

	return render_to_response(
		request,
		'gogogo/transit/index.html'
		,{
			'page_title': _("Transit information"),
			'pathbar' : pathbar,
		   "agency_list" : agency_list,
		   })

def agency(request,agency_id):
	"""
	Browse the information of a transport agency
	"""

	try:
		key = db.Key.from_path(Agency.kind(),agency_id)

		record = db.get(key)
	except (db.BadArgumentError,db.BadValueError):
		raise Http404

	if record == None:
		raise Http404

	agency = create_entity(record,request)
	agency['key_name'] = record.key().name()

	gql = db.GqlQuery("SELECT * FROM gogogo_route where type = :1 and agency=:2",2,record)
	rail_list = []
	for row in gql:
		entity = create_entity(row,request)
		entity['key_name'] = row.key().name()
		rail_list.append(entity)

	pathbar = Pathbar()
	pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	pathbar.append(agency['name'] , 'gogogo.views.transit.agency' , [agency['key_name']])

	t = loader.get_template('gogogo/transit/agency.html')
	c = RequestContext(
		request,
	{
        'page_title': agency['name'] ,
        'pathbar' : pathbar,
        'agency' : agency,
        'rail_list' : rail_list
    })

	return HttpResponse(t.render(c))

def route(request,agency_id,route_id):
	"""
	Browse the information of a route data.

	It is the most expensive call in the system.
	"""

	agency_record = get_object_or_404(Agency,key_name=agency_id)
	record = get_object_or_404(Route,key_name=route_id)

	lang = MLStringProperty.get_current_lang(request)

	agency_entity = {
		'key_name': agency_record.key().name(),
		'name' : MLStringProperty.trans(agency_record.name,lang),
	}

	route_entity = {
		'key_name' : record.key().name(),
		'long_name' : MLStringProperty.trans(record.long_name,lang),
		'short_name' : record.short_name,
		'desc' : record.desc,

		# Reduce the no. of database access
		'agency' : agency_id
		#'agency' : record.agency.key().name()
	}

	gql = db.GqlQuery("SELECT * from gogogo_trip where route = :1",record)

	trip_list = []
	travel_list = []

	# Endpoint in all trip
	endpoint_list = []

	for trip_record in gql:
		stop_list = []
		try:
			first_stop = db.get(trip_record.stop_list[0])
			last_stop = db.get(trip_record.stop_list[len(trip_record.stop_list)-1])
			travel_list.append( (MLStringProperty.trans(first_stop.name,lang),MLStringProperty.trans(last_stop.name,lang)))

			pt = LatLng(first_stop.latlng.lat,first_stop.latlng.lon )
			if pt not in endpoint_list:
				endpoint_list.append( pt )

			pt = LatLng(last_stop.latlng.lat,last_stop.latlng.lon )
			if pt not in endpoint_list:
				endpoint_list.append( pt )
		except IndexError:
			pass

		for key in trip_record.stop_list:
			stop = db.get(key)
			stop_list.append( (key.name() , MLStringProperty.trans(stop.name,lang) ) )
		trip_entity = {
			'headsign' : MLStringProperty.trans(trip_record.headsign,lang),
			'key_name' : trip_record.key().name(),
			'stop_list' : stop_list
		}
		trip_list.append(trip_entity)

	pathbar = Pathbar()
	pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	pathbar.append(agency_entity['name'] , 'gogogo.views.transit.agency' , [agency_entity['key_name']])
	pathbar.append(route_entity['long_name'] , 'gogogo.views.transit.route' , [agency_entity['key_name'] , route_entity['key_name']])

	return render_to_response(
		request,
		'gogogo/transit/route.html'
		,{
			'page_title': route_entity['long_name'],
			'pathbar' : pathbar,
		   "route" : route_entity,
		   "trip_list" : trip_list,
		   "travel_list" : travel_list,
		   "endpoint_list" : endpoint_list,
		   })

def trip(request,agency_id,route_id,trip_id):
	"""
	Browse the information of a trip
	"""

	trip_record = get_object_or_404(Trip,key_name=trip_id)
	agency_record = get_object_or_404(Agency,key_name=agency_id)
	route_record = get_object_or_404(Route,key_name=route_id)

	lang = MLStringProperty.get_current_lang(request)

	trip_entity = {
		'key_name' : trip_record.key().name(),
		'short_name' : MLStringProperty.trans(trip_record.short_name,lang),
		'headsign' : MLStringProperty.trans(trip_record.headsign,lang),
	}

	agency_entity = {
		'key_name': agency_record.key().name(),
		'name' : MLStringProperty.trans(agency_record.name,lang),
	}

	route_entity = {
		'key_name' : route_record.key().name(),
		'long_name' : MLStringProperty.trans(route_record.long_name,lang),
	}

	shape_entity = None
	if trip_record.shape:
		shape_entity = {'points' : trip_record.shape.points,
			'color': trip_record.shape.color	}

	pathbar = Pathbar()
	pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	pathbar.append(agency_entity['name'] , 'gogogo.views.transit.agency' , [agency_entity['key_name']])
	pathbar.append(route_entity['long_name'] , 'gogogo.views.transit.route' , [agency_entity['key_name'] , route_entity['key_name']])
	pathbar.append(trip_entity['headsign'] , 'gogogo.views.transit.trip' , [agency_entity['key_name'] , route_entity['key_name'],trip_entity['key_name']])

	return render_to_response(
		request,
		'gogogo/transit/trip.html'
		,{
			'page_title': trip_entity['headsign'],
			'pathbar': pathbar,
			"agency" : agency_entity,
			"route" : route_entity,
		   "trip" : trip_entity,
		   "shape" : shape_entity
		   })

def stop(request,stop_id):
	"""
	Browse stop information
	"""
	try:
		key = db.Key.from_path(Stop.kind(),stop_id)

		record = db.get(key)
	except (db.BadArgumentError,db.BadValueError):
		raise Http404

	lang = MLStringProperty.get_current_lang(request)
	stop_entity = {
		'key_name' : record.key().name(),
		'name' : MLStringProperty.trans(record.name,lang),
		'desc' : MLStringProperty.trans(record.desc,lang),
		'url' : record.url
	}

	pathbar = Pathbar()
	pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	pathbar.append(stop_entity['name'] , 'gogogo.views.transit.stop', [ stop_entity['key_name']]  )

	return render_to_response(
		request,
		'gogogo/transit/stop.html'
		,{
			'page_title': _("Transit information"),
			'pathbar' : pathbar,
		   "stop" : stop_entity,
		   })

<commit_msg>Added gogogo.views.transit.Pathbar - A Pathbar designed for transit information<commit_after>"""
Show transit information for normal user. (Not a contributor)
"""
from django.http import HttpResponse
from django.http import Http404
from django.template import Context, loader , RequestContext
from django.conf import settings
from django.utils.translation import ugettext_lazy as _
from django import forms
from django.forms import ModelForm

from ragendja.auth.decorators import staff_only
from ragendja.template import render_to_response
from ragendja.dbutils import get_object_or_404

from gogogo.models import *
from gogogo.geo.LatLng import LatLng
from widgets import Pathbar as _Pathbar
#import gogogo.views.widgets.Pathbar

class Pathbar(_Pathbar):
	"""
	A Pathbar for transit information
	"""
	def __init__(self,agency=None , stop = None):
		_Pathbar.__init__(self)

		self.append(_("Transit information") , 'gogogo.views.transit.index',None)

		if agency:
			try:
				self.append(agency[0]['name'] ,
					'gogogo.views.transit.agency' , [agency[0]['key_name']])
				self.append(agency[1]['long_name'] ,
					'gogogo.views.transit.route' ,
						[agency[0]['key_name'] , agency[1]['key_name']])
				self.append(agency[2]['headsign'] ,
					'gogogo.views.transit.trip' ,
						[agency[0]['key_name'] , agency[1]['key_name'],agency[2]['key_name']])
			except IndexError:
				pass

		elif stop:
			self.append(stop['name'] , 'gogogo.views.transit.stop', [ stop['key_name']]  )

def index(request):
	"""
	Show transit information
	"""

	pathbar = Pathbar()
	#pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)

	query = Agency.all()

	agency_list = []
	for row in query:
		entity = create_entity(row,request)
		entity['key_name'] = row.key().name()
		agency_list.append(entity)

	return render_to_response(
		request,
		'gogogo/transit/index.html'
		,{
			'page_title': _("Transit information"),
			'pathbar' : pathbar,
		   "agency_list" : agency_list,
		   })

def agency(request,agency_id):
	"""
	Browse the information of a transport agency
	"""

	try:
		key = db.Key.from_path(Agency.kind(),agency_id)

		record = db.get(key)
	except (db.BadArgumentError,db.BadValueError):
		raise Http404

	if record == None:
		raise Http404

	agency = create_entity(record,request)
	agency['key_name'] = record.key().name()

	gql = db.GqlQuery("SELECT * FROM gogogo_route where type = :1 and agency=:2",2,record)
	rail_list = []
	for row in gql:
		entity = create_entity(row,request)
		entity['key_name'] = row.key().name()
		rail_list.append(entity)

	pathbar = Pathbar(agency=(agency,))
	#pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	#pathbar.append(agency['name'] , 'gogogo.views.transit.agency' , [agency['key_name']])

	t = loader.get_template('gogogo/transit/agency.html')
	c = RequestContext(
		request,
	{
        'page_title': agency['name'] ,
        'pathbar' : pathbar,
        'agency' : agency,
        'rail_list' : rail_list
    })

	return HttpResponse(t.render(c))

def route(request,agency_id,route_id):
	"""
	Browse the information of a route data.

	It is the most expensive call in the system.
	"""

	agency_record = get_object_or_404(Agency,key_name=agency_id)
	record = get_object_or_404(Route,key_name=route_id)

	lang = MLStringProperty.get_current_lang(request)

	agency_entity = {
		'key_name': agency_record.key().name(),
		'name' : MLStringProperty.trans(agency_record.name,lang),
	}

	route_entity = {
		'key_name' : record.key().name(),
		'long_name' : MLStringProperty.trans(record.long_name,lang),
		'short_name' : record.short_name,
		'desc' : record.desc,

		# Reduce the no. of database access
		'agency' : agency_id
		#'agency' : record.agency.key().name()
	}

	gql = db.GqlQuery("SELECT * from gogogo_trip where route = :1",record)

	trip_list = []
	travel_list = []

	# Endpoint in all trip
	endpoint_list = []

	for trip_record in gql:
		stop_list = []
		try:
			first_stop = db.get(trip_record.stop_list[0])
			last_stop = db.get(trip_record.stop_list[len(trip_record.stop_list)-1])
			travel_list.append( (MLStringProperty.trans(first_stop.name,lang),MLStringProperty.trans(last_stop.name,lang)))

			pt = LatLng(first_stop.latlng.lat,first_stop.latlng.lon )
			if pt not in endpoint_list:
				endpoint_list.append( pt )

			pt = LatLng(last_stop.latlng.lat,last_stop.latlng.lon )
			if pt not in endpoint_list:
				endpoint_list.append( pt )
		except IndexError:
			pass

		for key in trip_record.stop_list:
			stop = db.get(key)
			stop_list.append( (key.name() , MLStringProperty.trans(stop.name,lang) ) )
		trip_entity = {
			'headsign' : MLStringProperty.trans(trip_record.headsign,lang),
			'key_name' : trip_record.key().name(),
			'stop_list' : stop_list
		}
		trip_list.append(trip_entity)

	pathbar = Pathbar(agency=(agency_entity,route_entity))
	#pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	#pathbar.append(agency_entity['name'] , 'gogogo.views.transit.agency' , [agency_entity['key_name']])
	#pathbar.append(route_entity['long_name'] , 'gogogo.views.transit.route' , [agency_entity['key_name'] , route_entity['key_name']])

	return render_to_response(
		request,
		'gogogo/transit/route.html'
		,{
			'page_title': route_entity['long_name'],
			'pathbar' : pathbar,
		   "route" : route_entity,
		   "trip_list" : trip_list,
		   "travel_list" : travel_list,
		   "endpoint_list" : endpoint_list,
		   })

def trip(request,agency_id,route_id,trip_id):
	"""
	Browse the information of a trip
	"""

	trip_record = get_object_or_404(Trip,key_name=trip_id)
	agency_record = get_object_or_404(Agency,key_name=agency_id)
	route_record = get_object_or_404(Route,key_name=route_id)

	lang = MLStringProperty.get_current_lang(request)

	trip_entity = {
		'key_name' : trip_record.key().name(),
		'short_name' : MLStringProperty.trans(trip_record.short_name,lang),
		'headsign' : MLStringProperty.trans(trip_record.headsign,lang),
	}

	agency_entity = {
		'key_name': agency_record.key().name(),
		'name' : MLStringProperty.trans(agency_record.name,lang),
	}

	route_entity = {
		'key_name' : route_record.key().name(),
		'long_name' : MLStringProperty.trans(route_record.long_name,lang),
	}

	shape_entity = None
	if trip_record.shape:
		shape_entity = {'points' : trip_record.shape.points,
			'color': trip_record.shape.color	}

	pathbar = Pathbar(agency=(agency_entity,route_entity,trip_entity))
	#pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	#pathbar.append(agency_entity['name'] , 'gogogo.views.transit.agency' , [agency_entity['key_name']])
	#pathbar.append(route_entity['long_name'] , 'gogogo.views.transit.route' , [agency_entity['key_name'] , route_entity['key_name']])
	#pathbar.append(trip_entity['headsign'] , 'gogogo.views.transit.trip' , [agency_entity['key_name'] , route_entity['key_name'],trip_entity['key_name']])

	return render_to_response(
		request,
		'gogogo/transit/trip.html'
		,{
			'page_title': trip_entity['headsign'],
			'pathbar': pathbar,
			"agency" : agency_entity,
			"route" : route_entity,
		   "trip" : trip_entity,
		   "shape" : shape_entity
		   })

def stop(request,stop_id):
	"""
	Browse stop information
	"""
	try:
		key = db.Key.from_path(Stop.kind(),stop_id)

		record = db.get(key)
	except (db.BadArgumentError,db.BadValueError):
		raise Http404

	lang = MLStringProperty.get_current_lang(request)
	stop_entity = {
		'key_name' : record.key().name(),
		'name' : MLStringProperty.trans(record.name,lang),
		'desc' : MLStringProperty.trans(record.desc,lang),
		'url' : record.url
	}

	pathbar = Pathbar(stop=stop_entity)
	#pathbar.append(_("Transit information") , 'gogogo.views.transit.index',None)
	#pathbar.append(stop_entity['name'] , 'gogogo.views.transit.stop', [ stop_entity['key_name']]  )

	return render_to_response(
		request,
		'gogogo/transit/stop.html'
		,{
			'page_title': _("Transit information"),
			'pathbar' : pathbar,
		   "stop" : stop_entity,
		   })

<|endoftext|>