<commit_before># -*- test-case-name: go.vumitools.tests.test_middleware -*-
import sys

from twisted.internet.defer import inlineCallbacks, returnValue

from vumi.middleware import (TransportMiddleware, TaggingMiddleware,
                                BaseMiddleware)
from vumi.application import TagpoolManager
from vumi.utils import normalize_msisdn
from vumi.persist.txriak_manager import TxRiakManager
from vumi.persist.message_store import MessageStore
from vumi import log

from go.vumitools.credit import CreditManager
from go.vumitools.account import AccountStore
from go.vumitools.conversation import ConversationStore

class NormalizeMsisdnMiddleware(TransportMiddleware):

    def setup_middleware(self):
        self.country_code = self.config['country_code']

    def handle_inbound(self, message, endpoint):
        from_addr = normalize_msisdn(message.get('from_addr'),
                        country_code=self.country_code)
        message['from_addr'] = from_addr
        return message

class DebitAccountError(Exception):
    """Exception raised if a message can't be paid for."""

class NoUserError(DebitAccountError):
    """Account could not be debited because no user was found."""

class NoTagError(DebitAccountError):
    """Account could not be debited because no tag was found."""

class BadTagPool(DebitAccountError):
    """Account could not be debited because the tag pool doesn't
       specify a cost."""

class InsufficientCredit(DebitAccountError):
    """Account could not be debited because the user account has
       insufficient credit."""

class GoApplicationRouterMiddleware(BaseMiddleware):
    """
    Base class for middlewares used by dispatchers using the
    `GoApplicationRouter`. It configures the `account_store` and the
    `message_store`.

    :type message_store: dict
    :param message_store:
        Dictionary containing the following values:

        *store_prefix*: the store prefix, defaults to 'message_store'

    :type redis: dict
    :param redis:
        Dictionary containing the configuration parameters for connecting
        to Redis with. Passed along as **kwargs to the Redis client.

    """
    def setup_middleware(self):
        from go.vumitools.api import get_redis
        r_server = get_redis(self.config)

        mdb_config = self.config.get('message_store', {})
        self.mdb_prefix = mdb_config.get('store_prefix', 'message_store')
        r_server = get_redis(self.config)
        self.manager = TxRiakManager.from_config({
                'bucket_prefix': self.mdb_prefix})
        self.account_store = AccountStore(self.manager)
        self.message_store = MessageStore(self.manager, r_server,
                                            self.mdb_prefix)

    def add_metadata_to_message(self, message):
        """
        Subclasses should override this method to appropriately set values
        on a message's `helper_metadata`. If specific message types or
        directions require different behaviour they can be overridden
        separately.
        """
        raise NotImplementedError("add_metadata_to_message should be "
                                    "implemented by the subclass")

    @inlineCallbacks
    def handle_inbound(self, message, endpoint):
        yield self.add_metadata_to_message(message)
        returnValue(message)

    @inlineCallbacks
    def handle_event(self, event, endpoint):
        yield self.add_metadata_to_message(event)
        returnValue(event)

    @inlineCallbacks
    def handle_outbound(self, message, endpoint):
        yield self.add_metadata_to_message(message)
        returnValue(message)

class LookupAccountMiddleware(GoApplicationRouterMiddleware):
    """
    Look up the account_key for a given message by retrieving
    this from the message tag's info.

    *NOTE*  This requires the `TaggingMiddleware` to be configured and placed
            before this middleware for this to work as it expects certain
            values to be set in the `helper_metadata`
    """

    @inlineCallbacks
    def find_account_key_for_message(self, message):
        tag = TaggingMiddleware.map_msg_to_tag(message)
        current_tag = yield self.message_store.get_tag_info(tag)
        if current_tag:
            batch = yield current_tag.current_batch.get()
            if batch:
                returnValue(batch.metadata['user_account'])

    @inlineCallbacks
    def add_metadata_to_message(self, message):
        account_key = yield self.find_account_key_for_message(message)
        if account_key:
            helper_metadata = message.get('helper_metadata', {})
            go_metadata = helper_metadata.setdefault('go', {})
            go_metadata['user_account'] = account_key

    @staticmethod
    def map_message_to_account_key(message):
        go_metadata = message.get('helper_metadata', {}).setdefault('go', {})
        return go_metadata.get('user_account')

class LookupBatchMiddleware(GoApplicationRouterMiddleware):
    """
    Look up a `batch_key` by inspecting the tag for a given message.

    *NOTE*  This requires the `TaggingMiddleware` to be configured and placed
            before this middleware to ensure that the appropriate tagging
            values are set in the `helper_metadata`
    """

    @inlineCallbacks
    def find_batch_for_message(self, message):
        tag = TaggingMiddleware.map_msg_to_tag(message)
        current_tag = yield self.message_store.get_tag_info(tag)
        if current_tag:
            batch = yield current_tag.current_batch.get()
            returnValue(batch)

    @inlineCallbacks
    def add_metadata_to_message(self, message):
        batch = yield self.find_batch_for_message(message)
        if batch:
            helper_metadata = message.get('helper_metadata', {})
            go_metadata = helper_metadata.setdefault('go', {})
            go_metadata['batch_key'] = batch.key

    @staticmethod
    def map_message_to_batch_key(message):
        go_metadata = message.get('helper_metadata', {}).get('go', {})
        return go_metadata.get('batch_key')

class LookupConversationMiddleware(GoApplicationRouterMiddleware):
    """
    Look up a conversation based on the `account_key` and `batch_key` that
    have been stored in the `helper_metadata` by the `LookupAccountMiddleware`
    and the `LookupBatchMiddleware` middlewares.

    *NOTE*  This middleware depends on the `LookupAccountMiddleware`,
            `LookupBatchMiddleware` and the `TaggingMiddleware` being
            configured and placed before this middleware to ensure that the
            appropriate variables are set in the `helper_metadata`
    """

    @inlineCallbacks
    def find_conversation_for_message(self, message):
        account_key = LookupAccountMiddleware.map_message_to_account_key(
                                                                    message)
        batch_key = LookupBatchMiddleware.map_message_to_batch_key(message)
        if account_key and batch_key:
            conversation_store = ConversationStore(self.manager, account_key)
            account_submanager = conversation_store.manager
            batch = self.message_store.batches(batch_key)
            all_conversations = yield batch.backlinks.conversations(
                                                            account_submanager)
            conversations = [c for c in all_conversations if not
                                c.ended()]
            if conversations:
                if len(conversations) > 1:
                    conv_keys = [c.key for c in conversations]
                    log.warning('Multiple conversations found '
                        'going with most recent: %r' % (conv_keys,))
                conversation = sorted(conversations, reverse=True,
                    key=lambda c: c.start_timestamp)[0]
                returnValue(conversation)

    @inlineCallbacks
    def add_metadata_to_message(self, message):
        conversation = yield self.find_conversation_for_message(message)
        if conversation:
            helper_metadata = message.get('helper_metadata', {})
            conv_metadata = helper_metadata.setdefault('conversations', {})
            conv_metadata['conversation_key'] = conversation.key
            conv_metadata['conversation_type'] = conversation.conversation_type

    @staticmethod
    def map_message_to_conversation_info(message):
        helper_metadata = message.get('helper_metadata', {})
        conv_metadata = helper_metadata.get('conversations', {})
        if conv_metadata:
            return (
                conv_metadata['conversation_key'],
                conv_metadata['conversation_type']
            )

class DebitAccountMiddleware(TransportMiddleware):

    def setup_middleware(self):
        # TODO: There really needs to be a helper function to
        #       turn this config into managers.
        from go.vumitools.api import get_redis
        r_server = get_redis(self.config)
        tpm_config = self.config.get('tagpool_manager', {})
        tpm_prefix = tpm_config.get('tagpool_prefix', 'tagpool_store')
        self.tpm = TagpoolManager(r_server, tpm_prefix)
        cm_config = self.config.get('credit_manager', {})
        cm_prefix = cm_config.get('credit_prefix', 'credit_store')
        self.cm = CreditManager(r_server, cm_prefix)

    def _credits_per_message(self, pool):
        tagpool_metadata = self.tpm.get_metadata(pool)
        credits_per_message = tagpool_metadata.get('credits_per_message')
        try:
            credits_per_message = int(credits_per_message)
            assert credits_per_message >= 0
        except Exception:
            exc_tb = sys.exc_info()[2]
            raise (BadTagPool,
                   BadTagPool("Invalid credits_per_message for pool %r"
                              % (pool,)),
                   exc_tb)
        return credits_per_message

    @staticmethod
    def map_msg_to_user(msg):
        """Convenience method for retrieving a user that was added
        to a message.
        """
        user_account = msg['helper_metadata'].get('go', {}).get('user_account')
        return user_account

    @staticmethod
    def map_payload_to_user(payload):
        """Convenience method for retrieving a user from a payload."""
        go_metadata = payload.get('helper_metadata', {}).get('go', {})
        return go_metadata.get('user_account')

    @staticmethod
    def add_user_to_message(msg, user_account_key):
        """Convenience method for adding a user to a message."""
        go_metadata = msg['helper_metadata'].setdefault('go', {})
        go_metadata['user_account'] = user_account_key

    @staticmethod
    def add_user_to_payload(payload, user_account_key):
        """Convenience method for adding a user to a message payload."""
        helper_metadata = payload.setdefault('helper_metadata', {})
        go_metadata = helper_metadata.setdefault('go', {})
        go_metadata['user_account'] = user_account_key

    def handle_outbound(self, msg, endpoint):
        # TODO: what actually happens when we raise an exception from
        #       inside middleware?
        user_account_key = self.map_msg_to_user(msg)
        if user_account_key is None:
            raise NoUserError(msg)
        tag = TaggingMiddleware.map_msg_to_tag(msg)
        if tag is None:
            raise NoTagError(msg)
        credits_per_message = self._credits_per_message(tag[0])
        self._debit_account(user_account_key, credits_per_message)
        success = self.cm.debit(user_account_key, credits_per_message)
        if not success:
            raise InsufficientCredit("User %r has insufficient credit"
                                     " to debit %r." %
                                     (user_account_key, credits_per_message))
        return msg
<commit_msg>GoApplicationRouterTestCase.test_no_tag test passing<commit_after># -*- test-case-name: go.vumitools.tests.test_middleware -*-
import sys

from twisted.internet.defer import inlineCallbacks, returnValue

from vumi.middleware import (TransportMiddleware, TaggingMiddleware,
                                BaseMiddleware)
from vumi.application import TagpoolManager
from vumi.utils import normalize_msisdn
from vumi.persist.txriak_manager import TxRiakManager
from vumi.persist.message_store import MessageStore
from vumi import log

from go.vumitools.credit import CreditManager
from go.vumitools.account import AccountStore
from go.vumitools.conversation import ConversationStore

class NormalizeMsisdnMiddleware(TransportMiddleware):

    def setup_middleware(self):
        self.country_code = self.config['country_code']

    def handle_inbound(self, message, endpoint):
        from_addr = normalize_msisdn(message.get('from_addr'),
                        country_code=self.country_code)
        message['from_addr'] = from_addr
        return message

class DebitAccountError(Exception):
    """Exception raised if a message can't be paid for."""

class NoUserError(DebitAccountError):
    """Account could not be debited because no user was found."""

class NoTagError(DebitAccountError):
    """Account could not be debited because no tag was found."""

class BadTagPool(DebitAccountError):
    """Account could not be debited because the tag pool doesn't
       specify a cost."""

class InsufficientCredit(DebitAccountError):
    """Account could not be debited because the user account has
       insufficient credit."""

class GoApplicationRouterMiddleware(BaseMiddleware):
    """
    Base class for middlewares used by dispatchers using the
    `GoApplicationRouter`. It configures the `account_store` and the
    `message_store`.

    :type message_store: dict
    :param message_store:
        Dictionary containing the following values:

        *store_prefix*: the store prefix, defaults to 'message_store'

    :type redis: dict
    :param redis:
        Dictionary containing the configuration parameters for connecting
        to Redis with. Passed along as **kwargs to the Redis client.

    """
    def setup_middleware(self):
        from go.vumitools.api import get_redis
        r_server = get_redis(self.config)

        mdb_config = self.config.get('message_store', {})
        self.mdb_prefix = mdb_config.get('store_prefix', 'message_store')
        r_server = get_redis(self.config)
        self.manager = TxRiakManager.from_config({
                'bucket_prefix': self.mdb_prefix})
        self.account_store = AccountStore(self.manager)
        self.message_store = MessageStore(self.manager, r_server,
                                            self.mdb_prefix)

    def add_metadata_to_message(self, message):
        """
        Subclasses should override this method to appropriately set values
        on a message's `helper_metadata`. If specific message types or
        directions require different behaviour they can be overridden
        separately.
        """
        raise NotImplementedError("add_metadata_to_message should be "
                                    "implemented by the subclass")

    @inlineCallbacks
    def handle_inbound(self, message, endpoint):
        yield self.add_metadata_to_message(message)
        returnValue(message)

    @inlineCallbacks
    def handle_event(self, event, endpoint):
        yield self.add_metadata_to_message(event)
        returnValue(event)

    @inlineCallbacks
    def handle_outbound(self, message, endpoint):
        yield self.add_metadata_to_message(message)
        returnValue(message)

class LookupAccountMiddleware(GoApplicationRouterMiddleware):
    """
    Look up the account_key for a given message by retrieving
    this from the message tag's info.

    *NOTE*  This requires the `TaggingMiddleware` to be configured and placed
            before this middleware for this to work as it expects certain
            values to be set in the `helper_metadata`
    """

    @inlineCallbacks
    def find_account_key_for_message(self, message):
        tag = TaggingMiddleware.map_msg_to_tag(message)
        if tag:
            current_tag = yield self.message_store.get_tag_info(tag)
            if current_tag:
                batch = yield current_tag.current_batch.get()
                if batch:
                    returnValue(batch.metadata['user_account'])

    @inlineCallbacks
    def add_metadata_to_message(self, message):
        account_key = yield self.find_account_key_for_message(message)
        if account_key:
            helper_metadata = message.get('helper_metadata', {})
            go_metadata = helper_metadata.setdefault('go', {})
            go_metadata['user_account'] = account_key

    @staticmethod
    def map_message_to_account_key(message):
        go_metadata = message.get('helper_metadata', {}).setdefault('go', {})
        return go_metadata.get('user_account')

class LookupBatchMiddleware(GoApplicationRouterMiddleware):
    """
    Look up a `batch_key` by inspecting the tag for a given message.

    *NOTE*  This requires the `TaggingMiddleware` to be configured and placed
            before this middleware to ensure that the appropriate tagging
            values are set in the `helper_metadata`
    """

    @inlineCallbacks
    def find_batch_for_message(self, message):
        tag = TaggingMiddleware.map_msg_to_tag(message)
        if tag:
            current_tag = yield self.message_store.get_tag_info(tag)
            if current_tag:
                batch = yield current_tag.current_batch.get()
                returnValue(batch)

    @inlineCallbacks
    def add_metadata_to_message(self, message):
        batch = yield self.find_batch_for_message(message)
        if batch:
            helper_metadata = message.get('helper_metadata', {})
            go_metadata = helper_metadata.setdefault('go', {})
            go_metadata['batch_key'] = batch.key

    @staticmethod
    def map_message_to_batch_key(message):
        go_metadata = message.get('helper_metadata', {}).get('go', {})
        return go_metadata.get('batch_key')

class LookupConversationMiddleware(GoApplicationRouterMiddleware):
    """
    Look up a conversation based on the `account_key` and `batch_key` that
    have been stored in the `helper_metadata` by the `LookupAccountMiddleware`
    and the `LookupBatchMiddleware` middlewares.

    *NOTE*  This middleware depends on the `LookupAccountMiddleware`,
            `LookupBatchMiddleware` and the `TaggingMiddleware` being
            configured and placed before this middleware to ensure that the
            appropriate variables are set in the `helper_metadata`
    """

    @inlineCallbacks
    def find_conversation_for_message(self, message):
        account_key = LookupAccountMiddleware.map_message_to_account_key(
                                                                    message)
        batch_key = LookupBatchMiddleware.map_message_to_batch_key(message)
        if account_key and batch_key:
            conversation_store = ConversationStore(self.manager, account_key)
            account_submanager = conversation_store.manager
            batch = self.message_store.batches(batch_key)
            all_conversations = yield batch.backlinks.conversations(
                                                            account_submanager)
            conversations = [c for c in all_conversations if not
                                c.ended()]
            if conversations:
                if len(conversations) > 1:
                    conv_keys = [c.key for c in conversations]
                    log.warning('Multiple conversations found '
                        'going with most recent: %r' % (conv_keys,))
                conversation = sorted(conversations, reverse=True,
                    key=lambda c: c.start_timestamp)[0]
                returnValue(conversation)

    @inlineCallbacks
    def add_metadata_to_message(self, message):
        conversation = yield self.find_conversation_for_message(message)
        if conversation:
            helper_metadata = message.get('helper_metadata', {})
            conv_metadata = helper_metadata.setdefault('conversations', {})
            conv_metadata['conversation_key'] = conversation.key
            conv_metadata['conversation_type'] = conversation.conversation_type

    @staticmethod
    def map_message_to_conversation_info(message):
        helper_metadata = message.get('helper_metadata', {})
        conv_metadata = helper_metadata.get('conversations', {})
        if conv_metadata:
            return (
                conv_metadata['conversation_key'],
                conv_metadata['conversation_type']
            )

class DebitAccountMiddleware(TransportMiddleware):

    def setup_middleware(self):
        # TODO: There really needs to be a helper function to
        #       turn this config into managers.
        from go.vumitools.api import get_redis
        r_server = get_redis(self.config)
        tpm_config = self.config.get('tagpool_manager', {})
        tpm_prefix = tpm_config.get('tagpool_prefix', 'tagpool_store')
        self.tpm = TagpoolManager(r_server, tpm_prefix)
        cm_config = self.config.get('credit_manager', {})
        cm_prefix = cm_config.get('credit_prefix', 'credit_store')
        self.cm = CreditManager(r_server, cm_prefix)

    def _credits_per_message(self, pool):
        tagpool_metadata = self.tpm.get_metadata(pool)
        credits_per_message = tagpool_metadata.get('credits_per_message')
        try:
            credits_per_message = int(credits_per_message)
            assert credits_per_message >= 0
        except Exception:
            exc_tb = sys.exc_info()[2]
            raise (BadTagPool,
                   BadTagPool("Invalid credits_per_message for pool %r"
                              % (pool,)),
                   exc_tb)
        return credits_per_message

    @staticmethod
    def map_msg_to_user(msg):
        """Convenience method for retrieving a user that was added
        to a message.
        """
        user_account = msg['helper_metadata'].get('go', {}).get('user_account')
        return user_account

    @staticmethod
    def map_payload_to_user(payload):
        """Convenience method for retrieving a user from a payload."""
        go_metadata = payload.get('helper_metadata', {}).get('go', {})
        return go_metadata.get('user_account')

    @staticmethod
    def add_user_to_message(msg, user_account_key):
        """Convenience method for adding a user to a message."""
        go_metadata = msg['helper_metadata'].setdefault('go', {})
        go_metadata['user_account'] = user_account_key

    @staticmethod
    def add_user_to_payload(payload, user_account_key):
        """Convenience method for adding a user to a message payload."""
        helper_metadata = payload.setdefault('helper_metadata', {})
        go_metadata = helper_metadata.setdefault('go', {})
        go_metadata['user_account'] = user_account_key

    def handle_outbound(self, msg, endpoint):
        # TODO: what actually happens when we raise an exception from
        #       inside middleware?
        user_account_key = self.map_msg_to_user(msg)
        if user_account_key is None:
            raise NoUserError(msg)
        tag = TaggingMiddleware.map_msg_to_tag(msg)
        if tag is None:
            raise NoTagError(msg)
        credits_per_message = self._credits_per_message(tag[0])
        self._debit_account(user_account_key, credits_per_message)
        success = self.cm.debit(user_account_key, credits_per_message)
        if not success:
            raise InsufficientCredit("User %r has insufficient credit"
                                     " to debit %r." %
                                     (user_account_key, credits_per_message))
        return msg
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
@author: Pieter Moris
'''

import os

def importGAF(path, geneSet):
    """
    Imports a GAF file (gene association format) and generates a dictionary mapping the gene uniprot AC to the GO ID.
    Only imports genes which are present in the provided (background) gene set.

    Information on the GAF 2.1 format can be found at
        http://geneontology.org/page/go-annotation-file-gaf-format-21

    Parameters
    ----------
    path : str
        The path to the file.
    geneSet : set
        A set containing the uniprot AC's of all the genes under consideration (background).

    Returns
    -------
    dict of str mapping to set
        A dictionary mapping gene uniprot AC's to GO ID's.

    Possible improvements:
        Check for `is_obsolete` and `replaced_by`, although the replacement term should be in OBO file as an entry.

        Check for inclusion in provided geneset afterwards using:
            gafDict = {key: value for key, value in gafDict.items() if key in geneSet }

        To do: double dictionary? already filter based on gene list? what to do with NOT?
    """

    gafPath = os.path.abspath(path)

    if not geneSet:
        print('Empty gene set was provided during creation of GAF dictionary.')
        exit()

    with open(gafPath, 'r') as gafFile:
        gafDict = {}
        for line in gafFile:
            if line.startswith('UniProtKB'):
                splitLine = line.split('\t')                # split column-wise
                uniprotAC = splitLine[1]
                goTerm = splitLine[4]
                goQualifier = splitLine[3]
                if 'NOT' not in goQualifier:                # ignore annotations with "NOT"
                    if uniprotAC in geneSet:                # only keep genes present in background gene set
                        if uniprotAC not in gafDict:        # Create new key if AC does not already appear in dictionary
                            # make dictionary containing uniprot AC as key and
                            gafDict[uniprotAC] = {goTerm}
                        else:                               # set of GO's as value
                            gafDict[uniprotAC].add(goTerm)

    print('Retrieved', len(gafDict),
          'annotated (background filtered) uniprot AC\'s from', gafPath + '\n')

    if len(gafDict) != len(geneSet):
        print('WARNING!\nNot every uniprot AC that was provided in the background set was found in the GAF file:')
        print([AC for AC in geneSet if AC not in gafDict])
        print('WARNING!\n')

    return gafDict

def createSubsetGafDict(subset, gafDict):
    """
    Generates a dictionary mapping the subset's gene uniprot AC's to the GO ID's,
    based on the provided gene subset and the gaf dictionary.

    Parameters
    ----------
    subset : set of str
        A subset of uniprot ACs of interest.

    Returns
    -------
    dict of str mapping to set
        A dictionary mapping the subset's gene uniprot AC's to GO ID's.
    """

    GafSubset = {}

    for gene, GOids in gafDict.items():
        if gene in subset:
            GafSubset[gene] = GOids

    if len(subset) != len(GafSubset):
        print('WARNING!\nNot every uniprot AC that was provided in the gene subset was found in the GAF file:')
        print([AC for AC in subset if AC not in GafSubset])
        print('WARNING!\n')

    return GafSubset
<commit_msg>Add import function for gaf file without specifiying background set<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
@author: Pieter Moris
'''

import os

def importGAF(path, geneSet):
    """
    Imports a GAF file (gene association format) and generates a dictionary mapping the gene uniprot AC to the GO ID.
    Only imports genes which are present in the provided (background) gene set.

    Information on the GAF 2.1 format can be found at
        http://geneontology.org/page/go-annotation-file-gaf-format-21

    Parameters
    ----------
    path : str
        The path to the file.
    geneSet : set
        A set containing the uniprot AC's of all the genes under consideration (background).

    Returns
    -------
    dict of str mapping to set
        A dictionary mapping gene uniprot AC's (str) to a set GO ID's.

    Possible improvements:
        Check for `is_obsolete` and `replaced_by`, although the replacement term should be in OBO file as an entry.

        Check for inclusion in provided gene set afterwards using:
            gafDict = {key: value for key, value in gafDict.items() if key in geneSet }

        To do: double dictionary? already filter based on gene list? what to do with NOT?
    """

    gafPath = os.path.abspath(path)

    if not geneSet:
        print('Empty gene set was provided during creation of GAF dictionary.')
        exit()

    with open(gafPath, 'r') as gafFile:
        gafDict = {}
        for line in gafFile:
            if line.startswith('UniProtKB'):
                splitLine = line.split('\t')                # split column-wise
                uniprotAC = splitLine[1]
                goTerm = splitLine[4]
                goQualifier = splitLine[3]
                if 'NOT' not in goQualifier:                # ignore annotations with "NOT"
                    if uniprotAC in geneSet:                # only keep genes present in background gene set
                        if uniprotAC not in gafDict:        # Create new key if AC does not already appear in dictionary
                            # make dictionary containing uniprot AC as key and
                            gafDict[uniprotAC] = {goTerm}
                        else:                               # set of GO's as value
                            gafDict[uniprotAC].add(goTerm)

    if len(gafDict) != len(geneSet):
        print('WARNING!\nNot every uniprot AC that was provided in the background set was found in the GAF file:')
        print([AC for AC in geneSet if AC not in gafDict])
        print('WARNING!\n')

    print('Retrieved', len(gafDict),
          'annotated (background filtered) uniprot AC\'s from', gafPath + '\n')

    return gafDict

def importFullGAF(path):
    """
    Imports a GAF file (gene association format) and generates a dictionary mapping the gene uniprot AC to the GO ID.
    Imports the full set of genes in the file.

    Information on the GAF 2.1 format can be found at
        http://geneontology.org/page/go-annotation-file-gaf-format-21

    Parameters
    ----------
    path : str
        The path to the file.

    Returns
    -------
    dict
        A dictionary mapping gene uniprot AC's (str) to a set of GO ID's.
   """

    gafPath = os.path.abspath(path)

    with open(gafPath, 'r') as gafFile:
        gafDict = {}
        for line in gafFile:
            if line.startswith('UniProtKB'):
                splitLine = line.split('\t')                # split column-wise
                uniprotAC = splitLine[1]
                goTerm = splitLine[4]
                goQualifier = splitLine[3]
                if 'NOT' not in goQualifier:                # ignore annotations with "NOT"
                    if uniprotAC not in gafDict:        # Create new key if AC does not already appear in dictionary
                        # make dictionary containing uniprot AC as key and
                        gafDict[uniprotAC] = {goTerm}
                    else:                               # set of GO's as value
                        gafDict[uniprotAC].add(goTerm)

    print('Retrieved', len(gafDict),
          'annotated uniprot AC\'s from', gafPath, 'after filtering on the background set.' + '\n')

    return gafDict

def createSubsetGafDict(subset, gafDict):
    """
    Generates a dictionary mapping the subset's gene uniprot AC's to the GO ID's,
    based on the provided gene subset and the gaf dictionary.

    Parameters
    ----------
    subset : set of str
        A subset of uniprot ACs of interest.

    Returns
    -------
    dict of str mapping to set
        A dictionary mapping the subset's gene uniprot AC's to GO ID's.
    """

    GafSubset = {}

    for gene, GOids in gafDict.items():
        if gene in subset:
            GafSubset[gene] = GOids

    if len(subset) != len(GafSubset):
        print('WARNING!\nNot every uniprot AC that was provided in the gene subset was found in the GAF file:')
        print([AC for AC in subset if AC not in GafSubset])
        print('WARNING!\n')

    return GafSubset
<|endoftext|>


<commit_before>"""Tenant views."""
# Copyright 2015 Solinea, Inc.
#
# Licensed under the Solinea Software License Agreement (goldstone),
# Version 1.0 (the "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at:
#
#     http://www.solinea.com/goldstone/LICENSE.pdf
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either expressed or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging

from rest_framework.serializers import ModelSerializer
from rest_framework.viewsets import ModelViewSet

from goldstone.utils import django_admin_only
from .models import Tenant

logger = logging.getLogger(__name__)

class TenantSerializer(ModelSerializer):
    """The Tenant model serializer."""

    class Meta:          # pylint: disable=C1001,C0111,W0232
        model = Tenant
        fields = ["name", "owner", "owner_contact", "uuid"]
        read_only_fields = ('uuid', )

class TenantsViewSet(ModelViewSet):
    """Provide all of the /tenants views."""

    queryset = Tenant.objects.all()
    serializer_class = TenantSerializer
    # lookup_field = "name"

    def get_queryset(self):
        """Return the queryset for list views."""

        return Tenant.objects.all()

    @django_admin_only
    def perform_create(self, serializer):
        """Add the system's default tenant_admin as the tenant_admin, and
        member, of the tenant we are creating.

        """
        from goldstone.accounts.models import Profile

        # Do what the superclass' perform_create() does, to get the newly
        # created row.
        tenant = serializer.save()

        # Get the tenant_admin. Use a filter in case there's erroneously more
        # than one in the system.
        admin_profile = Profile.objects.filter(default_tenant_admin=True)

        if not admin_profile:
            # No default tenant_admins is an error.
            logger.error("There are no default tenant_admins in the system."
                         " Using the Django administrator instead.")
            admin_profile = self.request.user.profile
        elif admin_profile.count() > 1:
            # More than one default tenant_admin is odd, but we'll continue.
            logger.warning("The system has more then one default tenant admin."
                           " There must be Only One: %s",
                           admin_profile)
            admin_profile = admin_profile[:1]
        else:
            # We found the default tenant_admin for this system.
            admin_profile = admin_profile[0]

        # Insert the default tenant_admin into the tenant and save it.
        admin_profile.tenant_admin = True
        admin_profile.tenant = tenant
        admin_profile.save()

    @django_admin_only
    def list(self, request, *args, **kwargs):
        """Provide a collection-of-objects GET response, for Django admins."""
        from rest_framework.response import Response

        # Return all the tenants to this Django admin.
        instance = self.filter_queryset(self.get_queryset())
        page = self.paginate_queryset(instance)

        serializer = \
            self.get_serializer(instance, many=True) if page is None else \
            self.get_pagination_serializer(page)

        return Response(serializer.data)
<commit_msg>GOLD-639: Bug, pep8 was not checking some files. Added system default_tenant_admin indicator, and added it to the tenant creation code. Created system initialization module, moved tasks private functions into it. Fixed pep8 errors.<commit_after>"""Tenant views."""
# Copyright 2015 Solinea, Inc.
#
# Licensed under the Solinea Software License Agreement (goldstone),
# Version 1.0 (the "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at:
#
#     http://www.solinea.com/goldstone/LICENSE.pdf
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either expressed or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging

from rest_framework.serializers import ModelSerializer
from rest_framework.viewsets import ModelViewSet

from goldstone.utils import django_admin_only
from .models import Tenant

logger = logging.getLogger(__name__)

class TenantSerializer(ModelSerializer):
    """The Tenant model serializer."""

    class Meta:          # pylint: disable=C1001,C0111,W0232
        model = Tenant
        fields = ["name", "owner", "owner_contact", "uuid"]
        read_only_fields = ('uuid', )

class TenantsViewSet(ModelViewSet):
    """Provide all of the /tenants views."""

    queryset = Tenant.objects.all()
    serializer_class = TenantSerializer
    lookup_field = "uuid"

    def get_queryset(self):
        """Return the queryset for list views."""

        return Tenant.objects.all()

    @django_admin_only
    def perform_create(self, serializer):
        """Add the system's default tenant_admin as the tenant_admin, and
        member, of the tenant we are creating.

        """
        from goldstone.accounts.models import Profile

        # Do what the superclass' perform_create() does, to get the newly
        # created row.
        tenant = serializer.save()

        # Get the tenant_admin. Use a filter in case there's erroneously more
        # than one in the system.
        admin_profile = Profile.objects.filter(default_tenant_admin=True)

        if not admin_profile:
            # No default tenant_admins is an error.
            logger.error("There are no default tenant_admins in the system."
                         " Using the Django administrator instead.")
            admin_profile = self.request.user.profile
        elif admin_profile.count() > 1:
            # More than one default tenant_admin is odd, but we'll continue.
            logger.warning("The system has more then one default tenant admin."
                           " There must be Only One: %s",
                           admin_profile)
            admin_profile = admin_profile[:1]
        else:
            # We found the default tenant_admin for this system.
            admin_profile = admin_profile[0]

        # Insert the default tenant_admin into the tenant and save it.
        admin_profile.tenant_admin = True
        admin_profile.tenant = tenant
        admin_profile.save()

    @django_admin_only
    def list(self, request, *args, **kwargs):
        """Provide a collection-of-objects GET response, for Django admins."""
        from rest_framework.response import Response

        # Return all the tenants to this Django admin.
        instance = self.filter_queryset(self.get_queryset())
        page = self.paginate_queryset(instance)

        serializer = \
            self.get_serializer(instance, many=True) if page is None else \
            self.get_pagination_serializer(page)

        return Response(serializer.data)

    @django_admin_only
    def retrieve(self, *args, **kwargs):
        """Return a single Tenant record, for a Django admin."""

        super(TenantsViewSet, self).retrieve(*args, **kwargs)
<|endoftext|>


<commit_before>from product import ProductType, Product, Package
from vector2d import Vector2D
from collections import deque

class Organization(object):
	'''
	Class representing a 'place' in the world, such as a factory or store.
	Virtual class, inherited only by Factory so far.
	'''
	name = 'Organization'

	def __init__(self, position):
		assert isinstance(position, Vector2D), 'position not a Vector2D.'
		self.position = position
		return None

	def single_round(self):
		pass

class Factory(Organization):
	'''
	Class representing a single manufacturing unit.
	The factory is placed on a certain position in the game world,
	and manufactures a product based on its specification.
	'''
	name = 'Factory'

	def __init__(self, product_type, capacity, position):
		'''
		Params:
			-- product_type - produced ProductType ref
			-- capacity - int, production per time unit
			-- position - 2D Vector that positions factory in game world
		'''
		Organization.__init__(self, position)

		assert isinstance(product_type, ProductType),\
			'product_type not a ProductType'
		self.product_type = product_type

		assert isinstance(capacity, int), 'capacity not an int.'
		self.capacity = int(capacity/self.product_type.difficulty)

		# storage of required materials
		self.storage = dict()
		for item in self.product_type.requirements:
			self.storage[item] = 0

		# list of factories supplying the materials
		self.suppliers = dict()

		# supply input queue
		self.input = deque()

		# list of customers for products
		self.customers = dict()

		# output queue
		self.output = deque()

		return None

	def __repr__(self):
		return '<Factory, ID={0}, position={1}, capacity={2}, product={3}>'.format(
			hex(id(self)), self.position, self.capacity, self.product_type.name)

	def __str__(self):
		return 'Factory at {0} producing {1}.\n'.format(
			self.position, self.product_type.name)\
			+ 'Storage state:\n{0}\n'.format('\n'.join(['-- {0}: {1} units'.format(
					key.name, self.storage[key]) for key in self.storage]))\
					+ 'Output queue length: {0}'.format(len(self.output))

	def update_suppliers(self, world):
		'''
		Method for updating the refs to factories supplying required materials.
		Param world -- ref to world object
		'''
		for m in self.storage:
			self.suppliers[m] = world.product_origin[m]
			self.customers[self.suppliers[m]] = self.requirements[m]
		return None

	def material_sufficient(self):
		'''
		Aux. method returning the number of product items that could
		be produced with current materials in storage (not taking into account
		the capacity of the factory).
		'''
		return int(min([self.storage[i]/self.product_type.requirements[i]\
			for i in self.storage]))

	def accept_supplies(self):
		'''
		Method for accepting the supplies, i.e. moving them from the first
		place in input queue into the storage.
		'''
		if len(self.input) == 0: # no supplies available
			 return None
		incoming_load = self.input.popleft()
		for package in incoming_load:
			while package.size() > 0:
				item = package.unpack
				assert item in self.storage, 'Accepted spam package.'
				self.storage[item] += 1
		return None

	def produce(self):
		'''
		Method for handling producing during one time unit.
		'''
		amount = min([self.material_sufficient(), self.capacity])
		for item in self.product_type.requirements:
			self.storage[item] -= amount * self.product_type.requirements[item]
		for p in range(amount):
			self.output.append(Product(self.product_type))
		return None

	def send(self, package, customer):
		dist = int((self.position - customer.position).length)
		while len(customer.input) < dist:
			customer.input.append(list())
		customer.input[dist - 1].append(package)
		return None
	def export(self):
		'''
		Method handling production export to customers.
		Current system: production is divided among customers in ratios
		equivalent to their product requirements.
		'''
		total_requirements = sum(self.customers.values())
		production_available = len(self.output)
		for cust in self.customers:
			package = Package()
			for i in range(int(production_available * \
				self.customers[cust]/total_requirements)):
				package.pack(self.output.popleft())
			self.send(package, cust)
		return None

	def single_round(self):
		'''
		Method for handling actions taking place in a single time unit.
		'''
		self.accept_supplies()
		self.produce()
		self.export()
		return None

ziemiaciek = ProductType('Ziemiaciek', {}, 1.00)
a = Factory(ziemiaciek, 220, Vector2D(10,10))
b = Factory(ProductType('Pluszowy Mort', {ziemiaciek: 100}, 1.02), 470, Vector2D(12,10))
<commit_msg>Modified production amount determination in Factory.produce method to account for cases when the Factory has no required materials (the amount is then set to be equal to factory capacity).<commit_after>from product import ProductType, Product, Package
from vector2d import Vector2D
from collections import deque

class Organization(object):
	'''
	Class representing a 'place' in the world, such as a factory or store.
	Virtual class, inherited only by Factory so far.
	'''
	name = 'Organization'

	def __init__(self, position):
		assert isinstance(position, Vector2D), 'position not a Vector2D.'
		self.position = position
		return None

	def single_round(self):
		pass

class Factory(Organization):
	'''
	Class representing a single manufacturing unit.
	The factory is placed on a certain position in the game world,
	and manufactures a product based on its specification.
	'''
	name = 'Factory'

	def __init__(self, product_type, capacity, position):
		'''
		Params:
			-- product_type - produced ProductType ref
			-- capacity - int, production per time unit
			-- position - 2D Vector that positions factory in game world
		'''
		Organization.__init__(self, position)

		assert isinstance(product_type, ProductType),\
			'product_type not a ProductType'
		self.product_type = product_type

		assert isinstance(capacity, int), 'capacity not an int.'
		self.capacity = int(capacity/self.product_type.difficulty)

		# storage of required materials
		self.storage = dict()
		for item in self.product_type.requirements:
			self.storage[item] = 0

		# list of factories supplying the materials
		self.suppliers = dict()

		# supply input queue
		self.input = deque()

		# list of customers for products
		self.customers = dict()

		# output queue
		self.output = deque()

		return None

	def __repr__(self):
		return '<Factory, ID={0}, position={1}, capacity={2}, product={3}>'.format(
			hex(id(self)), self.position, self.capacity, self.product_type.name)

	def __str__(self):
		return 'Factory at {0} producing {1}.\n'.format(
			self.position, self.product_type.name)\
			+ 'Storage state:\n{0}\n'.format('\n'.join(['-- {0}: {1} units'.format(
					key.name, self.storage[key]) for key in self.storage]))\
					+ 'Output queue length: {0}'.format(len(self.output))

	def update_suppliers(self, world):
		'''
		Method for updating the refs to factories supplying required materials.
		Param world -- ref to world object
		'''
		for m in self.storage:
			self.suppliers[m] = world.product_origin[m]
			self.customers[self.suppliers[m]] = self.requirements[m]
		return None

	def material_sufficient(self):
		'''
		Aux. method returning the number of product items that could
		be produced with current materials in storage (not taking into account
		the capacity of the factory).
		'''
		return int(min([self.storage[i]/self.product_type.requirements[i]\
			for i in self.storage]))

	def accept_supplies(self):
		'''
		Method for accepting the supplies, i.e. moving them from the first
		place in input queue into the storage.
		'''
		if len(self.input) == 0: # no supplies available
			 return None
		incoming_load = self.input.popleft()
		for package in incoming_load:
			while package.size() > 0:
				item = package.unpack
				assert item in self.storage, 'Accepted spam package.'
				self.storage[item] += 1
		return None

	def produce(self):
		'''
		Method for handling producing during one time unit.
		'''
		if len(self.storage) > 0:
			amount = min([self.material_sufficient(), self.capacity])
		else:
			amount = self.capacity # When no materials are required.

		for item in self.product_type.requirements:
			self.storage[item] -= amount * self.product_type.requirements[item]
		for p in range(amount):
			self.output.append(Product(self.product_type))
		return None

	def send(self, package, customer):
		dist = int((self.position - customer.position).length)
		while len(customer.input) < dist:
			customer.input.append(list())
		customer.input[dist - 1].append(package)
		return None
	def export(self):
		'''
		Method handling production export to customers.
		Current system: production is divided among customers in ratios
		equivalent to their product requirements.
		'''
		total_requirements = sum(self.customers.values())
		production_available = len(self.output)
		for cust in self.customers:
			package = Package()
			for i in range(int(production_available * \
				self.customers[cust]/total_requirements)):
				package.pack(self.output.popleft())
			self.send(package, cust)
		return None

	def single_round(self):
		'''
		Method for handling actions taking place in a single time unit.
		'''
		self.accept_supplies()
		self.produce()
		self.export()
		return None
<|endoftext|>


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

"""Provides function wrappers that implement page streaming and retrying."""

from __future__ import absolute_import
from . import config

def _add_timeout_arg(a_func, timeout):
    """Updates a_func so that it gets called with the timeout as its final arg.

    This converts a callable, a_func, into another callable with an additional
    positional arg.

    Args:
      a_func (callable): a callable to be updated
      timeout (int): to be added the original callable as it final positional
        args.

    Returns:
      callable: the original callable updated to the timeout arg
    """

    def inner(*args, **kw):
        """Updates args with the timeout."""
        updated_args = args + (timeout,)
        return a_func(*updated_args, **kw)

    return inner

def _retryable(a_func, max_attempts):
    """Creates a function equivalent to retrying, but that retries on certain
    exceptions.

    Args:
        call: A function.
        max_attempts: The maximum number of times that the call should be
            attempted; the call will always be attempted at least once.

    Returns:
        A function that will retry on exception.
    """

    def inner(*args, **kwargs):
        "Retries a_func upto max_attempt times"
        attempt_count = 0
        while 1:
            try:
                return a_func(*args, **kwargs)
            except config.RETRY_EXCEPTIONS:
                attempt_count += 1
                if attempt_count < max_attempts:
                    continue
                raise

    return inner

def _page_streamable(a_func,
                     request_page_token_field,
                     response_page_token_field,
                     resource_field,
                     timeout):
    """Creates a function that yields an iterable to performs page-streaming.

    Args:
        a_func: a call to a page streamable API
        request_page_token_field: The field of the page token in the request.
        response_page_token_field: The field of the next page token in the
            response.
        resource_field: The field to be streamed.
        timeout: the timeout to apply to the API call.

    Returns:
        A function that returns an iterable over the specified field.
    """
    with_timeout = _add_timeout_arg(a_func, timeout)

    def inner(request):
        """A generator that yields all the paged responses."""
        while True:
            response = with_timeout(request)
            for obj in getattr(response, resource_field):
                yield obj
            next_page_token = getattr(response, response_page_token_field)
            if not next_page_token:
                break
            setattr(request, request_page_token_field, next_page_token)

    return inner

class CallOptions(object):
    """Encapsulates the default settings for ApiCallable."""
    # pylint: disable=too-few-public-methods
    def __init__(self, timeout=30, is_idempotent_retrying=True,
                 max_attempts=16):
        """Constructor.

        Args:
            timeout: The client-side timeout for API calls.
            is_idempotent_retrying: If set, calls determined by configuration
                to be idempotent will retry upon transient error by default.
            max_attempts: The maximum number of attempts that should be made
                for a retrying call to this service.

        Returns:
            An CallOptions object.
        """
        self.timeout = timeout
        self.is_idempotent_retrying = is_idempotent_retrying
        self.max_attempts = max_attempts

def idempotent_callable(func, timeout=None, is_retrying=None,
                        page_streaming=None, max_attempts=None, defaults=None):
    """Creates an ApiCallable for an idempotent call.

    Args:
        func: The API call that this ApiCallable wraps.
        timeout: The timeout parameter to the API call. If not supplied, will
            default to the value in the defaults parameter.
        is_retrying: Boolean indicating whether this call should retry upon a
            transient error. If None, retrying will be determined by the
            defaults parameter.
        page_streaming: page_descriptor.PageDescriptor indicating the structure
            of page streaming to be performed. If None, this call will not
            perform page streaming.
        max_attempts: If is_retrying, the maximum number of times this call may
            be attempted. If not specified, will default to the value in the
            defaults parameter.
        defaults: A CallOptions object, from which default values will
            be drawn if not supplied by the other named parameters. The other
            named parameters always override those in the defaults. If neither
            the is_retrying nor defaults parameter is specified, a runtime
            error will result at callable creation time.

    Returns:
        An ApiCallable object.
    """
    if is_retrying is None:
        to_retry = defaults.is_idempotent_retrying
    else:
        to_retry = is_retrying
    return ApiCallable(
        func, timeout=timeout, page_streaming=page_streaming,
        max_attempts=max_attempts, defaults=defaults, is_retrying=to_retry)

class ApiCallable(object):
    """Represents zero or more API calls, with options to retry or perform
    page streaming.

    Calling an object of ApiCallable type causes these calls to be transmitted.
    """
    # pylint: disable=too-few-public-methods
    def __init__(self, func, timeout=None, is_retrying=False,
                 page_streaming=None, max_attempts=None, defaults=None):
        """Constructor.

        Args:
            func: The API call that this ApiCallable wraps.
            timeout: The timeout parameter to the API call. If not supplied,
                will default to the value in the defaults parameter.
            is_retrying: Boolean indicating whether this call should retry upon
                a transient error.
            page_streaming: page_descriptor.PageDescriptor indicating the
                structure of page streaming to be performed. If None, this call
                will not perform page streaming.
            max_attempts: If is_retrying, the maximum number of times this call
                may be attempted. If not specified, will default to the value
                in the defaults parameter.
            defaults: A CallOptions object, from which default values
                will be drawn if not supplied by the other named parameters.
                The other named parameters always override those in the
                defaults. If neither the defaults nor timeout parameter is
                specified, a runtime error will result at call time. If neither
                the defaults nor the max_attempts parameter is specified for a
                retrying call, a runtime error will result at call time.

        Returns:
            An ApiCallable object.
        """
        self.func = func
        self.is_retrying = is_retrying
        self.page_descriptor = page_streaming
        self.max_attempts = max_attempts
        self.timeout = timeout
        if defaults is not None:
            if max_attempts is None:
                self.max_attempts = defaults.max_attempts
            if timeout is None:
                self.timeout = defaults.timeout

    def __call__(self, *args, **kwargs):
        the_func = self.func

        # Update the_func using each of the applicable function decorators
        # before calling.
        if self.is_retrying:
            the_func = _retryable(the_func, self.max_attempts)
        if self.page_descriptor:
            the_func = _page_streamable(
                the_func,
                self.page_descriptor.request_page_token_field,
                self.page_descriptor.response_page_token_field,
                self.page_descriptor.resource_field,
                self.timeout)
        else:
            the_func = _add_timeout_arg(the_func, self.timeout)
        return the_func(*args, **kwargs)
<commit_msg>Update docstrings<commit_after># Copyright 2016, Google Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Provides function wrappers that implement page streaming and retrying."""

from __future__ import absolute_import
from . import config

def _add_timeout_arg(a_func, timeout):
    """Updates a_func so that it gets called with the timeout as its final arg.

    This converts a callable, a_func, into another callable with an additional
    positional arg.

    Args:
      a_func (callable): a callable to be updated
      timeout (int): to be added the original callable as it final positional
        args.

    Returns:
      callable: the original callable updated to the timeout arg
    """

    def inner(*args, **kw):
        """Updates args with the timeout."""
        updated_args = args + (timeout,)
        return a_func(*updated_args, **kw)

    return inner

def _retryable(a_func, max_attempts):
    """Creates a function equivalent to a_func, but that retries on certain
    exceptions.

    Args:
        a_func (callable): A callable.
        max_attempts (int): The maximum number of times that the call should
            be attempted; the call will always be attempted at least once.

    Returns:
        A function that will retry on exception.
    """

    def inner(*args, **kwargs):
        "Retries a_func upto max_attempt times"
        attempt_count = 0
        while 1:
            try:
                return a_func(*args, **kwargs)
            except config.RETRY_EXCEPTIONS:
                attempt_count += 1
                if attempt_count < max_attempts:
                    continue
                raise

    return inner

def _page_streamable(a_func,
                     request_page_token_field,
                     response_page_token_field,
                     resource_field,
                     timeout):
    """Creates a function that yields an iterable to performs page-streaming.

    Args:
        a_func: an API call that is page streaming.
        request_page_token_field: The field of the page token in the request.
        response_page_token_field: The field of the next page token in the
            response.
        resource_field: The field to be streamed.
        timeout: the timeout to apply to the API call.

    Returns:
        A function that returns an iterable over the specified field.
    """
    with_timeout = _add_timeout_arg(a_func, timeout)

    def inner(request):
        """A generator that yields all the paged responses."""
        while True:
            response = with_timeout(request)
            for obj in getattr(response, resource_field):
                yield obj
            next_page_token = getattr(response, response_page_token_field)
            if not next_page_token:
                break
            setattr(request, request_page_token_field, next_page_token)

    return inner

class CallOptions(object):
    """Encapsulates the default settings for ApiCallable."""
    # pylint: disable=too-few-public-methods
    def __init__(self, timeout=30, is_idempotent_retrying=True,
                 max_attempts=16):
        """Constructor.

        Args:
            timeout: The client-side timeout for API calls.
            is_idempotent_retrying: If set, calls determined by configuration
                to be idempotent will retry upon transient error by default.
            max_attempts: The maximum number of attempts that should be made
                for a retrying call to this service.

        Returns:
            An CallOptions object.
        """
        self.timeout = timeout
        self.is_idempotent_retrying = is_idempotent_retrying
        self.max_attempts = max_attempts

def idempotent_callable(func, timeout=None, is_retrying=None,
                        page_streaming=None, max_attempts=None, defaults=None):
    """Creates an ApiCallable for an idempotent call.

    Args:
        func: The API call that this ApiCallable wraps.
        timeout: The timeout parameter to the API call. If not supplied, will
            default to the value in the defaults parameter.
        is_retrying: Boolean indicating whether this call should retry upon a
            transient error. If None, retrying will be determined by the
            defaults parameter.
        page_streaming: page_descriptor.PageDescriptor indicating the structure
            of page streaming to be performed. If None, this call will not
            perform page streaming.
        max_attempts: If is_retrying, the maximum number of times this call may
            be attempted. If not specified, will default to the value in the
            defaults parameter.
        defaults: A CallOptions object, from which default values will
            be drawn if not supplied by the other named parameters. The other
            named parameters always override those in the defaults. If neither
            the is_retrying nor defaults parameter is specified, a runtime
            error will result at callable creation time.

    Returns:
        An ApiCallable object.
    """
    if is_retrying is None:
        to_retry = defaults.is_idempotent_retrying
    else:
        to_retry = is_retrying
    return ApiCallable(
        func, timeout=timeout, page_streaming=page_streaming,
        max_attempts=max_attempts, defaults=defaults, is_retrying=to_retry)

class ApiCallable(object):
    """Represents zero or more API calls, with options to retry or perform
    page streaming.

    Calling an object of ApiCallable type causes these calls to be transmitted.
    """
    # pylint: disable=too-few-public-methods
    def __init__(self, func, timeout=None, is_retrying=False,
                 page_streaming=None, max_attempts=None, defaults=None):
        """Constructor.

        Args:
            func: The API call that this ApiCallable wraps.
            timeout: The timeout parameter to the API call. If not supplied,
                will default to the value in the defaults parameter.
            is_retrying: Boolean indicating whether this call should retry upon
                a transient error.
            page_streaming: page_descriptor.PageDescriptor indicating the
                structure of page streaming to be performed. If None, this call
                will not perform page streaming.
            max_attempts: If is_retrying, the maximum number of times this call
                may be attempted. If not specified, will default to the value
                in the defaults parameter.
            defaults: A CallOptions object, from which default values
                will be drawn if not supplied by the other named parameters.
                The other named parameters always override those in the
                defaults. If neither the defaults nor timeout parameter is
                specified, a runtime error will result at call time. If neither
                the defaults nor the max_attempts parameter is specified for a
                retrying call, a runtime error will result at call time.

        Returns:
            An ApiCallable object.
        """
        self.func = func
        self.is_retrying = is_retrying
        self.page_descriptor = page_streaming
        self.max_attempts = max_attempts
        self.timeout = timeout
        if defaults is not None:
            if max_attempts is None:
                self.max_attempts = defaults.max_attempts
            if timeout is None:
                self.timeout = defaults.timeout

    def __call__(self, *args, **kwargs):
        the_func = self.func

        # Update the_func using each of the applicable function decorators
        # before calling.
        if self.is_retrying:
            the_func = _retryable(the_func, self.max_attempts)
        if self.page_descriptor:
            the_func = _page_streamable(
                the_func,
                self.page_descriptor.request_page_token_field,
                self.page_descriptor.response_page_token_field,
                self.page_descriptor.resource_field,
                self.timeout)
        else:
            the_func = _add_timeout_arg(the_func, self.timeout)
        return the_func(*args, **kwargs)
<|endoftext|>


<commit_before>import argparse

import log
from client import dispatcher
from communication import server, client, messages
from interfaces.cli import menu
from user import settings
from interfaces.gui.guimanager import GuiProcessManager
from client import cmwatcher
from communication.procs import EngineProcessManager

def main():
    clargs = get_args()
    if settings.user_settings['network'] == 'server':
        server.RemoteEngineServer().loop_forever()
        return
    bootup(clargs)

def bootup(clargs):
    ui_manager = GuiProcessManager() if clargs.interface == 'gui' else menu.MainMenu()
    if settings.user_settings['network'] == 'local':
        ref = EngineProcessManager()
    else:
        ref = client.RemoteEngineClient().connect()
    cmw = cmwatcher.CommandModuleWatcher()
    cmw.initialize_modules()
    cmw.start_watch_active_window()
    ui_manager.main_loop()
    ref.socket.close()
    messages.dispatch('shutdown')

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--interface', default='remote')
    parser.add_argument('--network', default='local') # or remote
    parser.add_argument('--engine_server', action='store_true')
    return parser.parse_args()

if __name__ == "__main__":
    main()<commit_msg>engine stop<commit_after>import argparse

import log
from client import dispatcher
from communication import server, client, messages
from interfaces.cli import menu
from user import settings
from interfaces.gui.guimanager import GuiProcessManager
from client import cmwatcher
from communication.procs import EngineProcessManager

def main():
    clargs = get_args()
    if settings.user_settings['network'] == 'server':
        server.RemoteEngineServer().loop_forever()
        return
    ui_manager = GuiProcessManager() if clargs.interface == 'gui' else menu.MainMenu()
    io_obj = get_io()
    try:
        cmw = cmwatcher.CommandModuleWatcher()
        cmw.initialize_modules()
        cmw.start_watch_active_window()
        ui_manager.main_loop()
    finally:
        messages.dispatch_sync('engine stop')
        io_obj.socket.close()

def get_io():
    if settings.user_settings['network'] == 'local':
        return EngineProcessManager()
    else:
        engine_client = client.RemoteEngineClient()
        engine_client.connect()
        return engine_client

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--interface', default='remote')
    parser.add_argument('--network', default='local') # or remote
    parser.add_argument('--engine_server', action='store_true')
    return parser.parse_args()

if __name__ == "__main__":
    main()<|endoftext|>