<commit_before>import abc

class RODataSet(object):
    """
    An :py:mod:`abc` that represents a read-only set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def __iter__(self):
        """
        Iterate over all the :py:class:`~.importtools.importables.Importable`
        instances in this set.

        """

    @abc.abstractmethod
    def __contains__(self, importable):
        """
        Check if the dataset contains an equal
        :py:class:`~.importtools.importables.Importable` instance.

        See :py:meth:`~.importtools.importables.Importable.__cmp__` method for
        more details about the equality test between two
        :py:class:`~.importtools.importables.Importable` instances.

        """

    @abc.abstractmethod
    def get(self, importable, default=None):
        """
        Return an equal :py:class:`~.importtools.importables.Importable` from
        the dataset or the default value if no such instances are found.

        """

class DataSet(RODataSet):
    """
    An :py:mod:`abc` that represents a mutable set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def add(self, importable):
        """
        Add the :py:class:`~.importtools.importables.Importable` in the dataset
        or replaces an existing and equal one.

        """

    @abc.abstractmethod
    def pop(self, importable):
        """
        Remove and return an equal
        :py:class:`~.importtools.importables.Importable` from the dataset.

        """

    @abc.abstractmethod
    def __len__(self):
        """``DataSet`` length."""

class MemoryDataSet(dict, DataSet):
    """
    A simple in-memory :py:class:`dict`-based :py:class:`DataSet`
    implementation.

    Since this particular :py:class:`DataSet` implementation only uses the
    :py:meth:`__hash__` and :py:meth:`__cmp__` methods of an
    :py:class:`~.importtools.importables.Importable` :py:mod:`abc` we can
    exemplify this class functionality using :py:class:`tuple` objects:

    .. testsetup::

    >>> from importtools.datasets import MemoryDataSet
    >>> item1 = (1, )
    >>> item2 = (1, )
    >>> item3 = (2, )

    At first, a newly created :py:class:`MemoryDataSet` instance has no
    elements:

    >>> mds = MemoryDataSet()
    >>> list(mds)
    []

    After creation, it can be populated and the elements in the dataset can be
    retrieved using other equal elements.

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item1) == item1
    True
    >>> mds.get(item1) is item1
    True
    >>> mds.get(item2) == item2
    True
    >>> mds.get(item2) is item1
    True

    Trying to get an inexistent item should return the default value or
    :py:class:`None`:

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item3) is None
    True
    >>> mds.get(item3, 'default')
    'default'

    Initial data can be passed when constructing instances:

    >>> mds = MemoryDataSet((item1, item2, item3))

    """
    def __init__(self, data_loader=None):
        data_iter = tuple()
        if data_loader is not None:
            data_iter = iter(data_loader)
        mapping = ((i, i) for i in data_iter)
        super(MemoryDataSet, self).__init__(mapping)

    def add(self, importable):
        self[importable] = importable

class DiffDataSet(DataSet):
    """
    A :py:class:`DataSet` wrapper that remembers all the changes, additions and
    removals done to the wrapped dataset trough itself.

    Using a :py:class:`DiffDataSet` as the destination of the import algorithm
    allows optimal persistence of the changes by grouping them in a way suited
    for batch processing.

    In order to exemplify this class functionality we start with a
    :py:class:`MemoryDataSet` instance that we wrap:

    .. testsetup::

    >>> from importtools.datasets import MemoryDataSet, DiffDataSet
    >>> mds = MemoryDataSet()
    >>> dds = DiffDataSet(mds)

    An :py:class:`~.importtools.importables.Importable` mock is needed to
    populate the datasets:

    .. testsetup::

    >>> from importtools.importables import Importable
    >>> class MockImportable(Importable):
    ...     def __init__(self, *args):
    ...         super(MockImportable, self).__init__()
    ...         self.args = tuple(args)
    ...         self.make_imported()
    ...     def __hash__(self):
    ...         return hash(self.args)
    ...     def __cmp__(self, other):
    ...         return cmp(self.args, other.args)
    ...     def __repr__(self):
    ...         smap = {1: 'IMPORTED', 2: 'FORCED', 3:'INVALID'}
    ...         return '<%r %s>' % (self.args, smap.get(self._status, 'N/A'))

    It's now possible to populate the wrapped :py:class:`MemoryDataSet`
    instance with some elements:

    >>> mds.add(MockImportable(0))
    >>> mds.add(MockImportable(1))
    >>> mds.add(MockImportable(2))
    >>> mds.add(MockImportable(3))
    >>> sorted(mds, key=lambda x: x.args)
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    The contents of ``dds`` should be the same as the content of the wrapped
    ``mds``:

    >>> sorted(dds, key=lambda x: x.args)
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    Fetching existing elements from ``dds`` should work the same as fetching
    elements from ``mds``:

    >>> mds.get(MockImportable(1)) == dds.get(MockImportable(1))
    True

    New :py:class:`~.importtools.importables.Importable` can still added in
    ``dds``:

    >>> dds.add(MockImportable(4))
    >>> dds.add(MockImportable(5))
    >>> sorted(dds, key=lambda x: x.args)
    ... # doctest: +NORMALIZE_WHITESPACE
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>,
     <(3,) IMPORTED>, <(4,) IMPORTED>, <(5,) IMPORTED>]

    Check if the new elements were added:

    >>> dds.get(MockImportable(4))
    <(4,) IMPORTED>
    >>> dds.get(MockImportable(5))
    <(5,) IMPORTED>

    A lookup with an inexistent element should return the default value:

    >>> dds.get(MockImportable(100), 'default')
    'default'

    We should be able to ``pop`` both original and recently added elements:

    >>> dds.pop(MockImportable(1))
    <(1,) IMPORTED>
    >>> dds.pop(MockImportable(4))
    <(4,) IMPORTED>
    >>> sorted(dds, key=lambda x: x.args)
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>, <(5,) IMPORTED>]

    And everything should have been registered:

    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    []

    Adding and removing the same
    :py:class:`~.importtools.importables.Importable` shouldn't change anything:

    >>> dds.add(MockImportable(100))
    >>> dds.pop(MockImportable(100))
    <(100,) IMPORTED>
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    []

    But removing an original :py:class:`~.importtools.importables.Importable`
    and adding a new one should be interpreted as a change:

    >>> dds.pop(MockImportable(2))
    <(2,) IMPORTED>
    >>> dds.add(MockImportable(2))
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(2,) IMPORTED>]

    If you change one of the existing elements it will be marked as changed:

    >>> dds.get(MockImportable(3)).register_change()
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(2,) IMPORTED>, <(3,) IMPORTED>]

    Even if the elements are accessed by iterating on over the dataset the
    changes are still tracked:

    >>> list(dds)[0].register_change()
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    When printing the object a list of all the added, removed and created
    elements will be shown, along with the total number for each operation:

    >>> print dds
    To be added: 1
    <(5,) IMPORTED>
    To be removed: 1
    <(1,) IMPORTED>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._added = MemoryDataSet()
    >>> dds._removed = MemoryDataSet()
    >>> print dds
    <BLANKLINE>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._changed = MemoryDataSet()
    >>> print dds
    No elements were affected.

    """
    def __init__(self, data_set):
        self._wrapped = data_set
        self._added = MemoryDataSet()
        self._removed = MemoryDataSet()
        self._changed = MemoryDataSet()
        super(DiffDataSet, self).__init__()

    def __iter__(self):
        for importable in self._wrapped:
            yield importable.register_listener(self.register_change)

    def __contains__(self, importable):
        return importable in self._wrapped

    def __str__(self):
        result = ''

        added_count = len(self._added)
        if added_count > 0:
            result += 'To be added: %s\n' % added_count
            result += '\n'.join(sorted(repr(x) for x in self.get_added()))

        removed_count = len(self._removed)
        if removed_count > 0:
            result += '\nTo be removed: %s\n' % removed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_removed()))

        changed_count = len(self._changed)
        if changed_count > 0:
            result += '\nTo be changed: %s\n' % changed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_changed()))

        if (added_count + removed_count + changed_count) == 0:
            result += 'No elements were affected.'

        return result

    def get(self, importable, default=None):
        importable = self._wrapped.get(importable)
        if importable is None:
            return default
        return importable.register_listener(self.register_change)

    def add(self, importable):
        self._wrapped.add(importable)

        was_removed = self._removed.get(importable)
        if was_removed is None:
            self._added.add(importable)
        else:
            self._removed.pop(importable)
            self._changed.add(importable)

    def pop(self, importable):
        was_added = self._added.get(importable)
        if was_added is None:
            self._removed.add(importable)
        else:
            self._added.pop(importable)

        return self._wrapped.pop(importable)

    def register_change(self, importable):
        """Mark an element in the current dataset as changed.

        This method is registered as a listener for changes in all the
        elements of the wrapped :py:class:`DataSet`.

        """
        assert importable in self._wrapped
        self._changed.add(importable)

    def get_added(self):
        """An iterable of all added elements in the wrapped dataset."""
        return iter(self._added)

    def get_removed(self):
        """An iterable of all removed elements in the wrapped dataset."""
        return iter(self._removed)

    def get_changed(self):
        """An iterable of all changed elements in the wrapped dataset."""
        return iter(self._changed)
<commit_msg>Refactored DiffDataSet.<commit_after>import abc

class RODataSet(object):
    """
    An :py:mod:`abc` that represents a read-only set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def __iter__(self):
        """
        Iterate over all the :py:class:`~.importtools.importables.Importable`
        instances in this set.

        """

    @abc.abstractmethod
    def __contains__(self, importable):
        """
        Check if the dataset contains an equal
        :py:class:`~.importtools.importables.Importable` instance.

        See :py:meth:`~.importtools.importables.Importable.__cmp__` method for
        more details about the equality test between two
        :py:class:`~.importtools.importables.Importable` instances.

        """

    @abc.abstractmethod
    def get(self, importable, default=None):
        """
        Return an equal :py:class:`~.importtools.importables.Importable` from
        the dataset or the default value if no such instances are found.

        """

class DataSet(RODataSet):
    """
    An :py:mod:`abc` that represents a mutable set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def add(self, importable):
        """
        Add the :py:class:`~.importtools.importables.Importable` in the dataset
        or replaces an existing and equal one.

        """

    @abc.abstractmethod
    def pop(self, importable):
        """
        Remove and return an equal
        :py:class:`~.importtools.importables.Importable` from the dataset.

        """

class MemoryDataSet(dict, DataSet):
    """
    A simple in-memory :py:class:`dict`-based :py:class:`DataSet`
    implementation.

    Since this particular :py:class:`DataSet` implementation only uses the
    :py:meth:`__hash__` and :py:meth:`__cmp__` methods of an
    :py:class:`~.importtools.importables.Importable` :py:mod:`abc` we can
    exemplify this class functionality using :py:class:`tuple` objects:

    .. testsetup::

    >>> from importtools.datasets import MemoryDataSet
    >>> item1 = (1, )
    >>> item2 = (1, )
    >>> item3 = (2, )

    At first, a newly created :py:class:`MemoryDataSet` instance has no
    elements:

    >>> mds = MemoryDataSet()
    >>> list(mds)
    []

    After creation, it can be populated and the elements in the dataset can be
    retrieved using other equal elements.

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item1) == item1
    True
    >>> mds.get(item1) is item1
    True
    >>> mds.get(item2) == item2
    True
    >>> mds.get(item2) is item1
    True

    Trying to get an inexistent item should return the default value or
    :py:class:`None`:

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item3) is None
    True
    >>> mds.get(item3, 'default')
    'default'

    Initial data can be passed when constructing instances:

    >>> mds = MemoryDataSet((item1, item2, item3))

    """
    def __init__(self, data_loader=None):
        data_iter = tuple()
        if data_loader is not None:
            data_iter = iter(data_loader)
        mapping = ((i, i) for i in data_iter)
        super(MemoryDataSet, self).__init__(mapping)

    def add(self, importable):
        self[importable] = importable

class DiffDataSet(MemoryDataSet):
    """
    A :py:class:`DataSet` implementation that remembers all the changes,
    additions and removals done to it.

    Using a :py:class:`DiffDataSet` as the destination of the import algorithm
    allows optimal persistence of the changes by grouping them in a way suited
    for batch processing.

    In order to exemplify this class functionality we start with a
    instance that we prepopulate:

    An :py:class:`~.importtools.importables.Importable` mock is needed to
    populate the datasets:

    .. testsetup::

    >>> from importtools.importables import Importable
    >>> class MockImportable(Importable):
    ...     def __init__(self, *args):
    ...         super(MockImportable, self).__init__()
    ...         self.args = tuple(args)
    ...         self.make_imported()
    ...     def __hash__(self):
    ...         return hash(self.args)
    ...     def __cmp__(self, other):
    ...         return cmp(self.args, other.args)
    ...     def __repr__(self):
    ...         smap = {1: 'IMPORTED', 2: 'FORCED', 3:'INVALID'}
    ...         return '<%r %s>' % (self.args, smap.get(self._status, 'N/A'))

    .. testsetup::

    >>> from importtools.datasets import MemoryDataSet, DiffDataSet
    >>> dds = DiffDataSet([
    ...     MockImportable(0),
    ...     MockImportable(1),
    ...     MockImportable(2),
    ...     MockImportable(3)
    ... ])

    >>> sorted(dds, key=lambda x: x.args)
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    Fetching existing elements from ``dds`` should work correctly:

    >>> dds.get(MockImportable(1)) == MockImportable(1)
    True

    New :py:class:`~.importtools.importables.Importable` can still added in
    ``dds``:

    >>> dds.add(MockImportable(4))
    >>> dds.add(MockImportable(5))
    >>> sorted(dds, key=lambda x: x.args)
    ... # doctest: +NORMALIZE_WHITESPACE
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>,
     <(3,) IMPORTED>, <(4,) IMPORTED>, <(5,) IMPORTED>]

    Check if the new elements were added:

    >>> dds.get(MockImportable(4))
    <(4,) IMPORTED>
    >>> dds.get(MockImportable(5))
    <(5,) IMPORTED>

    A lookup with an inexistent element should return the default value:

    >>> dds.get(MockImportable(100), 'default')
    'default'

    We should be able to ``pop`` both original and recently added elements:

    >>> dds.pop(MockImportable(1))
    <(1,) IMPORTED>
    >>> dds.pop(MockImportable(4))
    <(4,) IMPORTED>
    >>> sorted(dds, key=lambda x: x.args)
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>, <(5,) IMPORTED>]

    And everything should have been registered:

    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    []

    Adding and removing the same
    :py:class:`~.importtools.importables.Importable` shouldn't change anything:

    >>> dds.add(MockImportable(100))
    >>> dds.pop(MockImportable(100))
    <(100,) IMPORTED>
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    []

    But removing an original :py:class:`~.importtools.importables.Importable`
    and adding a new one should be interpreted as a change:

    >>> dds.pop(MockImportable(2))
    <(2,) IMPORTED>
    >>> dds.add(MockImportable(2))
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(2,) IMPORTED>]

    If you change one of the existing elements it will be marked as changed:

    >>> dds.get(MockImportable(3)).register_change()
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(2,) IMPORTED>, <(3,) IMPORTED>]

    Even if the elements are accessed by iterating on over the dataset the
    changes are still tracked:

    >>> list(dds)[0].register_change()
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    When printing the object a list of all the added, removed and created
    elements will be shown, along with the total number for each operation:

    >>> print dds
    To be added: 1
    <(5,) IMPORTED>
    To be removed: 1
    <(1,) IMPORTED>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._added = MemoryDataSet()
    >>> dds._removed = MemoryDataSet()
    >>> print dds
    <BLANKLINE>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._changed = MemoryDataSet()
    >>> print dds
    No elements were affected.

    """
    def __init__(self, data_loader=None):
        self._added = set()
        self._removed = set()
        self._changed = set()
        # XXX for some reason dict subclass methods can't be hashed
        self.rc = lambda x: self.register_change(x)
        super(DiffDataSet, self).__init__(
            element.register_listener(self.rc) for element in data_loader
        )

    def __str__(self):
        result = ''

        added_count = len(self._added)
        if added_count > 0:
            result += 'To be added: %s\n' % added_count
            result += '\n'.join(sorted(repr(x) for x in self.get_added()))

        removed_count = len(self._removed)
        if removed_count > 0:
            result += '\nTo be removed: %s\n' % removed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_removed()))

        changed_count = len(self._changed)
        if changed_count > 0:
            result += '\nTo be changed: %s\n' % changed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_changed()))

        if (added_count + removed_count + changed_count) == 0:
            result += 'No elements were affected.'

        return result

    def add(self, importable):
        super(DiffDataSet, self).add(
            importable.register_listener(self.rc)
        )

        if importable in self._removed:
            self._removed.discard(importable)
            self._changed.add(importable)
        else:
            self._added.add(importable)

    def pop(self, importable):
        if importable in self._added:
            self._added.discard(importable)
        else:
            self._removed.add(importable)

        return super(DiffDataSet, self).pop(importable)

    def register_change(self, importable):
        """Mark an element in the current dataset as changed.

        This method is registered as a listener for changes in all the
        elements of the wrapped :py:class:`DataSet`.

        """
        assert importable in self
        self._changed.add(importable)

    def get_added(self):
        """A set of all added elements in the wrapped dataset."""
        return self._added

    def get_removed(self):
        """A set of all removed elements in the wrapped dataset."""
        return self._removed

    def get_changed(self):
        """A set of all changed elements in the wrapped dataset."""
        return self._changed
<|endoftext|>


<commit_before>import sys
import argparse
import sys
import logging
import uuid

import zmq

from gtdoit import messages

class LogBookKeyError(KeyError):
    pass

class LogBookEventOrderError(IndexError):
    pass

class EventStore(object):
    """Stores events and keeps track of their order.

    Thread safe.

    TODO: Currently this uses an in-memory implementation. Needs to be
    persisted.
    """
    class EventKeyAlreadyExistError(LogBookKeyError):
        pass

    class EventKeyDoesNotExistError(LogBookKeyError):
        pass

    class EventIndexError(IndexError):
        pass

    def __init__(self):
        self.reset()

    def _reset(self):
        """Reset the event storage.

        Protected/private because it should only be called from tests.

        Calling this method makes it thread unsafe.
        """
        self.keys = []
        self.events = {}

    def add_event(self, key, event):
        if key in self.keys or key in self.events:
            raise EventKeyAlreadyExistError("The key already existed: %s" % key)
        self.keys.append(key)
        # Important no exceptions happen between these lines!
        self.events[key] = event

    def get_events(self, from_=None, to=None):
        if from_ and (from_ not in self.keys or from_ not in self.events):
            raise EventKeyDoesNotExistError("Could not find the from_ key: %s" %
                                            from_)
        if to and (to not in self.keys or to not in self.events):
            raise EventKeyDoesNotExistError("Could not find the from_ key: %s" %
                                            to)
        fromindex = self.keys.index(from_) if from_ else 0
        toindex = self.keys.index(to) if to else len(self.events)
        if fromindex > toindex:
            raise LogBookEventOrderError("'From' index came after 'To'. \
                                         Keys: (%s, %s) \
                                         Indices: (%s, %s)" % (from_, to,
                                                               fromindex,
                                                               toindex))
        return (self.events[key] for key in self.events[fromindex:toindex])

    def key_exists(self, key):
        return key in self.keys

class IdGenerator:
    """Generates unique strings."""
    def __init__(self, key_exists=lambda key: False):
        self.key_exists = key_exists

    def _propose_new_key(self):
        return str(uuid.uuid4())

    def generate(self):
        key = self._propose_new_key()
        while self.key_exists(key):
            # TODO: Make sure that we don't end up in a infinite loop here
            key = self._propose_new_key()
        return key

def build_socket(context, socket_type, bind_endpoints=[], connect_endpoints=[]):
    socket = context.socket(socket_type)
    for endpoint in bind_endpoints:
        socket.bind(endpoint)
    for endpoint in connect_endpoints:
        socket.connect(endpoint)
    return socket

def run(args):
    """Actually execute the program."""
    context = zmq.Context(1)
    incoming_socket = build_socket(context, zmq.PULL,
                                   bind_endpoints=args.incoming_bind_endpoints,
                                   connect_endpoints=args.incoming_connect_endpoints)
    query_socket = build_socket(context, zmq.REP,
                                bind_endpoints=args.query_bind_endpoints,
                                connect_endpoints=args.query_connect_endpoints)
    streaming_socket = build_socket(context, xmq.PUB,
                                    bind_endpoints=args.streaming_bind_endpoints,
                                    connect_endpoints=args.streaming_connect_endpoints)

    eventstore = EventStore()
    id_generator = IdGenerator(key_exists=lambda key:
                               eventstore.key_exists(key))

    # Initialize poll set
    poller = zmq.Poller()
    poller.register(incoming_socket, zmq.POLLIN)
    poller.register(query_socket, zmq.POLLIN)

    while True:
        socks = dict(poller.poll())

        if incoming_socket in socks and socks[incoming_socket]==zmq.POLLIN:
            event = incoming_socket.recv()
            # TODO: Parse the incoming event
            # TODO: Create a new StoredEvent
            # TODO: Give the StoredEvent a unique ID
            # TODO: Store the recently received event in the new StoredEvent
            # TODO: Serialize the StoredEvent
            # TODO: Save the serialized StoredEvent to event store
            # TODO: Sent the serialized StoredEvent to streaming_socket

        if query_socket in socks and socks[query_socket]==zmq.POLLIN:
            # TODO: Parse the incoming request
            # TODO: For-loop over the eventstore query result and send using
            #       SENDMORE
            pass

    return 0

def main(argv=None, exit=True):
    """Entry point for the logbook.

    Parses input and calls run() for the real work.

    Parameters:
        argv -- sys.argv arguments. Can be set for testing purposes.
        exit -- whether to call sys.exit(...) when this function is done, or
                not.

    returns -- the return code of the programif exit is set to True. Otherwise
               it exits the Python interpreter before returning.
    """
    parser = argparse.ArgumentParser(
        description='Event storage and event proxy.'
    )

    incoming_group = parser.add_argument_group(
        title='Incoming event endpoints',
        description='ZeroMQ endpoint for incoming events.'
    )
    incoming_group.add_argument('--incoming-bind-endpoint', action='append',
                                metavar='ZEROMQ-ENDPOINT',
                                help='the bind address for incoming events',
                                dest='incoming_bind_endpoints')
    incoming_group.add_argument('--incoming-connect-endpoint', action='append',
                                metavar='ZEROMQ-ENDPOINT',
                                help='the connect address for incoming events',
                                dest='incoming_connect_endpoints')
    query_group = parser.add_argument_group(
        title='Querying endpoints',
        description='Endpoints listening for event queries.'
    )
    query_group.add_argument('--query-bind-endpoint',
                             metavar='ZEROMQ-ENDPOINT',
                             help='the bind address for querying of events',
                             action='append', dest='query_bind_endpoints')
    query_group.add_argument('--query-connect-endpoint',
                             metavar='ZEROMQ-ENDPOINT',
                             help='the connect address for querying of events',
                             action='append', dest='query_connect_endpoints')
    stream_group = parser.add_argument_group(
        title='Streaming endpoints',
        description='Endpoints for streaming incoming events.'
    )
    stream_group.add_argument('--streaming-bind-endpoint',
                              metavar='ZEROMQ-ENDPOINT',
                              help='the bind address for streaming of events',
                              action='append',
                              dest='streaming_bind_endpoints')
    stream_group.add_argument('--streaming-connect-endpoint',
                              metavar='ZEROMQ-ENDPOINT',
                              help='the connect address for streaming of events',
                              action='append',
                              dest='streaming_connect_endpoints')

    args = argv if argv else sys.argv[1:]
    args = parser.parse_args(args)

    if not args.incoming_bind_endpoints \
       and not args.incoming_connect_endpoints \
       and not args.query_bind_endpoints \
       and not args.query_connect_endpoints:
        errmsg = "You must either specify an incoming or query endpoint.\n" \
                "(there's no use in simply having a streaming endpoint)"
        if exit:
            parser.error(errmsg)
        else:
            print errmsg
            return 2

    exitcode = run(args)

    if exit:
        sys.exit(exitcode)
    else:
        return exitcode
<commit_msg>Fixing typo<commit_after>import sys
import argparse
import sys
import logging
import uuid

import zmq

from gtdoit import messages

class LogBookKeyError(KeyError):
    pass

class LogBookEventOrderError(IndexError):
    pass

class EventStore(object):
    """Stores events and keeps track of their order.

    Thread safe.

    TODO: Currently this uses an in-memory implementation. Needs to be
    persisted.
    """
    class EventKeyAlreadyExistError(LogBookKeyError):
        pass

    class EventKeyDoesNotExistError(LogBookKeyError):
        pass

    class EventIndexError(IndexError):
        pass

    def __init__(self):
        self.reset()

    def _reset(self):
        """Reset the event storage.

        Protected/private because it should only be called from tests.

        Calling this method makes it thread unsafe.
        """
        self.keys = []
        self.events = {}

    def add_event(self, key, event):
        if key in self.keys or key in self.events:
            raise EventKeyAlreadyExistError("The key already existed: %s" % key)
        self.keys.append(key)
        # Important no exceptions happen between these lines!
        self.events[key] = event

    def get_events(self, from_=None, to=None):
        if from_ and (from_ not in self.keys or from_ not in self.events):
            raise EventKeyDoesNotExistError("Could not find the from_ key: %s" %
                                            from_)
        if to and (to not in self.keys or to not in self.events):
            raise EventKeyDoesNotExistError("Could not find the from_ key: %s" %
                                            to)
        fromindex = self.keys.index(from_) if from_ else 0
        toindex = self.keys.index(to) if to else len(self.events)
        if fromindex > toindex:
            raise LogBookEventOrderError("'From' index came after 'To'. \
                                         Keys: (%s, %s) \
                                         Indices: (%s, %s)" % (from_, to,
                                                               fromindex,
                                                               toindex))
        return (self.events[key] for key in self.events[fromindex:toindex])

    def key_exists(self, key):
        return key in self.keys

class IdGenerator:
    """Generates unique strings."""
    def __init__(self, key_exists=lambda key: False):
        self.key_exists = key_exists

    def _propose_new_key(self):
        return str(uuid.uuid4())

    def generate(self):
        key = self._propose_new_key()
        while self.key_exists(key):
            # TODO: Make sure that we don't end up in a infinite loop here
            key = self._propose_new_key()
        return key

def build_socket(context, socket_type, bind_endpoints=[], connect_endpoints=[]):
    socket = context.socket(socket_type)
    for endpoint in bind_endpoints:
        socket.bind(endpoint)
    for endpoint in connect_endpoints:
        socket.connect(endpoint)
    return socket

def run(args):
    """Actually execute the program."""
    context = zmq.Context(1)
    incoming_socket = build_socket(context, zmq.PULL,
                                   bind_endpoints=args.incoming_bind_endpoints,
                                   connect_endpoints=args.incoming_connect_endpoints)
    query_socket = build_socket(context, zmq.REP,
                                bind_endpoints=args.query_bind_endpoints,
                                connect_endpoints=args.query_connect_endpoints)
    streaming_socket = build_socket(context, zmq.PUB,
                                    bind_endpoints=args.streaming_bind_endpoints,
                                    connect_endpoints=args.streaming_connect_endpoints)

    eventstore = EventStore()
    id_generator = IdGenerator(key_exists=lambda key:
                               eventstore.key_exists(key))

    # Initialize poll set
    poller = zmq.Poller()
    poller.register(incoming_socket, zmq.POLLIN)
    poller.register(query_socket, zmq.POLLIN)

    while True:
        socks = dict(poller.poll())

        if incoming_socket in socks and socks[incoming_socket]==zmq.POLLIN:
            event = incoming_socket.recv()
            # TODO: Parse the incoming event
            # TODO: Create a new StoredEvent
            # TODO: Give the StoredEvent a unique ID
            # TODO: Store the recently received event in the new StoredEvent
            # TODO: Serialize the StoredEvent
            # TODO: Save the serialized StoredEvent to event store
            # TODO: Sent the serialized StoredEvent to streaming_socket

        if query_socket in socks and socks[query_socket]==zmq.POLLIN:
            # TODO: Parse the incoming request
            # TODO: For-loop over the eventstore query result and send using
            #       SENDMORE
            pass

    return 0

def main(argv=None, exit=True):
    """Entry point for the logbook.

    Parses input and calls run() for the real work.

    Parameters:
        argv -- sys.argv arguments. Can be set for testing purposes.
        exit -- whether to call sys.exit(...) when this function is done, or
                not.

    returns -- the return code of the programif exit is set to True. Otherwise
               it exits the Python interpreter before returning.
    """
    parser = argparse.ArgumentParser(
        description='Event storage and event proxy.'
    )

    incoming_group = parser.add_argument_group(
        title='Incoming event endpoints',
        description='ZeroMQ endpoint for incoming events.'
    )
    incoming_group.add_argument('--incoming-bind-endpoint', action='append',
                                metavar='ZEROMQ-ENDPOINT',
                                help='the bind address for incoming events',
                                dest='incoming_bind_endpoints')
    incoming_group.add_argument('--incoming-connect-endpoint', action='append',
                                metavar='ZEROMQ-ENDPOINT',
                                help='the connect address for incoming events',
                                dest='incoming_connect_endpoints')
    query_group = parser.add_argument_group(
        title='Querying endpoints',
        description='Endpoints listening for event queries.'
    )
    query_group.add_argument('--query-bind-endpoint',
                             metavar='ZEROMQ-ENDPOINT',
                             help='the bind address for querying of events',
                             action='append', dest='query_bind_endpoints')
    query_group.add_argument('--query-connect-endpoint',
                             metavar='ZEROMQ-ENDPOINT',
                             help='the connect address for querying of events',
                             action='append', dest='query_connect_endpoints')
    stream_group = parser.add_argument_group(
        title='Streaming endpoints',
        description='Endpoints for streaming incoming events.'
    )
    stream_group.add_argument('--streaming-bind-endpoint',
                              metavar='ZEROMQ-ENDPOINT',
                              help='the bind address for streaming of events',
                              action='append',
                              dest='streaming_bind_endpoints')
    stream_group.add_argument('--streaming-connect-endpoint',
                              metavar='ZEROMQ-ENDPOINT',
                              help='the connect address for streaming of events',
                              action='append',
                              dest='streaming_connect_endpoints')

    args = argv if argv else sys.argv[1:]
    args = parser.parse_args(args)

    if not args.incoming_bind_endpoints \
       and not args.incoming_connect_endpoints \
       and not args.query_bind_endpoints \
       and not args.query_connect_endpoints:
        errmsg = "You must either specify an incoming or query endpoint.\n" \
                "(there's no use in simply having a streaming endpoint)"
        if exit:
            parser.error(errmsg)
        else:
            print errmsg
            return 2

    exitcode = run(args)

    if exit:
        sys.exit(exitcode)
    else:
        return exitcode
<|endoftext|>


<commit_before>from gevent import event, Greenlet, sleep

from inbox.log import get_logger
logger = get_logger()

from inbox.util.concurrency import retry_and_report_killed
from inbox.heartbeat.status import HeartbeatStatusProxy, clear_heartbeat_status
from inbox.basicauth import ConnectionError, ValidationError
from inbox.models import Account
from inbox.models.session import session_scope

class BaseSyncMonitor(Greenlet):
    """
    Abstracted sync monitor, based on BaseMailSyncMonitor but not mail-specific

    Subclasses should run
    bind_context(self, 'mailsyncmonitor', account.id)

    poll_frequency : int
        How often to check for commands.
    retry_fail_classes : list
        Exceptions to *not* retry on.

    """
    def __init__(self, account_id, namespace_id, folder_id, poll_frequency=1,
                 retry_fail_classes=[]):

        self.account_id = account_id
        self.namespace_id = namespace_id
        self.poll_frequency = poll_frequency
        self.retry_fail_classes = retry_fail_classes
        self.folder_id = folder_id

        self.log = logger.new(account_id=account_id)

        self.shutdown = event.Event()
        self.heartbeat_status = HeartbeatStatusProxy(self.account_id,
                                                     self.folder_id)

        Greenlet.__init__(self)

    def _run(self):
        return retry_and_report_killed(self._run_impl,
                                       account_id=self.account_id,
                                       logger=self.log,
                                       fail_classes=self.retry_fail_classes)

    def _run_impl(self):
        # Return true/false based on whether the greenlet succeeds or throws
        # and error. Note this is not how the mailsync monitor works
        try:
            while True:
                # Check to see if this greenlet should exit
                if self.shutdown.is_set():
                    self._cleanup()
                    return False

                try:
                    self.sync()
                    self.heartbeat_status.publish(state='poll')

                # If we get a connection or API permissions error, then sleep
                # 2x poll frequency.
                except ConnectionError:
                    self.log.error('Error while polling', exc_info=True)
                    self.heartbeat_status.publish(state='poll error')
                    sleep(self.poll_frequency)
                sleep(self.poll_frequency)
        except ValidationError:
            # Bad account credentials; exit.
            self.log.error('Error while establishing the connection',
                           exc_info=True)
            self._cleanup()
            with session_scope() as db_session:
                account = db_session.query(Account).get(self.account_id)
                account.sync_state = 'invalid'

            return False

    def sync(self):
        """ Subclasses should override this to do work """
        raise NotImplementedError

    def _cleanup(self):
        clear_heartbeat_status(self.account_id, folder_id=self.folder_id)
<commit_msg>Start heartbeat proxy with correct params<commit_after>from gevent import event, Greenlet, sleep

from inbox.log import get_logger
logger = get_logger()

from inbox.util.concurrency import retry_and_report_killed
from inbox.heartbeat.status import HeartbeatStatusProxy, clear_heartbeat_status
from inbox.basicauth import ConnectionError, ValidationError
from inbox.models import Account
from inbox.models.session import session_scope

class BaseSyncMonitor(Greenlet):
    """
    Abstracted sync monitor, based on BaseMailSyncMonitor but not mail-specific

    Subclasses should run
    bind_context(self, 'mailsyncmonitor', account.id)

    poll_frequency : int
        How often to check for commands.
    retry_fail_classes : list
        Exceptions to *not* retry on.

    """
    def __init__(self, account_id, namespace_id, folder_id, poll_frequency=1,
                 retry_fail_classes=[]):

        self.account_id = account_id
        self.namespace_id = namespace_id
        self.poll_frequency = poll_frequency
        self.retry_fail_classes = retry_fail_classes
        self.folder_id = folder_id

        self.log = logger.new(account_id=account_id)

        self.shutdown = event.Event()
        self.heartbeat_status = HeartbeatStatusProxy(self.account_id,
                                                     self.folder_id)
        self.heartbeat_status.publish(email_address=self.email_address,
                                      provider_name=self._provider_name,
                                      folder_name=self.folder_name)

        Greenlet.__init__(self)

    def _run(self):
        return retry_and_report_killed(self._run_impl,
                                       account_id=self.account_id,
                                       logger=self.log,
                                       fail_classes=self.retry_fail_classes)

    def _run_impl(self):
        # Return true/false based on whether the greenlet succeeds or throws
        # and error. Note this is not how the mailsync monitor works
        try:
            while True:
                # Check to see if this greenlet should exit
                if self.shutdown.is_set():
                    self._cleanup()
                    return False

                try:
                    self.sync()
                    self.heartbeat_status.publish(state='poll')

                # If we get a connection or API permissions error, then sleep
                # 2x poll frequency.
                except ConnectionError:
                    self.log.error('Error while polling', exc_info=True)
                    self.heartbeat_status.publish(state='poll error')
                    sleep(self.poll_frequency)
                sleep(self.poll_frequency)
        except ValidationError:
            # Bad account credentials; exit.
            self.log.error('Error while establishing the connection',
                           exc_info=True)
            self._cleanup()
            with session_scope() as db_session:
                account = db_session.query(Account).get(self.account_id)
                account.sync_state = 'invalid'

            return False

    def sync(self):
        """ Subclasses should override this to do work """
        raise NotImplementedError

    def _cleanup(self):
        clear_heartbeat_status(self.account_id, folder_id=self.folder_id)
<|endoftext|>


<commit_before># Copyright 2020 KMEE
# Copyright (C) 2020 - TODAY Renato Lima - Akretion
# License AGPL-3.0 or later (http://www.gnu.org/lic enses/agpl).

{
    "name": "Base do plano de conta",
    "summary": """
        Base do plano de conta brasileiro""",
    "version": "13.0.1.0.0",
    "license": "AGPL-3",
    "author": "Akretion, KMEE, Odoo Community Association (OCA)",
    "maitainers": ["renatonlima", "mileo"],
    "category": "Accounting",
    "website": "https://github.com/OCA/l10n-brazil",
    "depends": ["account"],
    "data": [
        # Data
        "data/l10n_br_coa_template.xml",
        "data/account_tax_tag.xml",
        "data/account_tax_group.xml",
        "data/account_tax_template.xml",
        "data/account_type_data.xml",
        # Views
        "views/account_tax_template.xml",
        "views/account_tax.xml",
    ],
    "development_status": "Production/Stable",
    "installable": True,
}
<commit_msg>[MIG] l10n_br_coa: Migration to 14.0<commit_after># Copyright 2020 KMEE
# Copyright (C) 2020 - TODAY Renato Lima - Akretion
# License AGPL-3.0 or later (http://www.gnu.org/lic enses/agpl).

{
    "name": "Base do plano de conta",
    "summary": """
        Base do plano de conta brasileiro""",
    "version": "14.0.1.0.0",
    "license": "AGPL-3",
    "author": "Akretion, KMEE, Odoo Community Association (OCA)",
    "maitainers": ["renatonlima", "mileo"],
    "category": "Accounting",
    "website": "https://github.com/OCA/l10n-brazil",
    "depends": ["account"],
    "data": [
        # Data
        "data/l10n_br_coa_template.xml",
        "data/account_tax_tag.xml",
        "data/account_tax_group.xml",
        "data/account_tax_template.xml",
        "data/account_type_data.xml",
        # Views
        "views/account_tax_template.xml",
        "views/account_tax.xml",
    ],
    "development_status": "Production/Stable",
    "installable": True,
}
<|endoftext|>


<commit_before># -*-*- encoding: utf-8 -*-*-
#
# gateway4labs is free software: you can redistribute it and/or modify
# it under the terms of the BSD 2-Clause License
# gateway4labs is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

import sys
import time
import threading
import datetime
import traceback
import requests

from flask import url_for

from labmanager.db import db
from labmanager.models import RLMS as dbRLMS
from labmanager.application import app
from .base import register_blueprint, BaseRLMS, BaseFormCreator, Capabilities, Versions
from .caches import GlobalCache, VersionCache, InstanceCache, EmptyCache, get_cached_session, CacheDisabler, clean_cache

assert BaseFormCreator or register_blueprint or Versions or Capabilities or BaseRLMS or True # Avoid pyflakes warnings

#
# Add the proper managers by pointing to a module
#

_RLMSs = {
    # "RLMS type" :  <module>,

    # e.g.
    # "WebLab-Deusto" : ( labmanager.rlms.ext.weblabdeusto, ['4.0'], record ),
}

_GLOBAL_PERIODIC_TASKS = [
    # {
    #     'name' : "WebLab-Deusto",
    #     'versions' : ["4.0", "5.0"],
    #     'func' : func,
    #     'when' : datetime.timedelta
    # }
]

_LOCAL_PERIODIC_TASKS = [
    # {
    #     'name' : "WebLab-Deusto",
    #     'versions' : ["4.0", "5.0"],
    #     'func' : func,
    #     'when' : datetime.timedelta
    # }
]

class Laboratory(object):
    def __init__(self, name, laboratory_id, description=None, autoload=False, age_ranges=None, domains=None, long_description=None, keywords=None, preview_url=None, home_url=None, **kwargs):
        self.name             = name
        self.laboratory_id    = laboratory_id
        self.description      = description
        self.long_description = long_description
        self.autoload         = autoload
        self.age_ranges       = age_ranges
        self.domains          = domains
        self.keywords         = keywords
        self.preview_url      = preview_url
        self.home_url         = home_url

    def __repr__(self):
        return "Laboratory(%r, %r, %r, %r, %r, %r, %r)" % (self.name, self.laboratory_id, self.description, self.autoload, self.age_ranges, self.domains, self.long_description)

    def __hash__(self):
        return hash(self.laboratory_id)

class LabError(Exception):
    pass

class LabNotFoundError(LabError):
    pass

_current_task_id = 0
def _next_task_id():
    global _current_task_id
    _current_task_id += 1
    return _current_task_id

class _RegistrationRecord(object):
    def __init__(self, name, versions):
        self.name = name
        self.versions = versions
        global_key = '%s - %s' % (name, ', '.join(versions))
        self.global_cache = self.cache = GlobalCache(global_key)
        self.per_version_cache = {}
        for version in versions:
            # If a single version it's supported, the global cache is the same as the per version cache
            version_key = '%s - %s' % (name, version)
            self.per_version_cache[version] = VersionCache(version_key)

        self.per_thread = threading.local()

    def get_cache(self, version = None):
        if version is None:
            return self.cache
        else:
            return self.version_cache[version]

    @property
    def rlms_cache(self):
        current_rlms_id = getattr(self.per_thread, 'current_rlms_id', None)
        if current_rlms_id is not None:
            return InstanceCache(current_rlms_id)
        else:
            return EmptyCache()

    @property
    def cached_session(self):
        cached_session = getattr(self.per_thread, 'cached_session', None)
        if cached_session:
            return cached_session

        cached_session = get_cached_session()
        self.per_thread.cached_session = cached_session

        def timeout_get(url, timeout = (30, 60), max_retries = 3, **kwargs):
            count = 0
            while True:
                try:
                    return cached_session.get(url, timeout = timeout, **kwargs)
                except requests.Timeout:
                    count += 1
                    if count >= max_retries:
                        raise

        cached_session.timeout_get = timeout_get
        return cached_session

    def _add_periodic_task(self, where, task_name, function, hours, minutes, disable_cache):
        if hours == 0 and minutes == 0:
            raise ValueError("You must establish hours or minutes")

        if hours < 0 or minutes < 0:
            raise ValueError("hours and minutes must be positive numbers")

        where.append({
            'id'   : _next_task_id(),
            'name' : task_name,
            'rlms' : self.name,
            'versions' : self.versions,
            'func' : function,
            'when' : datetime.timedelta(hours = hours, minutes = minutes),
            'disable_cache' : disable_cache,
        })

    def add_global_periodic_task(self, task_name, function, hours = 0, minutes = 0, disable_cache = True):
        return self._add_periodic_task(_GLOBAL_PERIODIC_TASKS, task_name, function, hours, minutes, disable_cache)

    def add_local_periodic_task(self, task_name, function, hours = 0, minutes = 0, disable_cache = True):
        return self._add_periodic_task(_LOCAL_PERIODIC_TASKS, task_name, function, hours, minutes, disable_cache)

    def is_debug(self):
        return 'debug' in sys.argv or '-debug' in sys.argv or '--debug' in sys.argv

def _debug(msg):
    sys.stderr.flush()
    sys.stdout.flush()
    print u"[%s] - %s" % (time.ctime(), msg)
    sys.stdout.flush()
    sys.stderr.flush()

class TaskRunner(object):
    def __init__(self):
        # task_id : datetime.datetime
        self.latest_executions = {
        }
        self._stopping = False

    def _now(self):
        return datetime.datetime.now().replace(second = 0, microsecond = 0)

    def _must_be_run(self, now, task):
        if task['id'] not in self.latest_executions:
            return True

        time_elapsed = now - self.latest_executions[task['id']]
        if time_elapsed >= task['when']:
            return True

        return False

    def _run_all(self):
        # Run global tasks
        cache_disabler = CacheDisabler()

        for task in _GLOBAL_PERIODIC_TASKS:
            now = self._now()
            if self._must_be_run(now, task):
                for version in task['versions']:
                    _debug("Running task %r for rlms %s %s..." % (task['name'], task['rlms'], version))
                    if task['disable_cache']:
                        cache_disabler.disable()
                    try:
                        task['func']()
                    except Exception:
                        traceback.print_exc()
                    finally:
                        cache_disabler.reenable()
                    _debug(u"Finished")

                self.latest_executions[task['id']] = now

        # Same for regular
        for task in _LOCAL_PERIODIC_TASKS:
            now = self._now()
            if self._must_be_run(now, task):
                with app.app_context():
                    for version in task['versions']:
                        rlmss = db.session.query(dbRLMS).filter_by(kind = task['rlms'], version = version).all()
                        for db_rlms in rlmss:
                            ManagerClass = get_manager_class(db_rlms.kind, db_rlms.version, db_rlms.id)
                            rlms = ManagerClass(db_rlms.configuration)
                            _debug(u"Running task %r for rlms %s %s (%s)..." % (repr(task['name']), repr(task['rlms']), repr(version), repr(db_rlms.location)))
                            if task['disable_cache']:
                                cache_disabler.disable()

                            try:
                                task['func'](rlms)
                            except Exception:
                                traceback.print_exc()
                            finally:
                                cache_disabler.reenable()
                            _debug(u"Finished")

                self.latest_executions[task['id']] = now

    def _step(self, initial):
        before = self._now()
        if before.hour == initial.hour and before.minute == initial.minute:
            clean_cache()

        future = before + datetime.timedelta(minutes = 1)
        future = future.replace(second = 0, microsecond = 0)
        self._run_all()
        after = datetime.datetime.now()

        remaining_seconds = (future - after).total_seconds()
        if remaining_seconds > 0:
            PIECES = 60
            per_second = remaining_seconds / PIECES
            for _ in xrange(PIECES):
                if self._stopping:
                    break
                time.sleep(per_second)
        else:
            print "Warning: the last run_all took: %s time" % (after - before).total_seconds()

    def run_forever(self):
        initial = self._now()
        while not self._stopping:
            self._step(initial)

    def stop(self):
        self._stopping = True

def register(name, versions, module_name):
    record = _RegistrationRecord(name, versions)
    _RLMSs[name] = (module_name, versions, record)
    return record

def get_supported_types():
    return _RLMSs.keys()

def get_supported_versions(rlms_type):
    if rlms_type in _RLMSs:
        _, versions, _ = _RLMSs[rlms_type]
        return versions
    return []

def get_all_rlms():
    rlmss = []
    for db_rlms in db.session.query(dbRLMS).all():
        try:
            ManagerClass = get_manager_class(db_rlms.kind, db_rlms.version, db_rlms.id)
            rlms = ManagerClass(db_rlms.configuration)
        except:
            traceback.print_exc()
            continue
        rlmss.append(rlms)
    return rlmss

def is_supported(rlms_type, rlms_version):
    _, versions, _ = _RLMSs.get(rlms_type, (None, []))
    return rlms_version in versions

def _get_module(rlms_type, rlms_version):
    module_name, versions, _ = _RLMSs.get(rlms_type, (None, []))
    if rlms_version in versions:
        if hasattr(sys.modules[module_name], 'get_module'):
            return sys.modules[module_name].get_module(rlms_version)
        else:
            return sys.modules[module_name]
    else:
        raise Exception(u"Misconfiguration: %(rlmstype)s %(rmlsversion)s does not exist" % dict(rlmstype = rlms_type, rmlsversion = rlms_version))

def _get_form_creator(rlms_type, rlms_version):
    return _get_module(rlms_type, rlms_version).FORM_CREATOR

def get_form_class(rlms_type, rlms_version):
    form_creator = _get_form_creator(rlms_type, rlms_version)
    return form_creator.get_add_form()

def get_permissions_form_class(rlms_type, rlms_version):
    form_creator = _get_form_creator(rlms_type, rlms_version)
    return form_creator.get_permission_form()

def get_lms_permissions_form_class(rlms_type, rlms_version):
    form_creator = _get_form_creator(rlms_type, rlms_version)
    return form_creator.get_lms_permission_form()

def get_manager_class(rlms_type, rlms_version, current_rlms_id = None):
    module = _get_module(rlms_type, rlms_version)
    _, _, record = _RLMSs[rlms_type]
    if current_rlms_id:
        record.per_thread.current_rlms_id = current_rlms_id
    return module.RLMS

def find_smartgateway_link(url, return_url):
    rlms_by_id = {}
    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=True).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    if lab is not None:
                        return url_for('bookmarklet.public_rlms', rlms_id=db_rlms.public_identifier, lab_name=lab.laboratory_id, url=return_url)

    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=False).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    db_lab = db.session.query(Laboratory).filter_by(rlms=db_rlms, laboratory_id=lab.laboratory_id, publicly_available=True).first()
                    if db_lab is not None:
                        return url_for('bookmarklet.public_lab', public_identifier=db_lab.public_identifier, url=return_url)
    return None

def _get_widgets(rlms, laboratory_id):
    if Capabilities.WIDGET in rlms.get_capabilities():
        return rlms.list_widgets(laboratory_id)

    return [ { 'name' : 'lab', 'description' : 'Main view of the laboratory' } ]

def find_smartgateway_opensocial_link(url):
    rlms_by_id = {}
    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=True).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    if lab is not None:
                        for widget in _get_widgets(rlms, lab.laboratory_id):
                            # First widget wins...
                            return url_for('opensocial.public_rlms_widget_xml', rlms_identifier=db_rlms.public_identifier, lab_name=lab.laboratory_id, widget_name=widget['name'], _external=True)

    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=False).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    db_lab = db.session.query(Laboratory).filter_by(rlms=db_rlms, laboratory_id=lab.laboratory_id, publicly_available=True).first()
                    if db_lab is not None:
                        for widget in _get_widgets(rlms, lab.laboratory_id):
                            # First widget wins...
                            return url_for('opensocial.public_widget_xml', lab_name=db_lab.public_identifier, widget_name=widget['name'], _external=True)
    return None

<commit_msg>Fix Laboratory for dbLaboratory<commit_after># -*-*- encoding: utf-8 -*-*-
#
# gateway4labs is free software: you can redistribute it and/or modify
# it under the terms of the BSD 2-Clause License
# gateway4labs is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

import sys
import time
import threading
import datetime
import traceback
import requests

from flask import url_for

from labmanager.db import db
from labmanager.models import RLMS as dbRLMS, Laboratory as dbLaboratory
from labmanager.application import app
from .base import register_blueprint, BaseRLMS, BaseFormCreator, Capabilities, Versions
from .caches import GlobalCache, VersionCache, InstanceCache, EmptyCache, get_cached_session, CacheDisabler, clean_cache

assert BaseFormCreator or register_blueprint or Versions or Capabilities or BaseRLMS or True # Avoid pyflakes warnings

#
# Add the proper managers by pointing to a module
#

_RLMSs = {
    # "RLMS type" :  <module>,

    # e.g.
    # "WebLab-Deusto" : ( labmanager.rlms.ext.weblabdeusto, ['4.0'], record ),
}

_GLOBAL_PERIODIC_TASKS = [
    # {
    #     'name' : "WebLab-Deusto",
    #     'versions' : ["4.0", "5.0"],
    #     'func' : func,
    #     'when' : datetime.timedelta
    # }
]

_LOCAL_PERIODIC_TASKS = [
    # {
    #     'name' : "WebLab-Deusto",
    #     'versions' : ["4.0", "5.0"],
    #     'func' : func,
    #     'when' : datetime.timedelta
    # }
]

class Laboratory(object):
    def __init__(self, name, laboratory_id, description=None, autoload=False, age_ranges=None, domains=None, long_description=None, keywords=None, preview_url=None, home_url=None, **kwargs):
        self.name             = name
        self.laboratory_id    = laboratory_id
        self.description      = description
        self.long_description = long_description
        self.autoload         = autoload
        self.age_ranges       = age_ranges
        self.domains          = domains
        self.keywords         = keywords
        self.preview_url      = preview_url
        self.home_url         = home_url

    def __repr__(self):
        return "Laboratory(%r, %r, %r, %r, %r, %r, %r)" % (self.name, self.laboratory_id, self.description, self.autoload, self.age_ranges, self.domains, self.long_description)

    def __hash__(self):
        return hash(self.laboratory_id)

class LabError(Exception):
    pass

class LabNotFoundError(LabError):
    pass

_current_task_id = 0
def _next_task_id():
    global _current_task_id
    _current_task_id += 1
    return _current_task_id

class _RegistrationRecord(object):
    def __init__(self, name, versions):
        self.name = name
        self.versions = versions
        global_key = '%s - %s' % (name, ', '.join(versions))
        self.global_cache = self.cache = GlobalCache(global_key)
        self.per_version_cache = {}
        for version in versions:
            # If a single version it's supported, the global cache is the same as the per version cache
            version_key = '%s - %s' % (name, version)
            self.per_version_cache[version] = VersionCache(version_key)

        self.per_thread = threading.local()

    def get_cache(self, version = None):
        if version is None:
            return self.cache
        else:
            return self.version_cache[version]

    @property
    def rlms_cache(self):
        current_rlms_id = getattr(self.per_thread, 'current_rlms_id', None)
        if current_rlms_id is not None:
            return InstanceCache(current_rlms_id)
        else:
            return EmptyCache()

    @property
    def cached_session(self):
        cached_session = getattr(self.per_thread, 'cached_session', None)
        if cached_session:
            return cached_session

        cached_session = get_cached_session()
        self.per_thread.cached_session = cached_session

        def timeout_get(url, timeout = (30, 60), max_retries = 3, **kwargs):
            count = 0
            while True:
                try:
                    return cached_session.get(url, timeout = timeout, **kwargs)
                except requests.Timeout:
                    count += 1
                    if count >= max_retries:
                        raise

        cached_session.timeout_get = timeout_get
        return cached_session

    def _add_periodic_task(self, where, task_name, function, hours, minutes, disable_cache):
        if hours == 0 and minutes == 0:
            raise ValueError("You must establish hours or minutes")

        if hours < 0 or minutes < 0:
            raise ValueError("hours and minutes must be positive numbers")

        where.append({
            'id'   : _next_task_id(),
            'name' : task_name,
            'rlms' : self.name,
            'versions' : self.versions,
            'func' : function,
            'when' : datetime.timedelta(hours = hours, minutes = minutes),
            'disable_cache' : disable_cache,
        })

    def add_global_periodic_task(self, task_name, function, hours = 0, minutes = 0, disable_cache = True):
        return self._add_periodic_task(_GLOBAL_PERIODIC_TASKS, task_name, function, hours, minutes, disable_cache)

    def add_local_periodic_task(self, task_name, function, hours = 0, minutes = 0, disable_cache = True):
        return self._add_periodic_task(_LOCAL_PERIODIC_TASKS, task_name, function, hours, minutes, disable_cache)

    def is_debug(self):
        return 'debug' in sys.argv or '-debug' in sys.argv or '--debug' in sys.argv

def _debug(msg):
    sys.stderr.flush()
    sys.stdout.flush()
    print u"[%s] - %s" % (time.ctime(), msg)
    sys.stdout.flush()
    sys.stderr.flush()

class TaskRunner(object):
    def __init__(self):
        # task_id : datetime.datetime
        self.latest_executions = {
        }
        self._stopping = False

    def _now(self):
        return datetime.datetime.now().replace(second = 0, microsecond = 0)

    def _must_be_run(self, now, task):
        if task['id'] not in self.latest_executions:
            return True

        time_elapsed = now - self.latest_executions[task['id']]
        if time_elapsed >= task['when']:
            return True

        return False

    def _run_all(self):
        # Run global tasks
        cache_disabler = CacheDisabler()

        for task in _GLOBAL_PERIODIC_TASKS:
            now = self._now()
            if self._must_be_run(now, task):
                for version in task['versions']:
                    _debug("Running task %r for rlms %s %s..." % (task['name'], task['rlms'], version))
                    if task['disable_cache']:
                        cache_disabler.disable()
                    try:
                        task['func']()
                    except Exception:
                        traceback.print_exc()
                    finally:
                        cache_disabler.reenable()
                    _debug(u"Finished")

                self.latest_executions[task['id']] = now

        # Same for regular
        for task in _LOCAL_PERIODIC_TASKS:
            now = self._now()
            if self._must_be_run(now, task):
                with app.app_context():
                    for version in task['versions']:
                        rlmss = db.session.query(dbRLMS).filter_by(kind = task['rlms'], version = version).all()
                        for db_rlms in rlmss:
                            ManagerClass = get_manager_class(db_rlms.kind, db_rlms.version, db_rlms.id)
                            rlms = ManagerClass(db_rlms.configuration)
                            _debug(u"Running task %r for rlms %s %s (%s)..." % (repr(task['name']), repr(task['rlms']), repr(version), repr(db_rlms.location)))
                            if task['disable_cache']:
                                cache_disabler.disable()

                            try:
                                task['func'](rlms)
                            except Exception:
                                traceback.print_exc()
                            finally:
                                cache_disabler.reenable()
                            _debug(u"Finished")

                self.latest_executions[task['id']] = now

    def _step(self, initial):
        before = self._now()
        if before.hour == initial.hour and before.minute == initial.minute:
            clean_cache()

        future = before + datetime.timedelta(minutes = 1)
        future = future.replace(second = 0, microsecond = 0)
        self._run_all()
        after = datetime.datetime.now()

        remaining_seconds = (future - after).total_seconds()
        if remaining_seconds > 0:
            PIECES = 60
            per_second = remaining_seconds / PIECES
            for _ in xrange(PIECES):
                if self._stopping:
                    break
                time.sleep(per_second)
        else:
            print "Warning: the last run_all took: %s time" % (after - before).total_seconds()

    def run_forever(self):
        initial = self._now()
        while not self._stopping:
            self._step(initial)

    def stop(self):
        self._stopping = True

def register(name, versions, module_name):
    record = _RegistrationRecord(name, versions)
    _RLMSs[name] = (module_name, versions, record)
    return record

def get_supported_types():
    return _RLMSs.keys()

def get_supported_versions(rlms_type):
    if rlms_type in _RLMSs:
        _, versions, _ = _RLMSs[rlms_type]
        return versions
    return []

def get_all_rlms():
    rlmss = []
    for db_rlms in db.session.query(dbRLMS).all():
        try:
            ManagerClass = get_manager_class(db_rlms.kind, db_rlms.version, db_rlms.id)
            rlms = ManagerClass(db_rlms.configuration)
        except:
            traceback.print_exc()
            continue
        rlmss.append(rlms)
    return rlmss

def is_supported(rlms_type, rlms_version):
    _, versions, _ = _RLMSs.get(rlms_type, (None, []))
    return rlms_version in versions

def _get_module(rlms_type, rlms_version):
    module_name, versions, _ = _RLMSs.get(rlms_type, (None, []))
    if rlms_version in versions:
        if hasattr(sys.modules[module_name], 'get_module'):
            return sys.modules[module_name].get_module(rlms_version)
        else:
            return sys.modules[module_name]
    else:
        raise Exception(u"Misconfiguration: %(rlmstype)s %(rmlsversion)s does not exist" % dict(rlmstype = rlms_type, rmlsversion = rlms_version))

def _get_form_creator(rlms_type, rlms_version):
    return _get_module(rlms_type, rlms_version).FORM_CREATOR

def get_form_class(rlms_type, rlms_version):
    form_creator = _get_form_creator(rlms_type, rlms_version)
    return form_creator.get_add_form()

def get_permissions_form_class(rlms_type, rlms_version):
    form_creator = _get_form_creator(rlms_type, rlms_version)
    return form_creator.get_permission_form()

def get_lms_permissions_form_class(rlms_type, rlms_version):
    form_creator = _get_form_creator(rlms_type, rlms_version)
    return form_creator.get_lms_permission_form()

def get_manager_class(rlms_type, rlms_version, current_rlms_id = None):
    module = _get_module(rlms_type, rlms_version)
    _, _, record = _RLMSs[rlms_type]
    if current_rlms_id:
        record.per_thread.current_rlms_id = current_rlms_id
    return module.RLMS

def find_smartgateway_link(url, return_url):
    rlms_by_id = {}
    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=True).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    if lab is not None:
                        return url_for('bookmarklet.public_rlms', rlms_id=db_rlms.public_identifier, lab_name=lab.laboratory_id, url=return_url)

    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=False).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    db_lab = db.session.query(dbLaboratory).filter_by(rlms=db_rlms, laboratory_id=lab.laboratory_id, publicly_available=True).first()
                    if db_lab is not None:
                        return url_for('bookmarklet.public_lab', public_identifier=db_lab.public_identifier, url=return_url)
    return None

def _get_widgets(rlms, laboratory_id):
    if Capabilities.WIDGET in rlms.get_capabilities():
        return rlms.list_widgets(laboratory_id)

    return [ { 'name' : 'lab', 'description' : 'Main view of the laboratory' } ]

def find_smartgateway_opensocial_link(url):
    rlms_by_id = {}
    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=True).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    if lab is not None:
                        for widget in _get_widgets(rlms, lab.laboratory_id):
                            # First widget wins...
                            return url_for('opensocial.public_rlms_widget_xml', rlms_identifier=db_rlms.public_identifier, lab_name=lab.laboratory_id, widget_name=widget['name'], _external=True)

    for db_rlms in db.session.query(dbRLMS).filter_by(publicly_available=False).all():
        rlms = db_rlms.get_rlms()
        rlms_by_id[db_rlms.id] = rlms
        if Capabilities.URL_FINDER in rlms.get_capabilities():
            base_urls = rlms.get_base_urls() or []
            for base_url in base_urls:
                if url.startswith(base_url):
                    lab = rlms.get_lab_by_url(url)
                    db_lab = db.session.query(dbLaboratory).filter_by(rlms=db_rlms, laboratory_id=lab.laboratory_id, publicly_available=True).first()
                    if db_lab is not None:
                        for widget in _get_widgets(rlms, lab.laboratory_id):
                            # First widget wins...
                            return url_for('opensocial.public_widget_xml', lab_name=db_lab.public_identifier, widget_name=widget['name'], _external=True)
    return None

<|endoftext|>


<commit_before>#
# Copyright (C) 2021 Satoru SATOH <satoru.satoh@gmail.com>
# License: MIT
#
# pylint: disable=missing-docstring
import pathlib
import unittest

import anyconfig.api._load as TT

from .common import RES_DIR

def datasets_itr():
    for rdir in pathlib.Path(RES_DIR / 'multi').glob('*'):
        if not rdir.is_dir():
            continue

        exp = TT.single_load(rdir / 'e' / 'exp.json')
        opts = TT.single_load(rdir / 'options' / '00.json')

        yield (rdir, exp, opts)

def gen_datasets():
    datasets = datasets_itr()
    if not datasets:
        raise RuntimeError('No test data was found!')

    return datasets

class TestCase(unittest.TestCase):

    def test_multi_load_from_path_objects(self):
        for rdir, exp, opts in gen_datasets():
            res = TT.multi_load(rdir.glob('*.json'), **opts)
            self.assertEqual(res, exp)

    def test_multi_load_from_glob_path_str(self):
        for rdir, exp, opts in gen_datasets():
            res = TT.multi_load(str(rdir / '*.json'), **opts)
            self.assertEqual(res, exp)

# vim:sw=4:ts=4:et:
<commit_msg>change: make test resource dir dumped if the test case failed in tests.api.load.multi_load<commit_after>#
# Copyright (C) 2021 Satoru SATOH <satoru.satoh@gmail.com>
# License: MIT
#
# pylint: disable=missing-docstring
import pathlib
import unittest

import anyconfig.api._load as TT

from .common import RES_DIR

def datasets_itr():
    for rdir in pathlib.Path(RES_DIR / 'multi').glob('*'):
        if not rdir.is_dir():
            continue

        exp = TT.single_load(rdir / 'e' / 'exp.json')
        opts = TT.single_load(rdir / 'options' / '00.json')

        yield (rdir, exp, opts)

def gen_datasets():
    datasets = datasets_itr()
    if not datasets:
        raise RuntimeError('No test data was found!')

    return datasets

class TestCase(unittest.TestCase):

    def test_multi_load_from_path_objects(self):
        for rdir, exp, opts in gen_datasets():
            res = TT.multi_load(rdir.glob('*.json'), **opts)
            self.assertEqual(res, exp)

    def test_multi_load_from_glob_path_str(self):
        for rdir, exp, opts in gen_datasets():
            res = TT.multi_load(str(rdir / '*.json'), **opts)
            self.assertEqual(res, exp, f'{rdir!s}')

# vim:sw=4:ts=4:et:
<|endoftext|>


<commit_before>"""Print results of statistical checks and observations. Iterates only once
over the dataset. Implemented as a class to allow use during capture as well as
standalone."""

from __future__ import annotations
from itertools import product
from typing import Dict

from tabulate import tabulate
from tqdm import tqdm

from scaaml.io import Dataset
from .ap_counter import APCounter
from .ap_checker import APChecker
from .example_iterator import ExampleIterator
from .trace_stddev_of_stat import STDDEVofAVGofTraces
from .trace_stddev_of_stat import STDDEVofMAXofTraces
from .trace_stddev_of_stat import STDDEVofMINofTraces

class PrintStats:
    """Gather and print warnings from attack point checks and trace
    statistics.

    Example use:

    # On an existing dataset:
    print_stats = PrintStats.from_config(dataset_path='/mnt/path/to/dataset')

    # During capture:
    print_stats = PrintStats(measurements_info, attack_points_info)
    for example in example_iterator:
        print_stats.add_example(example=example,
                                split=split,
                                group=group,
                                part=part)
    print_stats.print()
    """
    def __init__(self, measurements_info: Dict,
                 attack_points_info: Dict) -> None:
        self._measurements_info = measurements_info
        self._attack_point_info = attack_points_info
        self._split_ap_counters = {
            split: {
                ap_name: [
                    APCounter(attack_point_info=ap_info),
                ]
                for ap_name, ap_info in attack_points_info.items()
            }
            for split in Dataset.SPLITS
        }
        self._split_trace_stats = {
            split: {
                trace_name: [
                    STDDEVofAVGofTraces(),
                    STDDEVofMAXofTraces(),
                    STDDEVofMINofTraces(),
                ]
                for trace_name in measurements_info
            }
            for split in Dataset.SPLITS
        }

    @staticmethod
    def from_config(dataset_path: str) -> PrintStats:
        """Read a dataset and print all attack point check warnings and trace
        statistics.

        Args:
          dataset_path: Root path of the dataset.
        """
        dataset = Dataset.from_config(dataset_path)
        print_stats = PrintStats(measurements_info=dataset.measurements_info,
                                 attack_points_info=dataset.attack_points_info)
        all_splits = set(dataset.shards_list.keys())
        all_groups = set(shard['group'] for split in Dataset.SPLITS
                         for shard in dataset.shards_list[split])
        all_parts = set(shard['part'] for split in Dataset.SPLITS
                        for shard in dataset.shards_list[split])
        prod = list(product(all_splits, all_groups, all_parts))
        for split, group, part in tqdm(prod):
            example_iterator = ExampleIterator(dataset_path=dataset_path,
                                               split=split,
                                               group=group,
                                               part=part)
            for example in tqdm(example_iterator, leave=False):
                print_stats.add_example(example=example,
                                        split=split,
                                        group=group,
                                        part=part)
        print_stats.print()
        return print_stats

    def add_example(self, example: Dict, split: str, group: int,
                    part: int) -> None:
        """Add another example.

        Args:
          example: A dictionary containing attack points and measurements
            (traces).
          split: The split this example belongs to.
          group: The group this example belongs to.
          part: The part this example belongs to.
        """
        # Group and part are unused now, but are there when we need more
        # granularity.
        _ = group
        _ = part
        for ap_name in self._attack_point_info:
            for ap_counter in self._split_ap_counters[split][ap_name]:
                ap_counter.update(attack_point=example[ap_name])
        for trace_name in self._measurements_info:
            for trace_stat in self._split_trace_stats[split][trace_name]:
                trace_stat.update(trace=example[trace_name])

    def print(self) -> None:
        """Print warnings from attack point checkers and statistics about
        traces."""
        print('Attack point check warnings:')
        for split in Dataset.SPLITS:
            print('')
            print(split)
            for ap_name in self._attack_point_info:
                for ap_count in self._split_ap_counters[split][ap_name]:
                    _ = APChecker(counts=ap_count.get_counts(),
                                  attack_point_name=ap_name)

        print('')
        print('-------------------------------------------')
        print('')

        print('Trace statistics:')
        for trace_name in self._measurements_info:
            print('')
            table = [[f'{trace_name} stats', *Dataset.SPLITS]]
            # Fill the statistic names.
            for trace_stat in self._split_trace_stats[
                    Dataset.SPLITS[0]][trace_name]:
                trace_stat_name = trace_stat.__class__.__name__
                table.append([trace_stat_name, *[0 for _ in Dataset.SPLITS]])
            # Fill the values.
            for i, split in enumerate(Dataset.SPLITS):
                # Check that we are filling the right column.
                assert split == table[0][i + 1]
                for j, trace_stat in enumerate(
                        self._split_trace_stats[split][trace_name]):
                    # Check that we are filling the right row.
                    assert trace_stat.__class__.__name__ == table[j + 1][0]
                    table[j + 1][i + 1] = trace_stat.result()
        print(tabulate(table, headers='firstrow'))
<commit_msg>Enable to take statistics with group granularity<commit_after>"""Print results of statistical checks and observations. Iterates only once
over the dataset. Implemented as a class to allow use during capture as well as
standalone."""

from __future__ import annotations
from collections import defaultdict
from itertools import product
from typing import Dict

from tabulate import tabulate
from tqdm import tqdm

from scaaml.io import Dataset
from .ap_counter import APCounter
from .ap_checker import APChecker
from .example_iterator import ExampleIterator
from .trace_stddev_of_stat import STDDEVofAVGofTraces
from .trace_stddev_of_stat import STDDEVofMAXofTraces
from .trace_stddev_of_stat import STDDEVofMINofTraces

class PrintStats:
    """Gather and print warnings from attack point checks and trace
    statistics.

    Example use:

    # On an existing dataset:
    print_stats = PrintStats.from_config(dataset_path='/mnt/path/to/dataset')

    # During capture:
    print_stats = PrintStats(measurements_info, attack_points_info)
    for example in example_iterator:
        print_stats.add_example(example=example,
                                split=split,
                                group=group,
                                part=part)
    print_stats.print()
    """
    def __init__(self, measurements_info: Dict,
                 attack_points_info: Dict) -> None:
        self._measurements_info = measurements_info
        self._attack_point_info = attack_points_info
        self._split_ap_counters = {
            split: {
                ap_name: [
                    APCounter(attack_point_info=ap_info),
                ]
                for ap_name, ap_info in attack_points_info.items()
            }
            for split in Dataset.SPLITS
        }

        self._all_trace_stats = lambda: [
            # Add trace statistics here.
            STDDEVofAVGofTraces(),
            STDDEVofMAXofTraces(),
            STDDEVofMINofTraces(),
        ]
        self._split_trace_stats = {
            split: {
                trace_name: self._all_trace_stats()
                for trace_name in measurements_info
            }
            for split in Dataset.SPLITS
        }
        self._split_trace_group_stats = {
            split: {
                trace_name: defaultdict(self._all_trace_stats)
                for trace_name in measurements_info
            }
            for split in Dataset.SPLITS
        }
        self._all_groups = set()

    @staticmethod
    def from_config(dataset_path: str) -> PrintStats:
        """Read a dataset and print all attack point check warnings and trace
        statistics.

        Args:
          dataset_path: Root path of the dataset.
        """
        dataset = Dataset.from_config(dataset_path)
        print_stats = PrintStats(measurements_info=dataset.measurements_info,
                                 attack_points_info=dataset.attack_points_info)
        all_splits = set(dataset.shards_list.keys())
        all_groups = set(shard['group'] for split in Dataset.SPLITS
                         for shard in dataset.shards_list[split])
        all_parts = set(shard['part'] for split in Dataset.SPLITS
                        for shard in dataset.shards_list[split])
        prod = list(product(all_splits, all_groups, all_parts))
        for split, group, part in tqdm(prod):
            example_iterator = ExampleIterator(dataset_path=dataset_path,
                                               split=split,
                                               group=group,
                                               part=part)
            for example in tqdm(example_iterator, leave=False):
                print_stats.add_example(example=example,
                                        split=split,
                                        group=group,
                                        part=part)
        print_stats.print()
        return print_stats

    def add_example(self, example: Dict, split: str, group: int,
                    part: int) -> None:
        """Add another example.

        Args:
          example: A dictionary containing attack points and measurements
            (traces).
          split: The split this example belongs to.
          group: The group this example belongs to.
          part: The part this example belongs to.
        """
        # Part is unused now, but is there when we need more granularity.
        _ = part
        self._all_groups.add(group)
        for ap_name in self._attack_point_info:
            for ap_counter in self._split_ap_counters[split][ap_name]:
                ap_counter.update(attack_point=example[ap_name])
        for trace_name in self._measurements_info:
            for trace_stat in self._split_trace_stats[split][trace_name]:
                trace_stat.update(trace=example[trace_name])
            for trace_stat in self._split_trace_group_stats[split][trace_name][
                    group]:
                trace_stat.update(trace=example[trace_name])

    def _print_attack_point_warnings(self):
        """Print warnings for attack points."""
        print('Attack point check warnings:')
        for split in Dataset.SPLITS:
            print('')
            print(split)
            for ap_name in self._attack_point_info:
                for ap_count in self._split_ap_counters[split][ap_name]:
                    # Add attack point checks here.
                    _ = APChecker(counts=ap_count.get_counts(),
                                  attack_point_name=ap_name)

    def _print_trace_statistics(self):
        """Print table of statistics. Rows are statistics, columns are splits.
        """
        print('Trace statistics:')
        class_name = lambda x: x.__class__.__name__
        for trace_name in self._measurements_info:
            print('')
            table = [[f'{trace_name} stats', *Dataset.SPLITS]]
            # Fill the statistic names.
            for trace_stat in self._all_trace_stats():
                trace_stat_name = class_name(trace_stat)
                table.append([trace_stat_name, *[0 for _ in Dataset.SPLITS]])
            # Fill the values.
            for i, split in enumerate(Dataset.SPLITS):
                # Check that we are filling the right column.
                assert split == table[0][i + 1]
                for j, trace_stat in enumerate(
                        self._split_trace_stats[split][trace_name]):
                    # Check that we are filling the right row.
                    assert class_name(trace_stat) == table[j + 1][0]
                    table[j + 1][i + 1] = trace_stat.result()
        print(tabulate(table, headers='firstrow'))

    def _print_trace_statistics_per_group(self):
        """Print a table of results for each statistics. Rows are groups,
        columns are splits."""
        print('Trace statistics per group:')
        class_name = lambda x: x.__class__.__name__
        for trace_name in self._measurements_info:
            for i, trace_stat in enumerate(self._all_trace_stats()):
                print('')
                print(f'{trace_name} {class_name(trace_stat)}:')
                table = [['group', *Dataset.SPLITS]]
                for group in self._all_groups:
                    table.append([str(group), *['NA' for _ in Dataset.SPLITS]])
                    for j, split in enumerate(Dataset.SPLITS):
                        if group in self._split_trace_group_stats[split][
                                trace_name]:
                            curr_trace_stat = self._split_trace_group_stats[
                                split][trace_name][group][i]
                            # Check that we are filling the right table.
                            assert class_name(curr_trace_stat) == class_name(
                                trace_stat)
                            table[-1][j + 1] = curr_trace_stat.result()
                print(tabulate(table, headers='firstrow'))

    @staticmethod
    def _print_separator():
        """Print separator between two blocks of warnings / statistics."""
        print('')
        print('-------------------------------------------')
        print('')

    def print(self) -> None:
        """Print warnings from attack point checkers and statistics about
        traces."""
        self._print_attack_point_warnings()
        PrintStats._print_separator()
        self._print_trace_statistics()
        PrintStats._print_separator()
        self._print_trace_statistics_per_group()
<|endoftext|>