<commit_before># Copyright (C) 2010-2011 Mathijs de Bruin <mathijs@mathijsfietst.nl>
#
# This file is part of django-shopkit.
#
# django-shopkit is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 2, or (at your option)
# any later version.
#
# This program 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.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

import logging
logger = logging.getLogger(__name__)

import datetime

from django.utils.translation import ugettext_lazy as _
from django.db import models

from shopkit.core.settings import MAX_NAME_LENGTH
from shopkit.core.managers import ActiveItemManager

"""
Generic abstract base classes for:
* :class:`Customers <shopkit.core.basemodels.AbstractCustomerBase>`

"""

class AbstractCustomerBase(models.Model):
    """ Abstract base class for customers of the shop. """

    class Meta:
        verbose_name = _('customer')
        verbose_name_plural = _('customers')
        abstract = True

    def get_all_orders(self):
        """ Get all orders by the customer """
        return self.order_set.all()

    def get_confirmed_orders(self):
        """
        Get all completed orders for this customer

        .. todo::
            We should consider adding a manager to :class:`OrderBase` which
            can filter on the completed states.
        """
        orders = self.get_all_orders()

        completed = orders.filter(confirmed=True)

        return completed

    def get_latest_order(self):
        """ Return the lastest confirmed order """
        completed = self.get_confirmed_orders()

        try:
            return completed.order_by('-date_added')[0]
        except IndexError:
            return None

class AbstractPricedItemBase(models.Model):
    """ Abstract base class for items with a price. This only contains
        a `get_price` dummy function yielding a NotImplementedError. An
        actual `price` field is contained in the `PricedItemBase` class.

        This is because we might want to get our prices somewhere else, ie.
        using some kind of algorithm, web API or database somewhere.
    """

    class Meta:
        abstract = True

    def get_price(self, **kwargs):
        """ Get price for the current product.

            This method _should_ be implemented in a subclass. """

        raise NotImplementedError

class QuantizedItemBase(models.Model):
    """ Abstract base class for items with a quantity field. """

    class Meta:
        abstract = True

    quantity = models.IntegerField(default=0, verbose_name=_('quantity'))
    """ Number of items of this kind. The default is 0: this is necessary so
    that we can add a certain quantity to a new object without knowing its
    initial value.
    """

class NamedItemBase(models.Model):
    """ Abstract base class for items with a name. """

    class Meta:
        abstract = True

    name = models.CharField(max_length=MAX_NAME_LENGTH,
                            verbose_name=_('name'))
    """ Name of this item. """

    def __unicode__(self):
        """ Returns the item's name. """

        return self.name

def get_next_ordering(qs, ordering_increase=10):
    max_ordering = qs.aggregate(models.Max('sort_order'))['sort_order__max']

    if max_ordering:
        return max_ordering+ordering_increase
    else:
        return ordering_increase

class OrderedItemBase(models.Model):
    """
    Abstract base class for items that have explicit ordering.
    """

    class Meta:
        abstract = True
        ordering = ('sort_order', )

    @classmethod
    def get_next_ordering(cls):
        return get_next_ordering(cls.objects.all())

    def save(self):
        """
        If no `sort_order` has been specified, make sure we calculate the
        it based on the highest available current `sort_order`.
        """

        if not self.sort_order:
            self.sort_order = self.get_next_ordering()

        super(OrderedItemBase, self).save()

    sort_order = models.PositiveIntegerField(
        verbose_name=('sort order'),
        unique=True, blank=True, db_index=True,
        help_text=_(
            'Change this to alter the order in which items are displayed.'
        )
    )

class OrderedInlineItemBase(models.Model):
    """
    This base class does what, actually, order_with_respect_to should do
    but (for now) doesn't implement very well: ordering of objects with
    a fk-relation to some class.

    As we do not know what the class with the respective relation is, it
    is important to note that something like the following is added::

        class MyOrderedInline(OrderedInlineItemBase):

            <related> = models.ForeignKey(RelatedModel)

            class Meta(OrderedInlineItemBase.Meta):
                unique_together = ('sort_order', '<related>')

            def get_related_ordering(self):
                return self.__class__.objects.filter(<related>=self.<related>)

        ... Or we could simply wait for the Django developers to fix
        `order_with_respect_to` once and for all. (Work in progress...
        See `Ticket #13 <http://code.djangoproject.com/ticket/13>`.)

    """

    class Meta:
        abstract = True
        ordering = ('sort_order', )

    def get_related_ordering(self):
        """
        Get a :class:`QuerySet <django.db.models.QuerySet.QuerySet`
        with related items to be considered for calculating the next
        `sort_order`.

        As we do not know in this base class what the related field(s)
        are, this raises a NotImplementedError. It should be subclassed
        with something like::

            return self.objects.filter(<related>=self.<related>)

        """
        raise NotImplementedError

    @staticmethod
    def get_next_ordering(related):
        """
        Get the next ordering based upon the :class:`QuerySet <django.db.models.QuerySet.QuerySet`
        with related items.
        """
        return get_next_ordering(related)

    def save(self):
        """
        If no `sort_order` has been specified, make sure we calculate the
        it based on the highest available current `sort_order`.
        """

        if not self.sort_order:
            related = self.get_related_ordering()
            self.sort_order = self.get_next_ordering(related)

            logger.debug(
                u'Generated sort_order %d for object %s',
                self.sort_order, self
            )

        super(OrderedInlineItemBase, self).save()

    sort_order = models.PositiveSmallIntegerField(
        verbose_name=('sort order'),
        blank=True, db_index=True,
        help_text=_(
            'Change this to alter the order in which items are displayed.'
        )
    )

class ActiveItemBase(models.Model):
    """
    Abstract base class for items which can be activated or deactivated.
    """

    class Meta:
        abstract = True

    active = models.BooleanField(verbose_name=_('active'),
                                 default=True, db_index=True)
    """ Whether the item is active in the frontend. """

class ActiveItemInShopBase(ActiveItemBase):
    """
    This is a subclass of :class:`ActiveItemBase` with an
    :class:`ActiveItemManager <shopkit.core.managers.ActiveItemManager>` called `in_shop`.

    The main purpose of this class is allowing for items to be enabled or
    disabled in the shop's frontend
    """

    class Meta:
        abstract = True

    objects = models.Manager()
    in_shop = ActiveItemManager()
    """ An instance of :class:`ActiveItemManager <shopkit.core.managers.ActiveItemManager>`,
        returning only items with `active=True`.
    """

class DatedItemBase(models.Model):
    """
    Item for which the add and modification date are automatically
    tracked.
    """

    class Meta:
        abstract = True
        ordering = ['-date_added', '-date_modified']
        get_latest_by = 'date_added'

    date_added = models.DateTimeField(auto_now_add=True,
                                      verbose_name=_('creation date'))
    date_modified = models.DateTimeField(auto_now=True,
                                         verbose_name=_('modification date'))

class PublishDateItemBase(models.Model):
    """ Item with a publish date. """

    class Meta:
        abstract = True
        ordering = ['-date_publish']
        get_latest_by = 'date_publish'

    date_publish = models.DateTimeField(default=datetime.datetime.now(),
                                        verbose_name=_('publication date'),
                                        db_index=True)

class NumberedOrderBase(models.Model):
    """ Base class for `Order` with invoice and order numbers. """

    class Meta:
        abstract = True

    invoice_number = models.SlugField(_('invoice number'), db_index=True,
                                      editable=False, max_length=255,
                                      unique=True, null=True, default=None)
    order_number = models.SlugField(_('order number'), db_index=True,
                                    editable=False, max_length=255,
                                    unique=True)

    def generate_invoice_number(self):
        """
        Generates an invoice number for the current order. Should be
        overridden in subclasses.
        """
        raise NotImplementedError

    def generate_order_number(self):
        """
        Generates an order number for the current order. Should be
        overridden in subclasses.
        """
        raise NotImplementedError

    def save(self, *args, **kwargs):
        """ Generate an order number upon saving the order. """

        if not self.order_number:
            self.order_number = self.generate_order_number()

            logger.debug('Generated order number %s for %s',
                         self.order_number, self.order_number)

        super(NumberedOrderBase, self).save(*args, **kwargs)

    def confirm(self):
        """ Make sure we set an invoice number upon order confirmation. """

        super(NumberedOrderBase, self).confirm()

        assert not self.invoice_number
        self.invoice_number = self.generate_invoice_number()

        logger.debug('Generated invoice number %s for %s',
                     self.invoice_number, self.order_number)

        # Note: This save operation might not me necessary
        self.save()
<commit_msg>Use clean() instead of save() for setting ordering<commit_after># Copyright (C) 2010-2011 Mathijs de Bruin <mathijs@mathijsfietst.nl>
#
# This file is part of django-shopkit.
#
# django-shopkit is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 2, or (at your option)
# any later version.
#
# This program 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.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

import logging
logger = logging.getLogger(__name__)

import datetime

from django.utils.translation import ugettext_lazy as _
from django.db import models

from shopkit.core.settings import MAX_NAME_LENGTH
from shopkit.core.managers import ActiveItemManager

"""
Generic abstract base classes for:
* :class:`Customers <shopkit.core.basemodels.AbstractCustomerBase>`

"""

class AbstractCustomerBase(models.Model):
    """ Abstract base class for customers of the shop. """

    class Meta:
        verbose_name = _('customer')
        verbose_name_plural = _('customers')
        abstract = True

    def get_all_orders(self):
        """ Get all orders by the customer """
        return self.order_set.all()

    def get_confirmed_orders(self):
        """
        Get all completed orders for this customer

        .. todo::
            We should consider adding a manager to :class:`OrderBase` which
            can filter on the completed states.
        """
        orders = self.get_all_orders()

        completed = orders.filter(confirmed=True)

        return completed

    def get_latest_order(self):
        """ Return the lastest confirmed order """
        completed = self.get_confirmed_orders()

        try:
            return completed.order_by('-date_added')[0]
        except IndexError:
            return None

class AbstractPricedItemBase(models.Model):
    """ Abstract base class for items with a price. This only contains
        a `get_price` dummy function yielding a NotImplementedError. An
        actual `price` field is contained in the `PricedItemBase` class.

        This is because we might want to get our prices somewhere else, ie.
        using some kind of algorithm, web API or database somewhere.
    """

    class Meta:
        abstract = True

    def get_price(self, **kwargs):
        """ Get price for the current product.

            This method _should_ be implemented in a subclass. """

        raise NotImplementedError

class QuantizedItemBase(models.Model):
    """ Abstract base class for items with a quantity field. """

    class Meta:
        abstract = True

    quantity = models.IntegerField(default=0, verbose_name=_('quantity'))
    """ Number of items of this kind. The default is 0: this is necessary so
    that we can add a certain quantity to a new object without knowing its
    initial value.
    """

class NamedItemBase(models.Model):
    """ Abstract base class for items with a name. """

    class Meta:
        abstract = True

    name = models.CharField(max_length=MAX_NAME_LENGTH,
                            verbose_name=_('name'))
    """ Name of this item. """

    def __unicode__(self):
        """ Returns the item's name. """

        return self.name

def get_next_ordering(qs, ordering_increase=10):
    max_ordering = qs.aggregate(models.Max('sort_order'))['sort_order__max']

    if max_ordering:
        return max_ordering+ordering_increase
    else:
        return ordering_increase

class OrderedItemBase(models.Model):
    """
    Abstract base class for items that have explicit ordering.
    """

    class Meta:
        abstract = True
        ordering = ('sort_order', )

    @classmethod
    def get_next_ordering(cls):
        return get_next_ordering(cls.objects.all())

    def clean(self):
        """
        If no `sort_order` has been specified, make sure we calculate the
        it based on the highest available current `sort_order`.
        """

        if not self.sort_order:
            self.sort_order = self.get_next_ordering()

    sort_order = models.PositiveIntegerField(
        verbose_name=('sort order'),
        unique=True, blank=True, db_index=True,
        help_text=_(
            'Change this to alter the order in which items are displayed.'
        )
    )

class OrderedInlineItemBase(models.Model):
    """
    This base class does what, actually, order_with_respect_to should do
    but (for now) doesn't implement very well: ordering of objects with
    a fk-relation to some class.

    As we do not know what the class with the respective relation is, it
    is important to note that something like the following is added::

        class MyOrderedInline(OrderedInlineItemBase):

            <related> = models.ForeignKey(RelatedModel)

            class Meta(OrderedInlineItemBase.Meta):
                unique_together = ('sort_order', '<related>')

            def get_related_ordering(self):
                return self.__class__.objects.filter(<related>=self.<related>)

        ... Or we could simply wait for the Django developers to fix
        `order_with_respect_to` once and for all. (Work in progress...
        See `Ticket #13 <http://code.djangoproject.com/ticket/13>`.)

    """

    class Meta:
        abstract = True
        ordering = ('sort_order', )

    def get_related_ordering(self):
        """
        Get a :class:`QuerySet <django.db.models.QuerySet.QuerySet`
        with related items to be considered for calculating the next
        `sort_order`.

        As we do not know in this base class what the related field(s)
        are, this raises a NotImplementedError. It should be subclassed
        with something like::

            return self.objects.filter(<related>=self.<related>)

        """
        raise NotImplementedError

    @staticmethod
    def get_next_ordering(related):
        """
        Get the next ordering based upon the :class:`QuerySet <django.db.models.QuerySet.QuerySet`
        with related items.
        """
        return get_next_ordering(related)

    def save(self):
        """
        If no `sort_order` has been specified, make sure we calculate the
        it based on the highest available current `sort_order`.
        """

        if not self.sort_order:
            related = self.get_related_ordering()
            self.sort_order = self.get_next_ordering(related)

            logger.debug(
                u'Generated sort_order %d for object %s',
                self.sort_order, self
            )

        super(OrderedInlineItemBase, self).save()

    sort_order = models.PositiveSmallIntegerField(
        verbose_name=('sort order'),
        blank=True, db_index=True,
        help_text=_(
            'Change this to alter the order in which items are displayed.'
        )
    )

class ActiveItemBase(models.Model):
    """
    Abstract base class for items which can be activated or deactivated.
    """

    class Meta:
        abstract = True

    active = models.BooleanField(verbose_name=_('active'),
                                 default=True, db_index=True)
    """ Whether the item is active in the frontend. """

class ActiveItemInShopBase(ActiveItemBase):
    """
    This is a subclass of :class:`ActiveItemBase` with an
    :class:`ActiveItemManager <shopkit.core.managers.ActiveItemManager>` called `in_shop`.

    The main purpose of this class is allowing for items to be enabled or
    disabled in the shop's frontend
    """

    class Meta:
        abstract = True

    objects = models.Manager()
    in_shop = ActiveItemManager()
    """ An instance of :class:`ActiveItemManager <shopkit.core.managers.ActiveItemManager>`,
        returning only items with `active=True`.
    """

class DatedItemBase(models.Model):
    """
    Item for which the add and modification date are automatically
    tracked.
    """

    class Meta:
        abstract = True
        ordering = ['-date_added', '-date_modified']
        get_latest_by = 'date_added'

    date_added = models.DateTimeField(auto_now_add=True,
                                      verbose_name=_('creation date'))
    date_modified = models.DateTimeField(auto_now=True,
                                         verbose_name=_('modification date'))

class PublishDateItemBase(models.Model):
    """ Item with a publish date. """

    class Meta:
        abstract = True
        ordering = ['-date_publish']
        get_latest_by = 'date_publish'

    date_publish = models.DateTimeField(default=datetime.datetime.now(),
                                        verbose_name=_('publication date'),
                                        db_index=True)

class NumberedOrderBase(models.Model):
    """ Base class for `Order` with invoice and order numbers. """

    class Meta:
        abstract = True

    invoice_number = models.SlugField(_('invoice number'), db_index=True,
                                      editable=False, max_length=255,
                                      unique=True, null=True, default=None)
    order_number = models.SlugField(_('order number'), db_index=True,
                                    editable=False, max_length=255,
                                    unique=True)

    def generate_invoice_number(self):
        """
        Generates an invoice number for the current order. Should be
        overridden in subclasses.
        """
        raise NotImplementedError

    def generate_order_number(self):
        """
        Generates an order number for the current order. Should be
        overridden in subclasses.
        """
        raise NotImplementedError

    def save(self, *args, **kwargs):
        """ Generate an order number upon saving the order. """

        if not self.order_number:
            self.order_number = self.generate_order_number()

            logger.debug('Generated order number %s for %s',
                         self.order_number, self.order_number)

        super(NumberedOrderBase, self).save(*args, **kwargs)

    def confirm(self):
        """ Make sure we set an invoice number upon order confirmation. """

        super(NumberedOrderBase, self).confirm()

        assert not self.invoice_number
        self.invoice_number = self.generate_invoice_number()

        logger.debug('Generated invoice number %s for %s',
                     self.invoice_number, self.order_number)

        # Note: This save operation might not me necessary
        self.save()
<|endoftext|>


<commit_before>"""Simple, Elegant Argument parsing.
@author: Fabrice Normandin
"""
from . import helpers, utils, wrappers
from .conflicts import ConflictResolution
from .help_formatter import SimpleHelpFormatter
from .helpers import (
    MutableField,
    Serializable,
    choice,
    field,
    flag,
    list_field,
    mutable_field,
    subgroups,
    subparsers,
)

from .parsing import (
    ArgumentParser,
    DashVariant,
    ParsingError,
    ArgumentGenerationMode,
    NestedMode,
)
from .utils import InconsistentArgumentError

from . import _version

__version__ = _version.get_versions()["version"]
<commit_msg>Re-add the __all__ in __init__.py<commit_after>"""Simple, Elegant Argument parsing.
@author: Fabrice Normandin
"""
from . import helpers, utils, wrappers
from .conflicts import ConflictResolution
from .help_formatter import SimpleHelpFormatter
from .helpers import (
    MutableField,
    Serializable,
    choice,
    field,
    flag,
    list_field,
    mutable_field,
    subgroups,
    subparsers,
)
from .parsing import ArgumentParser, DashVariant, ParsingError, ArgumentGenerationMode, NestedMode
from .utils import InconsistentArgumentError

__all__ = [
    "helpers",
    "utils",
    "wrappers",
    "ConflictResolution",
    "SimpleHelpFormatter",
    "MutableField",
    "Serializable",
    "choice",
    "field",
    "flag",
    "list_field",
    "mutable_field",
    "subgroups",
    "subparsers",
    "ArgumentParser",
    "DashVariant",
    "ParsingError",
    "ArgumentGenerationMode",
    "NestedMode",
    "InconsistentArgumentError",
]

from . import _version

__version__ = _version.get_versions()["version"]
<|endoftext|>


<commit_before># encoding: utf-8

from __future__ import absolute_import, unicode_literals

import re

from pathvalidate import unprintable_ascii_chars
from pathvalidate.error import (
    InvalidCharError,
    InvalidReservedNameError,
    NullNameError,
    ValidReservedNameError,
)

__SQLITE_VALID_RESERVED_KEYWORDS = [
    "ABORT",
    "ACTION",
    "AFTER",
    "ANALYZE",
    "ASC",
    "ATTACH",
    "BEFORE",
    "BEGIN",
    "BY",
    "CASCADE",
    "CAST",
    "COLUMN",
    "CONFLICT",
    "CROSS",
    "CURRENT_DATE",
    "CURRENT_TIME",
    "CURRENT_TIMESTAMP",
    "DATABASE",
    "DEFERRED",
    "DESC",
    "DETACH",
    "EACH",
    "END",
    "EXCLUSIVE",
    "EXPLAIN",
    "FAIL",
    "FOR",
    "FULL",
    "GLOB",
    "IGNORE",
    "IMMEDIATE",
    "INDEXED",
    "INITIALLY",
    "INNER",
    "INSTEAD",
    "KEY",
    "LEFT",
    "LIKE",
    "MATCH",
    "NATURAL",
    "NO",
    "OF",
    "OFFSET",
    "OUTER",
    "PLAN",
    "PRAGMA",
    "QUERY",
    "RAISE",
    "RECURSIVE",
    "REGEXP",
    "REINDEX",
    "RELEASE",
    "RENAME",
    "REPLACE",
    "RESTRICT",
    "RIGHT",
    "ROLLBACK",
    "ROW",
    "SAVEPOINT",
    "TEMP",
    "TEMPORARY",
    "TRIGGER",
    "VACUUM",
    "VIEW",
    "VIRTUAL",
    "WITH",
    "WITHOUT",
]
__SQLITE_INVALID_RESERVED_KEYWORDS = [
    "ADD",
    "ALL",
    "ALTER",
    "AND",
    "AS",
    "AUTOINCREMENT",
    "BETWEEN",
    "CASE",
    "CHECK",
    "COLLATE",
    "COMMIT",
    "CONSTRAINT",
    "CREATE",
    "DEFAULT",
    "DEFERRABLE",
    "DELETE",
    "DISTINCT",
    "DROP",
    "ELSE",
    "ESCAPE",
    "EXCEPT",
    "EXISTS",
    "FOREIGN",
    "FROM",
    "GROUP",
    "HAVING",
    "IN",
    "INDEX",
    "INSERT",
    "INTERSECT",
    "INTO",
    "IS",
    "ISNULL",
    "JOIN",
    "LIMIT",
    "NOT",
    "NOTNULL",
    "NULL",
    "ON",
    "OR",
    "ORDER",
    "PRIMARY",
    "REFERENCES",
    "SELECT",
    "SET",
    "TABLE",
    "THEN",
    "TO",
    "TRANSACTION",
    "UNION",
    "UNIQUE",
    "UPDATE",
    "USING",
    "VALUES",
    "WHEN",
    "WHERE",
]

__SQLITE_VALID_RESERVED_KEYWORDS_TABLE = __SQLITE_VALID_RESERVED_KEYWORDS
__SQLITE_INVALID_RESERVED_KEYWORDS_TABLE = __SQLITE_INVALID_RESERVED_KEYWORDS + ["IF"]

__SQLITE_VALID_RESERVED_KEYWORDS_ATTR = __SQLITE_VALID_RESERVED_KEYWORDS + ["IF"]
__SQLITE_INVALID_RESERVED_KEYWORDS_ATTR = __SQLITE_INVALID_RESERVED_KEYWORDS

__RE_INVALID_CHARS = re.compile(
    "[{:s}]".format(re.escape("".join(unprintable_ascii_chars))), re.UNICODE
)

def validate_sqlite_table_name(name):
    """
    :param str name: Name to validate.
    :raises pathvalidate.NullNameError: If the ``name`` is empty.
    :raises pathvalidate.InvalidCharError:
        If the ``name`` includes unprintable character(s).
    :raises pathvalidate.InvalidReservedNameError:
        |raises_sqlite_keywords|
        And invalid as a table name.
    :raises pathvalidate.ValidReservedNameError:
        |raises_sqlite_keywords|
        However, valid as a table name.
    """

    if not name:
        raise NullNameError("null name")

    if __RE_INVALID_CHARS.search(name):
        raise InvalidCharError("unprintable character found")

    if name.upper() in __SQLITE_INVALID_RESERVED_KEYWORDS_TABLE:
        raise InvalidReservedNameError("'{:s}' is a reserved keyword by sqlite".format(name))

    if name.upper() in __SQLITE_VALID_RESERVED_KEYWORDS_TABLE:
        raise ValidReservedNameError("'{:s}' is a reserved keyword by sqlite".format(name))

def validate_sqlite_attr_name(name):
    """
    :param str name: Name to validate.
    :raises pathvalidate.NullNameError: If the ``name`` is empty.
    :raises pathvalidate.InvalidCharError:
        If the ``name`` includes unprintable character(s).
    :raises pathvalidate.InvalidReservedNameError:
        |raises_sqlite_keywords|
        And invalid as an attribute name.
    :raises pathvalidate.ValidReservedNameError:
        |raises_sqlite_keywords|
        However, valid as an attribute name.
    """

    if not name:
        raise NullNameError("null name")

    if __RE_INVALID_CHARS.search(name):
        raise InvalidCharError("unprintable character found")

    if name.upper() in __SQLITE_INVALID_RESERVED_KEYWORDS_ATTR:
        raise InvalidReservedNameError("'{}' is a reserved keyword by sqlite".format(name))

    if name.upper() in __SQLITE_VALID_RESERVED_KEYWORDS_ATTR:
        raise ValidReservedNameError("'{}' is a reserved keyword by sqlite".format(name))
<commit_msg>Refactoring<commit_after># encoding: utf-8

from __future__ import absolute_import, unicode_literals

import re

from pathvalidate import unprintable_ascii_chars
from pathvalidate.error import (
    InvalidCharError,
    InvalidReservedNameError,
    NullNameError,
    ValidReservedNameError,
)

__SQLITE_VALID_RESERVED_KEYWORDS = [
    "ABORT",
    "ACTION",
    "AFTER",
    "ANALYZE",
    "ASC",
    "ATTACH",
    "BEFORE",
    "BEGIN",
    "BY",
    "CASCADE",
    "CAST",
    "COLUMN",
    "CONFLICT",
    "CROSS",
    "CURRENT_DATE",
    "CURRENT_TIME",
    "CURRENT_TIMESTAMP",
    "DATABASE",
    "DEFERRED",
    "DESC",
    "DETACH",
    "EACH",
    "END",
    "EXCLUSIVE",
    "EXPLAIN",
    "FAIL",
    "FOR",
    "FULL",
    "GLOB",
    "IGNORE",
    "IMMEDIATE",
    "INDEXED",
    "INITIALLY",
    "INNER",
    "INSTEAD",
    "KEY",
    "LEFT",
    "LIKE",
    "MATCH",
    "NATURAL",
    "NO",
    "OF",
    "OFFSET",
    "OUTER",
    "PLAN",
    "PRAGMA",
    "QUERY",
    "RAISE",
    "RECURSIVE",
    "REGEXP",
    "REINDEX",
    "RELEASE",
    "RENAME",
    "REPLACE",
    "RESTRICT",
    "RIGHT",
    "ROLLBACK",
    "ROW",
    "SAVEPOINT",
    "TEMP",
    "TEMPORARY",
    "TRIGGER",
    "VACUUM",
    "VIEW",
    "VIRTUAL",
    "WITH",
    "WITHOUT",
]
__SQLITE_INVALID_RESERVED_KEYWORDS = [
    "ADD",
    "ALL",
    "ALTER",
    "AND",
    "AS",
    "AUTOINCREMENT",
    "BETWEEN",
    "CASE",
    "CHECK",
    "COLLATE",
    "COMMIT",
    "CONSTRAINT",
    "CREATE",
    "DEFAULT",
    "DEFERRABLE",
    "DELETE",
    "DISTINCT",
    "DROP",
    "ELSE",
    "ESCAPE",
    "EXCEPT",
    "EXISTS",
    "FOREIGN",
    "FROM",
    "GROUP",
    "HAVING",
    "IN",
    "INDEX",
    "INSERT",
    "INTERSECT",
    "INTO",
    "IS",
    "ISNULL",
    "JOIN",
    "LIMIT",
    "NOT",
    "NOTNULL",
    "NULL",
    "ON",
    "OR",
    "ORDER",
    "PRIMARY",
    "REFERENCES",
    "SELECT",
    "SET",
    "TABLE",
    "THEN",
    "TO",
    "TRANSACTION",
    "UNION",
    "UNIQUE",
    "UPDATE",
    "USING",
    "VALUES",
    "WHEN",
    "WHERE",
]

__SQLITE_VALID_RESERVED_KEYWORDS_TABLE = __SQLITE_VALID_RESERVED_KEYWORDS
__SQLITE_INVALID_RESERVED_KEYWORDS_TABLE = __SQLITE_INVALID_RESERVED_KEYWORDS + ["IF"]

__SQLITE_VALID_RESERVED_KEYWORDS_ATTR = __SQLITE_VALID_RESERVED_KEYWORDS + ["IF"]
__SQLITE_INVALID_RESERVED_KEYWORDS_ATTR = __SQLITE_INVALID_RESERVED_KEYWORDS

__RE_INVALID_CHARS = re.compile(
    "[{:s}]".format(re.escape("".join(unprintable_ascii_chars))), re.UNICODE
)

def validate_sqlite_table_name(name):
    """
    :param str name: Name to validate.
    :raises pathvalidate.NullNameError: If the ``name`` is empty.
    :raises pathvalidate.InvalidCharError:
        If the ``name`` includes unprintable character(s).
    :raises pathvalidate.InvalidReservedNameError:
        |raises_sqlite_keywords|
        And invalid as a table name.
    :raises pathvalidate.ValidReservedNameError:
        |raises_sqlite_keywords|
        However, valid as a table name.
    """

    if not name:
        raise NullNameError("null name")

    if __RE_INVALID_CHARS.search(name):
        raise InvalidCharError("unprintable character found")

    name = name.upper()

    if name in __SQLITE_INVALID_RESERVED_KEYWORDS_TABLE:
        raise InvalidReservedNameError("'{:s}' is a reserved keyword by sqlite".format(name))

    if name in __SQLITE_VALID_RESERVED_KEYWORDS_TABLE:
        raise ValidReservedNameError("'{:s}' is a reserved keyword by sqlite".format(name))

def validate_sqlite_attr_name(name):
    """
    :param str name: Name to validate.
    :raises pathvalidate.NullNameError: If the ``name`` is empty.
    :raises pathvalidate.InvalidCharError:
        If the ``name`` includes unprintable character(s).
    :raises pathvalidate.InvalidReservedNameError:
        |raises_sqlite_keywords|
        And invalid as an attribute name.
    :raises pathvalidate.ValidReservedNameError:
        |raises_sqlite_keywords|
        However, valid as an attribute name.
    """

    if not name:
        raise NullNameError("null name")

    if __RE_INVALID_CHARS.search(name):
        raise InvalidCharError("unprintable character found")

    name = name.upper()

    if name in __SQLITE_INVALID_RESERVED_KEYWORDS_ATTR:
        raise InvalidReservedNameError("'{}' is a reserved keyword by sqlite".format(name))

    if name in __SQLITE_VALID_RESERVED_KEYWORDS_ATTR:
        raise ValidReservedNameError("'{}' is a reserved keyword by sqlite".format(name))
<|endoftext|>


<commit_before>"""
Tests for the ``simpy.events.Process``.

"""
# Pytest gets the parameters "env" and "log" from the *conftest.py* file
import pytest

from simpy import Interrupt

def test_start_non_process(env):
    """Check that you cannot start a normal function."""
    def foo():
        pass

    pytest.raises(ValueError, env.process, foo)

def test_get_state(env):
    """A process is alive until it's generator has not terminated."""
    def pem_a(env):
        yield env.timeout(3)

    def pem_b(env, pem_a):
        yield env.timeout(1)
        assert pem_a.is_alive

        yield env.timeout(3)
        assert not pem_a.is_alive

    proc_a = env.process(pem_a(env))
    env.process(pem_b(env, proc_a))
    env.run()

def test_target(env):
    def pem(env, event):
        yield event

    event = env.timeout(5)
    proc = env.process(pem(env, event))

    # Wait until "proc" is initialized and yielded the event
    while env.peek() < 5:
        env.step()
    assert proc.target is event
    proc.interrupt()

def test_wait_for_proc(env):
    """A process can wait until another process finishes."""
    def finisher(env):
        yield env.timeout(5)

    def waiter(env, finisher):
        proc = env.process(finisher(env))
        yield proc  # Waits until "proc" finishes

        assert env.now == 5

    env.process(waiter(env, finisher))
    env.run()

def test_exit(env):
    """Processes can set a return value via an ``exit()`` function,
    comparable to ``sys.exit()``.

    """
    def child(env):
        yield env.timeout(1)
        env.exit(env.now)

    def parent(env):
        result1 = yield env.process(child(env))
        result2 = yield env.process(child(env))

        assert [result1, result2] == [1, 2]

    env.process(parent(env))
    env.run()

@pytest.mark.skipif('sys.version_info[:2] < (3, 3)')
def test_return_value(env):
    """Processes can set a return value."""
    # Python < 3.2 would raise a SyntaxError if this was real code ...
    code = """def child(env):
        yield env.timeout(1)
        return env.now
    """
    globs, locs = {}, {}
    code = compile(code, '<string>', 'exec')
    eval(code, globs, locs)
    child = locs['child']

    def parent(env):
        result1 = yield env.process(child(env))
        result2 = yield env.process(child(env))

        assert [result1, result2] == [1, 2]

    env.process(parent(env))
    env.run()

def test_child_exception(env):
    """A child catches an exception and sends it to its parent."""
    def child(env):
        try:
            yield env.timeout(1)
            raise RuntimeError('Onoes!')
        except RuntimeError as err:
            env.exit(err)

    def parent(env):
        result = yield env.process(child(env))
        assert isinstance(result, Exception)

    env.process(parent(env))
    env.run()

def test_interrupted_join(env):
    """Tests that interrupts are raised while the victim is waiting for
    another process. The victim should get unregistered from the other
    process.

    """
    def interruptor(env, process):
        yield env.timeout(1)
        process.interrupt()

    def child(env):
        yield env.timeout(2)

    def parent(env):
        child_proc = env.process(child(env))
        try:
            yield child_proc
            pytest.fail('Did not receive an interrupt.')
        except Interrupt:
            assert env.now == 1
            assert child_proc.is_alive

            # We should not get resumed when child terminates.
            yield env.timeout(5)
            assert env.now == 6

    parent_proc = env.process(parent(env))
    env.process(interruptor(env, parent_proc))
    env.run()

def test_interrupted_join_and_rejoin(env):
    """Tests that interrupts are raised while the victim is waiting for
    another process. The victim tries to join again.

    """
    def interruptor(env, process):
        yield env.timeout(1)
        process.interrupt()

    def child(env):
        yield env.timeout(2)

    def parent(env):
        child_proc = env.process(child(env))
        try:
            yield child_proc
            pytest.fail('Did not receive an interrupt.')
        except Interrupt:
            assert env.now == 1
            assert child_proc.is_alive

            yield child_proc
            assert env.now == 2

    parent_proc = env.process(parent(env))
    env.process(interruptor(env, parent_proc))
    env.run()

def test_unregister_after_interrupt(env):
    """If a process is interrupted while waiting for another one, it
    should be unregistered from that process.

    """
    def interruptor(env, process):
        yield env.timeout(1)
        process.interrupt()

    def child(env):
        yield env.timeout(2)

    def parent(env):
        child_proc = env.process(child(env))
        try:
            yield child_proc
            pytest.fail('Did not receive an interrupt.')
        except Interrupt:
            assert env.now == 1
            assert child_proc.is_alive

        yield env.timeout(2)
        assert env.now == 3
        assert not child_proc.is_alive

    parent_proc = env.process(parent(env))
    env.process(interruptor(env, parent_proc))
    env.run()

def test_error_and_interrupted_join(env):
    def child_a(env, process):
        process.interrupt()
        env.exit()
        yield  # Dummy yield

    def child_b(env):
        raise AttributeError('spam')
        yield  # Dummy yield

    def parent(env):
        env.process(child_a(env, env.active_process))
        b = env.process(child_b(env))

        try:
            yield b
        # This interrupt unregisters me from b so I won't receive its
        # AttributeError
        except Interrupt:
            pass

        yield env.timeout(0)

    env.process(parent(env))
    pytest.raises(AttributeError, env.run)
<commit_msg>Remove duplicate test and tweak documentation.<commit_after>"""
Tests for the ``simpy.events.Process``.

"""
# Pytest gets the parameters "env" and "log" from the *conftest.py* file
import pytest

from simpy import Interrupt

def test_start_non_process(env):
    """Check that you cannot start a normal function."""
    def foo():
        pass

    pytest.raises(ValueError, env.process, foo)

def test_get_state(env):
    """A process is alive until it's generator has not terminated."""
    def pem_a(env):
        yield env.timeout(3)

    def pem_b(env, pem_a):
        yield env.timeout(1)
        assert pem_a.is_alive

        yield env.timeout(3)
        assert not pem_a.is_alive

    proc_a = env.process(pem_a(env))
    env.process(pem_b(env, proc_a))
    env.run()

def test_target(env):
    def pem(env, event):
        yield event

    event = env.timeout(5)
    proc = env.process(pem(env, event))

    # Wait until "proc" is initialized and yielded the event
    while env.peek() < 5:
        env.step()
    assert proc.target is event
    proc.interrupt()

def test_wait_for_proc(env):
    """A process can wait until another process finishes."""
    def finisher(env):
        yield env.timeout(5)

    def waiter(env, finisher):
        proc = env.process(finisher(env))
        yield proc  # Waits until "proc" finishes

        assert env.now == 5

    env.process(waiter(env, finisher))
    env.run()

def test_exit(env):
    """Processes can set a return value via an ``exit()`` function,
    comparable to ``sys.exit()``.

    """
    def child(env):
        yield env.timeout(1)
        env.exit(env.now)

    def parent(env):
        result1 = yield env.process(child(env))
        result2 = yield env.process(child(env))

        assert [result1, result2] == [1, 2]

    env.process(parent(env))
    env.run()

@pytest.mark.skipif('sys.version_info[:2] < (3, 3)')
def test_return_value(env):
    """Processes can set a return value."""
    # Python < 3.2 would raise a SyntaxError if this was real code ...
    code = """def child(env):
        yield env.timeout(1)
        return env.now
    """
    globs, locs = {}, {}
    code = compile(code, '<string>', 'exec')
    eval(code, globs, locs)
    child = locs['child']

    def parent(env):
        result1 = yield env.process(child(env))
        result2 = yield env.process(child(env))

        assert [result1, result2] == [1, 2]

    env.process(parent(env))
    env.run()

def test_child_exception(env):
    """A child catches an exception and sends it to its parent."""
    def child(env):
        try:
            yield env.timeout(1)
            raise RuntimeError('Onoes!')
        except RuntimeError as err:
            env.exit(err)

    def parent(env):
        result = yield env.process(child(env))
        assert isinstance(result, Exception)

    env.process(parent(env))
    env.run()

def test_interrupted_join(env):
    """Interrupts remove a process from the callbacks of its target."""

    def interruptor(env, process):
        yield env.timeout(1)
        process.interrupt()

    def child(env):
        yield env.timeout(2)

    def parent(env):
        child_proc = env.process(child(env))
        try:
            yield child_proc
            pytest.fail('Did not receive an interrupt.')
        except Interrupt:
            assert env.now == 1
            assert child_proc.is_alive

            # We should not get resumed when child terminates.
            yield env.timeout(5)
            assert env.now == 6

    parent_proc = env.process(parent(env))
    env.process(interruptor(env, parent_proc))
    env.run()

def test_interrupted_join_and_rejoin(env):
    """Tests that interrupts are raised while the victim is waiting for
    another process. The victim tries to join again.

    """
    def interruptor(env, process):
        yield env.timeout(1)
        process.interrupt()

    def child(env):
        yield env.timeout(2)

    def parent(env):
        child_proc = env.process(child(env))
        try:
            yield child_proc
            pytest.fail('Did not receive an interrupt.')
        except Interrupt:
            assert env.now == 1
            assert child_proc.is_alive

            yield child_proc
            assert env.now == 2

    parent_proc = env.process(parent(env))
    env.process(interruptor(env, parent_proc))
    env.run()

def test_error_and_interrupted_join(env):
    def child_a(env, process):
        process.interrupt()
        env.exit()
        yield  # Dummy yield

    def child_b(env):
        raise AttributeError('spam')
        yield  # Dummy yield

    def parent(env):
        env.process(child_a(env, env.active_process))
        b = env.process(child_b(env))

        try:
            yield b
        # This interrupt unregisters me from b so I won't receive its
        # AttributeError
        except Interrupt:
            pass

        yield env.timeout(0)

    env.process(parent(env))
    pytest.raises(AttributeError, env.run)
<|endoftext|>


<commit_before>import numpy as np
import warnings

need_state_equation_function_msg = ("if dim_state > 0, DynamicalSystem must"
                                    + " have a state_equation_function")

need_output_equation_function_msg = ("if dim_state == 0, DynamicalSystem must"
                                     + " have an output_equation_function")

zero_dim_output_msg = "A DynamicalSystem must provide an output"

def full_state_output(*args):
    """
    A drop-in ``output_equation_function`` for stateful systems that provide
    output the full state directly.
    """
    return np.r_[args][1:]

class DynamicalSystem(object):
    """
    A dynamical system which models systems of the form::

        xdot(t) = state_equation_function(t,x,u)
        y(t) = output_equation_function(t,x)

    or::

        y(t) = output_equation_function(t,u)

    These could also represent discrete-time systems, in which case xdot(t)
    represents x[k+1].

    This can also model discontinuous systems. Discontinuities must occur on
    zero-crossings of the ``event_equation_function``, which take the same
    arguments as ``output_equation_function``, depending on ``dim_state``.
    At the zero-crossing, ``update_equation_function`` is called with the same
    arguments. If ``dim_state`` > 0, the return value of
    ``update_equation_function`` is used as the state of the system immediately
    after the discontinuity.
    """
    def __init__(self, state_equation_function=None,
                 output_equation_function=None, event_equation_function=None,
                 update_equation_function=None, dim_state=0, dim_input=0,
                 dim_output=0, dt=0, initial_condition=None):
        """
        Parameters
        ----------
        state_equation_function : callable, optional
            The derivative (or update equation) of the system state. Not needed
            if ``dim_state`` is zero.
        output_equation_function : callable, optional
            The output equation of the system. A system must have an
            ``output_equation_function``. If not set, uses full state output.
        event_equation_function : callable, optional
            The function whose output determines when discontinuities occur.
        update_equation_function : callable, optional
            The function called when a discontinuity occurs.
        dim_state : int, optional
            Dimension of the system state. Optional, defaults to 0.
        dim_input : int, optional
            Dimension of the system input. Optional, defaults to 0.
        dim_output : int, optional
            Dimension of the system output. Optional, defaults to dim_state.
        dt : float, optional
            Sample rate of the system. Optional, defaults to 0 representing a
            continuous time system.
        initial_condition : array_like of numerical values, optional
            Array or Matrix used as the initial condition of the system.
            Defaults to zeros of the same dimension as the state.
        """
        self.dim_state = dim_state
        self.dim_input = dim_input
        self.dim_output = dim_output or dim_state

        self.state_equation_function = state_equation_function

        self.output_equation_function = (
            full_state_output
            if output_equation_function is None and self.dim_state > 0
            else output_equation_function
        )

        self.event_equation_function = event_equation_function
        self.update_equation_function = update_equation_function
        self.initial_condition = initial_condition
        self.dt = dt

        self.validate()

    @property
    def initial_condition(self):
        return (self._initial_condition
                if self._initial_condition is not None
                else np.zeros(self.dim_state))

    @initial_condition.setter
    def initial_condition(self, initial_condition):
        self._initial_condition = initial_condition

    def prepare_to_integrate(self):
        return

    def validate(self):
        if self.dim_output == 0:
            raise ValueError(zero_dim_output_msg)

        if (self.dim_state > 0
                and getattr(self, 'state_equation_function', None) is None):
            raise ValueError(need_state_equation_function_msg)

        if (self.dim_state == 0
                and getattr(self, 'output_equation_function', None) is None):
            raise ValueError(need_output_equation_function_msg)

def SystemFromCallable(incallable, dim_input, dim_output, dt=0):
    """
    Construct a memoryless system from a callable.

    Parameters
    ----------
    incallable : callable
        Function to use as the output_equation_function. Should have signature
        (t, u) if dim_input > 0 or (t) if dim_input = 0.
    dim_input : int
        Dimension of input.
    dim_output : int
        Dimension of output.
    """
    system = DynamicalSystem(output_equation_function=incallable,
                             dim_input=dim_input, dim_output=dim_output, dt=dt)
    return system

class SwitchedSystem(DynamicalSystem):
    """
    Provides a useful pattern for discontinuous systems where the state and
    output equations change depending on the value of a function of the state
    and/or input (``event_variable_equation_function``). Most of the usefulness
    comes from constructing the ``event_equation_function`` with a Bernstein
    basis polynomial with roots at the boundaries. This class also provides
    logic for outputting the correct state and output equation based on the
    ``event_variable_equation_function`` value.
    """
    def __init__(self, state_equations_functions=None,
                 output_equations_functions=None,
                 event_variable_equation_function=None, event_bounds=None,
                 state_update_equation_function=None, dim_state=0, dim_input=0,
                 dim_output=0, dt=0, initial_condition=None):
        """
        Parameters
        ----------
        state_equations_functions : array_like of callables, optional
            The derivative (or update equation) of the system state. Not needed
            if ``dim_state`` is zero. The array indexes the
            event-state and should be one more than the number of event bounds.
            This should also be indexed to match the boundaries (i.e., the
            first function is used when the event variable is below the first
            event_bounds value). If only one callable is provided, the callable
            is used in each condition.
        output_equations_functions : array_like of callables, optional
            The output equation of the system. A system must have an
            ``output_equation_function``. If not set, uses full state output.
            The array indexes the event-state and should be one more than the
            number of event bounds. This should also be indexed to match the
            boundaries (i.e., the first function is used when the event
            variable is below the first event_bounds value). If only one
            callable is provided, the callable is used in each condition.
        event_variable_equation_function : callable
            When the output of this function crosses the values in
            ``event_bounds``, a discontuity event occurs.
        event_bounds : array_like of floats
            Defines the boundary points the trigger discontinuity events based
            on the output of ``event_variable_equation_function``.
        state_update_equation_function : callable, optional
            When an event occurs, the state update equation function is called
            to determine the state update. If not set, uses full state output,
            so the state is not changed upon a zero-crossing of the event
            variable function.
        dim_state : int, optional
            Dimension of the system state. Optional, defaults to 0.
        dim_input : int, optional
            Dimension of the system input. Optional, defaults to 0.
        dim_output : int, optional
            Dimension of the system output. Optional, defaults to dim_state.
        dt : float, optional
            Sample rate of the system. Optional, defaults to 0 representing a
            continuous time system.
        """
        self.dim_state = dim_state
        self.dim_input = dim_input
        self.dim_output = dim_output or dim_state
        self.event_bounds = event_bounds

        self.state_equations_functions = np.empty(self.n_conditions,
                                                  dtype=object)
        self.state_equations_functions[:] = state_equations_functions

        self.output_equations_functions = np.empty(self.n_conditions,
                                                   dtype=object)
        self.output_equations_functions[:] = (
            full_state_output
            if output_equations_functions is None and self.dim_state > 0
            else output_equations_functions
        )

        self.event_variable_equation_function = \
            event_variable_equation_function

        self.state_update_equation_function = (
            state_update_equation_function or
            full_state_output
        )

        self.initial_condition = initial_condition or np.zeros(dim_state)
        self.dt = dt

        self.validate()

    def validate(self):
        super().validate()

        if (self.dim_state > 0
                and np.any(np.equal(self.state_equations_functions, None))):
            raise ValueError(need_state_equation_function_msg)

        if (self.dim_state == 0
                and np.any(np.equal(self.output_equations_functions, None))):
            raise ValueError(need_output_equation_function_msg)

        if self.event_variable_equation_function is None:
            raise ValueError("A SwitchedSystem requires " +
                             "event_variable_equation_function")

    @property
    def event_bounds(self):
        return self._event_bounds

    @event_bounds.setter
    def event_bounds(self, event_bounds):
        if event_bounds is None:
            raise ValueError("A SwitchedSystem requires event_bounds")
        self._event_bounds = np.array(event_bounds).reshape(1, -1)
        self.n_conditions = self._event_bounds.size + 1
        if self.n_conditions == 2:
            self.event_bounds_range = 1
        else:
            self.event_bounds_range = np.diff(self.event_bounds[0, [0, -1]])

    def output_equation_function(self, *args):
        return self.output_equations_functions[self.condition_idx](*args)

    def state_equation_function(self, *args):
        return self.state_equations_functions[self.condition_idx](*args)

    def event_equation_function(self, *args):
        event_var = self.event_variable_equation_function(*args)
        return np.prod(
            (self.event_bounds_range-self.event_bounds)*event_var -
            self.event_bounds*(self.event_bounds_range - event_var),
            axis=1
        )

    def update_equation_function(self, *args):
        event_var = self.event_variable_equation_function(*args)
        if self.condition_idx is None:
            self.condition_idx = np.where(np.all(np.r_[
                    np.c_[[[True]], event_var >= self.event_bounds],
                    np.c_[event_var <= self.event_bounds, [[True]]]
                    ], axis=0))[0][0]
        else:
            sq_dist = (event_var - self.event_bounds)**2
            crossed_root_idx = np.where(sq_dist == np.min(sq_dist))[1][0]
            if crossed_root_idx == self.condition_idx:
                self.condition_idx += 1
            elif crossed_root_idx == self.condition_idx-1:
                self.condition_idx -= 1
            else:
                warnings.warn("SwitchedSystem did not cross a neighboring " +
                              "boundary. This may indicate an integration " +
                              "error. Continuing without updating " +
                              "condition_idx", UserWarning)
        return self.state_update_equation_function(*args)

    def prepare_to_integrate(self):
        self.condition_idx = None

class LTISystem(DynamicalSystem):
    """
    A linear, time-invariant system.
    """
    def __init__(self, *args, initial_condition=None, dt=0):
        """
        Construct an LTI system with the following input formats:

        1. state matrix A, input matrix B, output matrix C for systems with
           state::

              dx_dt = Ax + Bu
              y = Hx

        2. state matrix A, input matrix B for systems with state, assume full
           state output::

              dx_dt = Ax + Bu
              y = Ix

        3. gain matrix K for systems without state::

              y = Kx

        The matrices should be numeric arrays of consistent shape.
        """
        self.dt = dt

        if len(args) not in (1, 2, 3):
            raise ValueError("LTI system expects 1, 2, or 3 args")

        # TODO: setup jacobian functions
        if len(args) == 1:
            self.gain_matrix = gain_matrix = args[0]
            self.dim_input = self.gain_matrix.shape[1] if len(gain_matrix.shape) > 1 else 1
            self.dim_output = self.gain_matrix.shape[0]
            self.dim_state = 0
            self.initial_condition = np.zeros(self.dim_state)
            self.output_equation_function = lambda t, x: (gain_matrix@x).reshape(-1)
            return

        if len(args) == 2:
            state_matrix, input_matrix = args
            output_matrix = np.eye(state_matrix.shape[0])

        elif len(args) == 3:
            state_matrix, input_matrix, output_matrix = args

        if len(input_matrix.shape) == 1:
            input_matrix = input_matrix.reshape(-1, 1)

        self.dim_state = state_matrix.shape[0]
        self.dim_input = input_matrix.shape[1]
        self.dim_output = output_matrix.shape[0]

        self.state_matrix = state_matrix
        self.input_matrix = input_matrix
        self.output_matrix = output_matrix

        self.initial_condition = initial_condition or np.zeros(self.dim_state)
        self.state_equation_function = lambda t, x, u: (state_matrix@x + input_matrix@u).reshape(-1)
        self.output_equation_function = lambda t, x: (output_matrix@x).reshape(-1)

        self.validate()

    def validate(self):
        super().validate()
        if self.dim_state:
            assert self.state_matrix.shape[1] == self.dim_state
            assert self.input_matrix.shape[0] == self.dim_state
            assert self.output_matrix.shape[1] == self.dim_state

    @property
    def data(self):
        if self.dim_state:
            return self.state_matrix, self.input_matrix, self.output_matrix
        else:
            return self.gain_matrix
<commit_msg>add aliases for LTI system matrices<commit_after>import numpy as np
import warnings

need_state_equation_function_msg = ("if dim_state > 0, DynamicalSystem must"
                                    + " have a state_equation_function")

need_output_equation_function_msg = ("if dim_state == 0, DynamicalSystem must"
                                     + " have an output_equation_function")

zero_dim_output_msg = "A DynamicalSystem must provide an output"

def full_state_output(*args):
    """
    A drop-in ``output_equation_function`` for stateful systems that provide
    output the full state directly.
    """
    return np.r_[args][1:]

class DynamicalSystem(object):
    """
    A dynamical system which models systems of the form::

        xdot(t) = state_equation_function(t,x,u)
        y(t) = output_equation_function(t,x)

    or::

        y(t) = output_equation_function(t,u)

    These could also represent discrete-time systems, in which case xdot(t)
    represents x[k+1].

    This can also model discontinuous systems. Discontinuities must occur on
    zero-crossings of the ``event_equation_function``, which take the same
    arguments as ``output_equation_function``, depending on ``dim_state``.
    At the zero-crossing, ``update_equation_function`` is called with the same
    arguments. If ``dim_state`` > 0, the return value of
    ``update_equation_function`` is used as the state of the system immediately
    after the discontinuity.
    """
    def __init__(self, state_equation_function=None,
                 output_equation_function=None, event_equation_function=None,
                 update_equation_function=None, dim_state=0, dim_input=0,
                 dim_output=0, dt=0, initial_condition=None):
        """
        Parameters
        ----------
        state_equation_function : callable, optional
            The derivative (or update equation) of the system state. Not needed
            if ``dim_state`` is zero.
        output_equation_function : callable, optional
            The output equation of the system. A system must have an
            ``output_equation_function``. If not set, uses full state output.
        event_equation_function : callable, optional
            The function whose output determines when discontinuities occur.
        update_equation_function : callable, optional
            The function called when a discontinuity occurs.
        dim_state : int, optional
            Dimension of the system state. Optional, defaults to 0.
        dim_input : int, optional
            Dimension of the system input. Optional, defaults to 0.
        dim_output : int, optional
            Dimension of the system output. Optional, defaults to dim_state.
        dt : float, optional
            Sample rate of the system. Optional, defaults to 0 representing a
            continuous time system.
        initial_condition : array_like of numerical values, optional
            Array or Matrix used as the initial condition of the system.
            Defaults to zeros of the same dimension as the state.
        """
        self.dim_state = dim_state
        self.dim_input = dim_input
        self.dim_output = dim_output or dim_state

        self.state_equation_function = state_equation_function

        self.output_equation_function = (
            full_state_output
            if output_equation_function is None and self.dim_state > 0
            else output_equation_function
        )

        self.event_equation_function = event_equation_function
        self.update_equation_function = update_equation_function
        self.initial_condition = initial_condition
        self.dt = dt

        self.validate()

    @property
    def initial_condition(self):
        return (self._initial_condition
                if self._initial_condition is not None
                else np.zeros(self.dim_state))

    @initial_condition.setter
    def initial_condition(self, initial_condition):
        self._initial_condition = initial_condition

    def prepare_to_integrate(self):
        return

    def validate(self):
        if self.dim_output == 0:
            raise ValueError(zero_dim_output_msg)

        if (self.dim_state > 0
                and getattr(self, 'state_equation_function', None) is None):
            raise ValueError(need_state_equation_function_msg)

        if (self.dim_state == 0
                and getattr(self, 'output_equation_function', None) is None):
            raise ValueError(need_output_equation_function_msg)

def SystemFromCallable(incallable, dim_input, dim_output, dt=0):
    """
    Construct a memoryless system from a callable.

    Parameters
    ----------
    incallable : callable
        Function to use as the output_equation_function. Should have signature
        (t, u) if dim_input > 0 or (t) if dim_input = 0.
    dim_input : int
        Dimension of input.
    dim_output : int
        Dimension of output.
    """
    system = DynamicalSystem(output_equation_function=incallable,
                             dim_input=dim_input, dim_output=dim_output, dt=dt)
    return system

class SwitchedSystem(DynamicalSystem):
    """
    Provides a useful pattern for discontinuous systems where the state and
    output equations change depending on the value of a function of the state
    and/or input (``event_variable_equation_function``). Most of the usefulness
    comes from constructing the ``event_equation_function`` with a Bernstein
    basis polynomial with roots at the boundaries. This class also provides
    logic for outputting the correct state and output equation based on the
    ``event_variable_equation_function`` value.
    """
    def __init__(self, state_equations_functions=None,
                 output_equations_functions=None,
                 event_variable_equation_function=None, event_bounds=None,
                 state_update_equation_function=None, dim_state=0, dim_input=0,
                 dim_output=0, dt=0, initial_condition=None):
        """
        Parameters
        ----------
        state_equations_functions : array_like of callables, optional
            The derivative (or update equation) of the system state. Not needed
            if ``dim_state`` is zero. The array indexes the
            event-state and should be one more than the number of event bounds.
            This should also be indexed to match the boundaries (i.e., the
            first function is used when the event variable is below the first
            event_bounds value). If only one callable is provided, the callable
            is used in each condition.
        output_equations_functions : array_like of callables, optional
            The output equation of the system. A system must have an
            ``output_equation_function``. If not set, uses full state output.
            The array indexes the event-state and should be one more than the
            number of event bounds. This should also be indexed to match the
            boundaries (i.e., the first function is used when the event
            variable is below the first event_bounds value). If only one
            callable is provided, the callable is used in each condition.
        event_variable_equation_function : callable
            When the output of this function crosses the values in
            ``event_bounds``, a discontuity event occurs.
        event_bounds : array_like of floats
            Defines the boundary points the trigger discontinuity events based
            on the output of ``event_variable_equation_function``.
        state_update_equation_function : callable, optional
            When an event occurs, the state update equation function is called
            to determine the state update. If not set, uses full state output,
            so the state is not changed upon a zero-crossing of the event
            variable function.
        dim_state : int, optional
            Dimension of the system state. Optional, defaults to 0.
        dim_input : int, optional
            Dimension of the system input. Optional, defaults to 0.
        dim_output : int, optional
            Dimension of the system output. Optional, defaults to dim_state.
        dt : float, optional
            Sample rate of the system. Optional, defaults to 0 representing a
            continuous time system.
        """
        self.dim_state = dim_state
        self.dim_input = dim_input
        self.dim_output = dim_output or dim_state
        self.event_bounds = event_bounds

        self.state_equations_functions = np.empty(self.n_conditions,
                                                  dtype=object)
        self.state_equations_functions[:] = state_equations_functions

        self.output_equations_functions = np.empty(self.n_conditions,
                                                   dtype=object)
        self.output_equations_functions[:] = (
            full_state_output
            if output_equations_functions is None and self.dim_state > 0
            else output_equations_functions
        )

        self.event_variable_equation_function = \
            event_variable_equation_function

        self.state_update_equation_function = (
            state_update_equation_function or
            full_state_output
        )

        self.initial_condition = initial_condition or np.zeros(dim_state)
        self.dt = dt

        self.validate()

    def validate(self):
        super().validate()

        if (self.dim_state > 0
                and np.any(np.equal(self.state_equations_functions, None))):
            raise ValueError(need_state_equation_function_msg)

        if (self.dim_state == 0
                and np.any(np.equal(self.output_equations_functions, None))):
            raise ValueError(need_output_equation_function_msg)

        if self.event_variable_equation_function is None:
            raise ValueError("A SwitchedSystem requires " +
                             "event_variable_equation_function")

    @property
    def event_bounds(self):
        return self._event_bounds

    @event_bounds.setter
    def event_bounds(self, event_bounds):
        if event_bounds is None:
            raise ValueError("A SwitchedSystem requires event_bounds")
        self._event_bounds = np.array(event_bounds).reshape(1, -1)
        self.n_conditions = self._event_bounds.size + 1
        if self.n_conditions == 2:
            self.event_bounds_range = 1
        else:
            self.event_bounds_range = np.diff(self.event_bounds[0, [0, -1]])

    def output_equation_function(self, *args):
        return self.output_equations_functions[self.condition_idx](*args)

    def state_equation_function(self, *args):
        return self.state_equations_functions[self.condition_idx](*args)

    def event_equation_function(self, *args):
        event_var = self.event_variable_equation_function(*args)
        return np.prod(
            (self.event_bounds_range-self.event_bounds)*event_var -
            self.event_bounds*(self.event_bounds_range - event_var),
            axis=1
        )

    def update_equation_function(self, *args):
        event_var = self.event_variable_equation_function(*args)
        if self.condition_idx is None:
            self.condition_idx = np.where(np.all(np.r_[
                    np.c_[[[True]], event_var >= self.event_bounds],
                    np.c_[event_var <= self.event_bounds, [[True]]]
                    ], axis=0))[0][0]
        else:
            sq_dist = (event_var - self.event_bounds)**2
            crossed_root_idx = np.where(sq_dist == np.min(sq_dist))[1][0]
            if crossed_root_idx == self.condition_idx:
                self.condition_idx += 1
            elif crossed_root_idx == self.condition_idx-1:
                self.condition_idx -= 1
            else:
                warnings.warn("SwitchedSystem did not cross a neighboring " +
                              "boundary. This may indicate an integration " +
                              "error. Continuing without updating " +
                              "condition_idx", UserWarning)
        return self.state_update_equation_function(*args)

    def prepare_to_integrate(self):
        self.condition_idx = None

class LTISystem(DynamicalSystem):
    """
    A linear, time-invariant system.
    """
    def __init__(self, *args, initial_condition=None, dt=0):
        """
        Construct an LTI system with the following input formats:

        1. state matrix A, input matrix B, output matrix C for systems with
           state::

              dx_dt = Ax + Bu
              y = Hx

        2. state matrix A, input matrix B for systems with state, assume full
           state output::

              dx_dt = Ax + Bu
              y = Ix

        3. gain matrix K for systems without state::

              y = Kx

        The matrices should be numeric arrays of consistent shape. The class
        provides ``A``, ``B``, ``C`` and ``F``, ``G``, ``H`` aliases for the
        matrices of systems with state, as well as a ``K`` alias for the gain
        matrix. The ``data`` alias provides the matrices as a tuple.
        """
        self.dt = dt

        if len(args) not in (1, 2, 3):
            raise ValueError("LTI system expects 1, 2, or 3 args")

        # TODO: setup jacobian functions
        if len(args) == 1:
            self.gain_matrix = gain_matrix = args[0]
            self.dim_input = (self.gain_matrix.shape[1]
                              if len(gain_matrix.shape) > 1
                              else 1)
            self.dim_output = self.gain_matrix.shape[0]
            self.dim_state = 0
            self.initial_condition = np.zeros(self.dim_state)
            self.output_equation_function = \
                lambda t, x: (gain_matrix@x).reshape(-1)
            return

        if len(args) == 2:
            state_matrix, input_matrix = args
            output_matrix = np.eye(state_matrix.shape[0])

        elif len(args) == 3:
            state_matrix, input_matrix, output_matrix = args

        if len(input_matrix.shape) == 1:
            input_matrix = input_matrix.reshape(-1, 1)

        self.dim_state = state_matrix.shape[0]
        self.dim_input = input_matrix.shape[1]
        self.dim_output = output_matrix.shape[0]

        self.state_matrix = state_matrix
        self.input_matrix = input_matrix
        self.output_matrix = output_matrix

        self.initial_condition = initial_condition or np.zeros(self.dim_state)
        self.state_equation_function = \
            lambda t, x, u: (state_matrix@x + input_matrix@u).reshape(-1)
        self.output_equation_function = \
            lambda t, x: (output_matrix@x).reshape(-1)

        self.validate()

    def validate(self):
        super().validate()
        if self.dim_state:
            assert self.state_matrix.shape[1] == self.dim_state
            assert self.input_matrix.shape[0] == self.dim_state
            assert self.output_matrix.shape[1] == self.dim_state

    @property
    def data(self):
        if self.dim_state:
            return self.state_matrix, self.input_matrix, self.output_matrix
        else:
            return self.gain_matrix

    @property
    def A(self):
        return self.state_matrix

    @property
    def F(self):
        return self.state_matrix

    @property
    def B(self):
        return self.input_matrix

    @property
    def G(self):
        return self.input_matrix

    @property
    def C(self):
        return self.output_matrix

    @property
    def H(self):
        return self.output_matrix

    @property
    def K(self):
        return self.gain_matrix
<|endoftext|>


<commit_before>from datetime import datetime

from django.core import mail
from django.db.models.signals import post_save

from testfixtures import LogCapture

import responses

from malaria24.ona.models import (
    ReportedCase, SMS, Digest, alert_new_case, MANAGER_DISTRICT, Facility)

from .base import MalariaTestCase

class ReportedCaseTest(MalariaTestCase):

    @responses.activate
    def test_capture_no_ehps(self):
        with LogCapture() as log:
            self.mk_case()
            log.check(('root',
                       'WARNING',
                       'No EHPs found for facility code facility_code.'))

    @responses.activate
    def test_capture_no_ehp_phonenumber(self):
        with LogCapture() as log:
            ehp = self.mk_ehp(phone_number='')
            case = self.mk_case(facility_code=ehp.facility_code)
            log.check(('root',
                       'WARNING',
                       ('Unable to SMS report for case %s. '
                        'Missing phone_number.') % (case.case_number,)))

    @responses.activate
    def test_capture_no_ehp_email_address(self):
        with LogCapture() as log:
            ehp = self.mk_ehp(email_address='')
            case = self.mk_case(facility_code=ehp.facility_code)
            log.check(('root',
                       'WARNING',
                       ('Unable to Email report for case %s. '
                        'Missing email_address.') % (case.case_number,)))

    @responses.activate
    def test_capture_no_reported_by(self):
        with LogCapture() as log:
            ehp = self.mk_ehp()
            case = self.mk_case(facility_code=ehp.facility_code,
                                reported_by='')
            log.check(('root',
                       'WARNING',
                       ('Unable to SMS case number for case %s. '
                        'Missing reported_by.') % (case.case_number,)))

    @responses.activate
    def test_capture_all_ok(self):
        self.assertEqual(SMS.objects.count(), 0)
        ehp = self.mk_ehp()
        case = self.mk_case(facility_code=ehp.facility_code)
        [ehp_sms, reporter_sms] = SMS.objects.all()
        self.assertEqual(ehp_sms.to, 'phone_number')
        self.assertEqual(ehp_sms.content,
                         'A new case has been reported, the full report will '
                         'be sent to you via email.')
        self.assertEqual(ehp_sms.message_id, 'the-message-id')
        self.assertEqual(reporter_sms.to, 'reported_by')
        self.assertEqual(
            reporter_sms.content,
            'Your reported case has been assigned case number %s.' % (
                case.case_number,))
        self.assertEqual(reporter_sms.message_id, 'the-message-id')

    @responses.activate
    def test_idempotency(self):
        self.assertEqual(SMS.objects.count(), 0)
        ehp = self.mk_ehp()
        case = self.mk_case(facility_code=ehp.facility_code)
        case.save()
        self.assertEqual(SMS.objects.count(), 2)

    @responses.activate
    def test_email_sending(self):
        facility = Facility.objects.create(facility_code='0001',
                                           facility_name='Facility 1',
                                           district='The District',
                                           subdistrict='The Subdistrict',
                                           province='The Province')
        ehp = self.mk_ehp(facility_code=facility.facility_code)
        case = self.mk_case(facility_code=facility.facility_code)
        [message] = mail.outbox
        self.assertEqual(message.subject,
                         'Malaria case number %s' % (case.case_number,))
        self.assertEqual(message.to, [ehp.email_address])
        self.assertTrue('does not support HTML' in message.body)
        [alternative] = message.alternatives
        content, content_type = alternative
        self.assertTrue(case.facility_code in content)
        self.assertTrue(case.sa_id_number in content)
        self.assertTrue('The District' in content)
        self.assertTrue('The Subdistrict' in content)
        self.assertTrue('The Province' in content)
        self.assertTrue('landmark' in content)
        self.assertTrue('landmark_description' in content)
        self.assertEqual('text/html', content_type)

    @responses.activate
    def test_age(self):
        case = self.mk_case(date_of_birth="820101")
        self.assertEqual(33, case.age)

    @responses.activate
    def test_facility_name(self):
        Facility.objects.create(facility_code='0001',
                                facility_name='Facility 1')
        case1 = self.mk_case(facility_code='0001')
        case2 = self.mk_case(facility_code='0002')
        self.assertEqual(case1.facility_names, 'Facility 1')
        self.assertEqual(case2.facility_names, 'Unknown')

class DigestTest(MalariaTestCase):

    def setUp(self):
        super(DigestTest, self).setUp()
        post_save.disconnect(alert_new_case, sender=ReportedCase)

    def tearDown(self):
        super(DigestTest, self).tearDown()
        post_save.connect(alert_new_case, sender=ReportedCase)

    @responses.activate
    def test_compile_digest(self):
        manager1 = self.mk_actor(role='MANAGER_DISTRICT')
        manager2 = self.mk_actor(role='MANAGER_DISTRICT')
        cases = [self.mk_case() for i in range(10)]
        digest = Digest.compile_digest()
        self.assertEqual([manager1, manager2],
                         list(digest.recipients.all().order_by('pk')))
        self.assertEqual([c.pk for c in cases],
                         [c.pk for c in digest.reportedcase_set.all()])

    @responses.activate
    def test_send_digest_email(self):
        manager1 = self.mk_actor(role=MANAGER_DISTRICT,
                                 email_address='manager@example.org')
        ehp1 = self.mk_ehp(name='EHP1', email_address='ehp1@example.org')
        ehp2 = self.mk_ehp(name='EHP2', email_address='ehp2@example.org')

        for i in range(10):
            case = self.mk_case()
            case.ehps.add(ehp1)
            case.ehps.add(ehp2)
            case.save()

        digest = Digest.compile_digest()
        digest.send_digest_email()
        [message] = mail.outbox
        [alternative] = message.alternatives
        html_content, content_type = alternative
        self.assertEqual(message.body.count('EHP1, EHP2'), 10)
        self.assertEqual(html_content.count('EHP1, EHP2'), 10)
        self.assertEqual(message.to, [manager1.email_address])
<commit_msg>flake8 fixes<commit_after>from django.core import mail
from django.db.models.signals import post_save

from testfixtures import LogCapture

import responses

from malaria24.ona.models import (
    ReportedCase, SMS, Digest, alert_new_case, MANAGER_DISTRICT, Facility)

from .base import MalariaTestCase

class ReportedCaseTest(MalariaTestCase):

    @responses.activate
    def test_capture_no_ehps(self):
        with LogCapture() as log:
            self.mk_case()
            log.check(('root',
                       'WARNING',
                       'No EHPs found for facility code facility_code.'))

    @responses.activate
    def test_capture_no_ehp_phonenumber(self):
        with LogCapture() as log:
            ehp = self.mk_ehp(phone_number='')
            case = self.mk_case(facility_code=ehp.facility_code)
            log.check(('root',
                       'WARNING',
                       ('Unable to SMS report for case %s. '
                        'Missing phone_number.') % (case.case_number,)))

    @responses.activate
    def test_capture_no_ehp_email_address(self):
        with LogCapture() as log:
            ehp = self.mk_ehp(email_address='')
            case = self.mk_case(facility_code=ehp.facility_code)
            log.check(('root',
                       'WARNING',
                       ('Unable to Email report for case %s. '
                        'Missing email_address.') % (case.case_number,)))

    @responses.activate
    def test_capture_no_reported_by(self):
        with LogCapture() as log:
            ehp = self.mk_ehp()
            case = self.mk_case(facility_code=ehp.facility_code,
                                reported_by='')
            log.check(('root',
                       'WARNING',
                       ('Unable to SMS case number for case %s. '
                        'Missing reported_by.') % (case.case_number,)))

    @responses.activate
    def test_capture_all_ok(self):
        self.assertEqual(SMS.objects.count(), 0)
        ehp = self.mk_ehp()
        case = self.mk_case(facility_code=ehp.facility_code)
        [ehp_sms, reporter_sms] = SMS.objects.all()
        self.assertEqual(ehp_sms.to, 'phone_number')
        self.assertEqual(ehp_sms.content,
                         'A new case has been reported, the full report will '
                         'be sent to you via email.')
        self.assertEqual(ehp_sms.message_id, 'the-message-id')
        self.assertEqual(reporter_sms.to, 'reported_by')
        self.assertEqual(
            reporter_sms.content,
            'Your reported case has been assigned case number %s.' % (
                case.case_number,))
        self.assertEqual(reporter_sms.message_id, 'the-message-id')

    @responses.activate
    def test_idempotency(self):
        self.assertEqual(SMS.objects.count(), 0)
        ehp = self.mk_ehp()
        case = self.mk_case(facility_code=ehp.facility_code)
        case.save()
        self.assertEqual(SMS.objects.count(), 2)

    @responses.activate
    def test_email_sending(self):
        facility = Facility.objects.create(facility_code='0001',
                                           facility_name='Facility 1',
                                           district='The District',
                                           subdistrict='The Subdistrict',
                                           province='The Province')
        ehp = self.mk_ehp(facility_code=facility.facility_code)
        case = self.mk_case(facility_code=facility.facility_code)
        [message] = mail.outbox
        self.assertEqual(message.subject,
                         'Malaria case number %s' % (case.case_number,))
        self.assertEqual(message.to, [ehp.email_address])
        self.assertTrue('does not support HTML' in message.body)
        [alternative] = message.alternatives
        content, content_type = alternative
        self.assertTrue(case.facility_code in content)
        self.assertTrue(case.sa_id_number in content)
        self.assertTrue('The District' in content)
        self.assertTrue('The Subdistrict' in content)
        self.assertTrue('The Province' in content)
        self.assertTrue('landmark' in content)
        self.assertTrue('landmark_description' in content)
        self.assertEqual('text/html', content_type)

    @responses.activate
    def test_age(self):
        case = self.mk_case(date_of_birth="820101")
        self.assertEqual(33, case.age)

    @responses.activate
    def test_facility_name(self):
        Facility.objects.create(facility_code='0001',
                                facility_name='Facility 1')
        case1 = self.mk_case(facility_code='0001')
        case2 = self.mk_case(facility_code='0002')
        self.assertEqual(case1.facility_names, 'Facility 1')
        self.assertEqual(case2.facility_names, 'Unknown')

class DigestTest(MalariaTestCase):

    def setUp(self):
        super(DigestTest, self).setUp()
        post_save.disconnect(alert_new_case, sender=ReportedCase)

    def tearDown(self):
        super(DigestTest, self).tearDown()
        post_save.connect(alert_new_case, sender=ReportedCase)

    @responses.activate
    def test_compile_digest(self):
        manager1 = self.mk_actor(role='MANAGER_DISTRICT')
        manager2 = self.mk_actor(role='MANAGER_DISTRICT')
        cases = [self.mk_case() for i in range(10)]
        digest = Digest.compile_digest()
        self.assertEqual([manager1, manager2],
                         list(digest.recipients.all().order_by('pk')))
        self.assertEqual([c.pk for c in cases],
                         [c.pk for c in digest.reportedcase_set.all()])

    @responses.activate
    def test_send_digest_email(self):
        manager1 = self.mk_actor(role=MANAGER_DISTRICT,
                                 email_address='manager@example.org')
        ehp1 = self.mk_ehp(name='EHP1', email_address='ehp1@example.org')
        ehp2 = self.mk_ehp(name='EHP2', email_address='ehp2@example.org')

        for i in range(10):
            case = self.mk_case()
            case.ehps.add(ehp1)
            case.ehps.add(ehp2)
            case.save()

        digest = Digest.compile_digest()
        digest.send_digest_email()
        [message] = mail.outbox
        [alternative] = message.alternatives
        html_content, content_type = alternative
        self.assertEqual(message.body.count('EHP1, EHP2'), 10)
        self.assertEqual(html_content.count('EHP1, EHP2'), 10)
        self.assertEqual(message.to, [manager1.email_address])
<|endoftext|>


<commit_before>"""
Rotate a fasta sequence so that the start is just before the terminase
"""

import os
import sys
import argparse
from roblib import stream_blast_results, rc, read_fasta

def read_blast_file(bf, evalue, verbose=False):
    """
    Read the blast file and store tuples of contig, start, stop for each hit
    :param bf: blast file to parse
    :param evalue: E value to keep matches
    :param verbose: more output
    :return:
    """

    results = {}

    for br in stream_blast_results(bf, verbose=verbose):
        if br.evalue > evalue:
            continue
        if br.query not in results:
            results[br.query] = { '-ve' : set(), '+ve' : set() }
        strand = '+ve'
        if br.query_end < br.query_start:
            strand = '-ve'
        results[br.query][strand] = results[br.query][strand] | set(range(br.query_start, br.query_end))
    return results

def is_consecutive(l):
    """
    Determines whether a list contains consecutive numbers. If it does, the sum of the
    numbers should be the same as n * (min + max) / 2 (where n is num elements)
    :param l: A set of numbers
    :return: True or False
    """

    total = sum(l)
    mathsum = len(l) * (min(l) + max(l)) / 2
    return total == mathsum

def find_gene(brs):
    """
    Test for where the genes are
    :param brs: The blast results
    :return:
    """

    contigs = brs.keys()
    if len(contigs) == 1:
        sys.stderr.write("There was only one contig with a match ({})\n".format(contigs[0]))
    else:
        sys.stderr.write("Multiple contigs had matches: {}\n".format(contigs))

    for c in contigs:
        for strand in ['-ve', '+ve']:
            if len(brs[c][strand]) == 0:
                continue
            if is_consecutive(brs[c][strand]):
                sys.stderr.write("There is a consecutive match to {} {} {}\n".format(c, max(brs[c][strand]), min(brs[c][strand])))
            else:
                sys.stderr.write("There are multiple discontinuous matches to {}. Try adjusting the evalue\n".format(c))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Recircularize the genome at the terminase")
    parser.add_argument('-f', help='fasta file of the genome', required=True)
    parser.add_argument('-l', help='blast output compared to terminase LARGE gene', required=True)
    parser.add_argument('-s', help='blast output compared to terminase SMALL gene', required=True)
    parser.add_argument('-e', help='evalue cutoff (default=all)', type=int, default=100)
    parser.add_argument('-v', help='verbose output', action="store_true")
    args = parser.parse_args()

    sys.stderr.write("Parsing small subunit\n")
    ss = read_blast_file(args.s, args.e, args.v)
    find_gene(ss)

    sys.stderr.write("Parsing large subunit\n")
    ls = read_blast_file(args.l, args.e, args.v)
    find_gene(ls)<commit_msg>blasting<commit_after>"""
Rotate a fasta sequence so that the start is just before the terminase
"""

import os
import sys
import argparse
from roblib import stream_blast_results, rc, read_fasta

def read_blast_file(bf, evalue, verbose=False):
    """
    Read the blast file and store tuples of contig, start, stop for each hit
    :param bf: blast file to parse
    :param evalue: E value to keep matches
    :param verbose: more output
    :return:
    """

    results = {}

    for br in stream_blast_results(bf, verbose=verbose):
        if br.evalue > evalue:
            continue
        if br.query not in results:
            results[br.query] = { '-ve' : set(), '+ve' : set() }
        strand = '+ve'
        (s, e) = (br.query_start, br.query_end)
        if br.query_end < br.query_start:
            strand = '-ve'
            (s, e) = (br.query_end, br.query_start)

        results[br.query][strand] = results[br.query][strand] | set(range(s, e))
    return results

def is_consecutive(l):
    """
    Determines whether a list contains consecutive numbers. If it does, the sum of the
    numbers should be the same as n * (min + max) / 2 (where n is num elements)
    :param l: A set of numbers
    :return: True or False
    """

    total = sum(l)
    mathsum = len(l) * (min(l) + max(l)) / 2
    return total == mathsum

def find_gene(brs):
    """
    Test for where the genes are
    :param brs: The blast results
    :return:
    """

    contigs = brs.keys()
    if len(contigs) == 1:
        sys.stderr.write("There was only one contig with a match ({})\n".format(list(contigs)[0]))
    else:
        sys.stderr.write("Multiple contigs had matches: {}\n".format(contigs))

    for c in contigs:
        sys.stderr.write(f"{c}\n")
        for strand in ['-ve', '+ve']:
            if len(brs[c][strand]) == 0:
                continue
            if is_consecutive(brs[c][strand]):
                sys.stderr.write("There is a consecutive match to {} {} {}\n".format(c, max(brs[c][strand]), min(brs[c][strand])))
            else:
                sys.stderr.write("There are multiple discontinuous matches to {}. Try adjusting the evalue\n".format(c))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Recircularize the genome at the terminase")
    parser.add_argument('-f', help='fasta file of the genome', required=True)
    parser.add_argument('-l', help='blast output compared to terminase LARGE gene', required=True)
    parser.add_argument('-s', help='blast output compared to terminase SMALL gene', required=True)
    parser.add_argument('-e', help='evalue cutoff (default=all)', type=int, default=100)
    parser.add_argument('-v', help='verbose output', action="store_true")
    args = parser.parse_args()

    sys.stderr.write("Parsing small subunit\n")
    ss = read_blast_file(args.s, args.e, args.v)
    find_gene(ss)

    sys.stderr.write("Parsing large subunit\n")
    ls = read_blast_file(args.l, args.e, args.v)
    find_gene(ls)
<|endoftext|>