<commit_before>from django.conf import settings

from casexml.apps.phone.models import SyncLogSQL
from dimagi.utils.parsing import string_to_utc_datetime
from pillowtop.checkpoints.manager import KafkaPillowCheckpoint
from pillowtop.feed.interface import Change
from pillowtop.pillow.interface import ConstructedPillow
from pillowtop.processors.interface import PillowProcessor
from pillowtop.reindexer.reindexer import Reindexer, ReindexerFactory

from corehq.apps.change_feed import topics
from corehq.apps.change_feed.consumer.feed import (
    KafkaChangeFeed,
    KafkaCheckpointEventHandler,
)
from corehq.apps.receiverwrapper.util import get_version_from_build_id
from corehq.apps.users.models import (
    CommCareUser,
    CouchUser,
    DeviceAppMeta,
    UserReportingMetadataStaging,
    WebUser,
)
from corehq.apps.users.util import (
    update_device_meta,
    update_last_sync,
    update_latest_builds,
)
from corehq.util.doc_processor.couch import CouchDocumentProvider
from corehq.util.doc_processor.interface import (
    BaseDocProcessor,
    DocumentProcessorController,
)

SYNCLOG_SQL_USER_SYNC_GROUP_ID = "synclog_sql_user_sync"

def get_user_sync_history_pillow(
        pillow_id='UpdateUserSyncHistoryPillow', num_processes=1, process_num=0, **kwargs):
    """
    This gets a pillow which iterates through all synclogs
    """
    change_feed = KafkaChangeFeed(
        topics=[topics.SYNCLOG_SQL], client_id=SYNCLOG_SQL_USER_SYNC_GROUP_ID,
        num_processes=num_processes, process_num=process_num)
    checkpoint = KafkaPillowCheckpoint(pillow_id, [topics.SYNCLOG_SQL])
    return ConstructedPillow(
        name=pillow_id,
        checkpoint=checkpoint,
        change_feed=change_feed,
        processor=UserSyncHistoryProcessor(),
        change_processed_event_handler=KafkaCheckpointEventHandler(
            checkpoint=checkpoint, checkpoint_frequency=100, change_feed=change_feed
        ),
    )

class UserSyncHistoryProcessor(PillowProcessor):

    def process_change(self, change):
        synclog = change.get_document()
        if not synclog:
            return

        user_id = synclog.get('user_id')
        domain = synclog.get('domain')

        if not user_id or not domain:
            return

        try:
            sync_date = string_to_utc_datetime(synclog.get('date'))
        except (ValueError, AttributeError):
            return

        build_id = synclog.get('build_id')
        device_id = synclog.get('device_id')
        app_id = synclog.get('app_id')

        if settings.USER_REPORTING_METADATA_BATCH_ENABLED:
            UserReportingMetadataStaging.add_sync(domain, user_id, app_id, build_id, sync_date, device_id)
        else:
            mark_last_synclog(domain, user_id, app_id, build_id, sync_date, device_id)

def mark_last_synclog(domain, user_id, app_id, build_id, sync_date, device_id):
    user = CouchUser.get_by_user_id(user_id)
    if not user:
        return

    version = None
    if build_id:
        version = get_version_from_build_id(domain, build_id)

    save = update_last_sync(user, app_id, sync_date, version)
    if version:
        save |= update_latest_builds(user, app_id, sync_date, version)

    app_meta = None
    if device_id:
        if app_id:
            app_meta = DeviceAppMeta(app_id=app_id, build_id=build_id, last_sync=sync_date)
        save |= update_device_meta(user, device_id, device_app_meta=app_meta, save=False)

    if save:
        user.save(fire_signals=False)

class UserSyncHistoryReindexerDocProcessor(BaseDocProcessor):

    def __init__(self, pillow_processor):
        self.pillow_processor = pillow_processor

    def process_doc(self, doc):
        synclog_changes = self._doc_to_changes(doc)
        for change in synclog_changes:
            try:
                self.pillow_processor.process_change(change)
            except Exception:
                return False
        return True

    def handle_skip(self, doc):
        print('Unable to process user {}'.format(
            doc['_id'],
        ))
        return True

    def _doc_to_changes(self, doc):
        # creates a change object for the last 10 synclogs
        # of the given user, for the synclog pillow to process.
        # this means we wont have to iterate through all synclogs
        # when reindexing.
        synclogs = SyncLogSQL.objects.filter(user_id=doc['_id']).order_by('date')[:10]
        changes = [Change(
            id=res.doc['_id'],
            sequence_id=None,
            document=res.doc
        ) for res in synclogs]
        return changes

class UserSyncHistoryReindexer(Reindexer):

    def __init__(self, doc_provider, chunk_size=1000, reset=False):
        self.reset = reset
        self.doc_provider = doc_provider
        self.chunk_size = chunk_size
        self.doc_processor = UserSyncHistoryReindexerDocProcessor(UserSyncHistoryProcessor())

    def reindex(self):
        processor = DocumentProcessorController(
            self.doc_provider,
            self.doc_processor,
            reset=self.reset,
            chunk_size=self.chunk_size,
        )
        processor.run()

class UpdateUserSyncHistoryReindexerFactory(ReindexerFactory):
    slug = 'user-sync-history'
    arg_contributors = [
        ReindexerFactory.resumable_reindexer_args,
    ]

    def build(self):
        iteration_key = "UpdateUserSyncHistoryPillow_reindexer"
        doc_provider = CouchDocumentProvider(iteration_key, doc_type_tuples=[
            CommCareUser,
            WebUser
        ])
        return UserSyncHistoryReindexer(doc_provider, **self.options)
<commit_msg>Process WebApps syncs synchronously<commit_after>from django.conf import settings

from casexml.apps.phone.models import SyncLogSQL
from dimagi.utils.parsing import string_to_utc_datetime
from pillowtop.checkpoints.manager import KafkaPillowCheckpoint
from pillowtop.feed.interface import Change
from pillowtop.pillow.interface import ConstructedPillow
from pillowtop.processors.interface import PillowProcessor
from pillowtop.reindexer.reindexer import Reindexer, ReindexerFactory

from corehq.apps.change_feed import topics
from corehq.apps.change_feed.consumer.feed import (
    KafkaChangeFeed,
    KafkaCheckpointEventHandler,
)
from corehq.apps.receiverwrapper.util import get_version_from_build_id
from corehq.apps.users.models import (
    CommCareUser,
    CouchUser,
    DeviceAppMeta,
    UserReportingMetadataStaging,
    WebUser,
)
from corehq.apps.users.util import (
    update_device_meta,
    update_last_sync,
    update_latest_builds,
)
from corehq.util.doc_processor.couch import CouchDocumentProvider
from corehq.util.doc_processor.interface import (
    BaseDocProcessor,
    DocumentProcessorController,
)

SYNCLOG_SQL_USER_SYNC_GROUP_ID = "synclog_sql_user_sync"

def get_user_sync_history_pillow(
        pillow_id='UpdateUserSyncHistoryPillow', num_processes=1, process_num=0, **kwargs):
    """
    This gets a pillow which iterates through all synclogs
    """
    change_feed = KafkaChangeFeed(
        topics=[topics.SYNCLOG_SQL], client_id=SYNCLOG_SQL_USER_SYNC_GROUP_ID,
        num_processes=num_processes, process_num=process_num)
    checkpoint = KafkaPillowCheckpoint(pillow_id, [topics.SYNCLOG_SQL])
    return ConstructedPillow(
        name=pillow_id,
        checkpoint=checkpoint,
        change_feed=change_feed,
        processor=UserSyncHistoryProcessor(),
        change_processed_event_handler=KafkaCheckpointEventHandler(
            checkpoint=checkpoint, checkpoint_frequency=100, change_feed=change_feed
        ),
    )

class UserSyncHistoryProcessor(PillowProcessor):

    def process_change(self, change):
        synclog = change.get_document()
        if not synclog:
            return

        user_id = synclog.get('user_id')
        domain = synclog.get('domain')

        if not user_id or not domain:
            return

        try:
            sync_date = string_to_utc_datetime(synclog.get('date'))
        except (ValueError, AttributeError):
            return

        build_id = synclog.get('build_id')
        device_id = synclog.get('device_id')
        app_id = synclog.get('app_id')

        # WebApps syncs do not provide the app_id.
        # For those syncs we go ahead and mark the last synclog synchronously.
        if app_id and settings.USER_REPORTING_METADATA_BATCH_ENABLED:
            UserReportingMetadataStaging.add_sync(domain, user_id, app_id, build_id, sync_date, device_id)
        else:
            mark_last_synclog(domain, user_id, app_id, build_id, sync_date, device_id)

def mark_last_synclog(domain, user_id, app_id, build_id, sync_date, device_id):
    user = CouchUser.get_by_user_id(user_id)
    if not user:
        return

    version = None
    if build_id:
        version = get_version_from_build_id(domain, build_id)

    save = update_last_sync(user, app_id, sync_date, version)
    if version:
        save |= update_latest_builds(user, app_id, sync_date, version)

    app_meta = None
    if device_id:
        if app_id:
            app_meta = DeviceAppMeta(app_id=app_id, build_id=build_id, last_sync=sync_date)
        save |= update_device_meta(user, device_id, device_app_meta=app_meta, save=False)

    if save:
        user.save(fire_signals=False)

class UserSyncHistoryReindexerDocProcessor(BaseDocProcessor):

    def __init__(self, pillow_processor):
        self.pillow_processor = pillow_processor

    def process_doc(self, doc):
        synclog_changes = self._doc_to_changes(doc)
        for change in synclog_changes:
            try:
                self.pillow_processor.process_change(change)
            except Exception:
                return False
        return True

    def handle_skip(self, doc):
        print('Unable to process user {}'.format(
            doc['_id'],
        ))
        return True

    def _doc_to_changes(self, doc):
        # creates a change object for the last 10 synclogs
        # of the given user, for the synclog pillow to process.
        # this means we wont have to iterate through all synclogs
        # when reindexing.
        synclogs = SyncLogSQL.objects.filter(user_id=doc['_id']).order_by('date')[:10]
        changes = [Change(
            id=res.doc['_id'],
            sequence_id=None,
            document=res.doc
        ) for res in synclogs]
        return changes

class UserSyncHistoryReindexer(Reindexer):

    def __init__(self, doc_provider, chunk_size=1000, reset=False):
        self.reset = reset
        self.doc_provider = doc_provider
        self.chunk_size = chunk_size
        self.doc_processor = UserSyncHistoryReindexerDocProcessor(UserSyncHistoryProcessor())

    def reindex(self):
        processor = DocumentProcessorController(
            self.doc_provider,
            self.doc_processor,
            reset=self.reset,
            chunk_size=self.chunk_size,
        )
        processor.run()

class UpdateUserSyncHistoryReindexerFactory(ReindexerFactory):
    slug = 'user-sync-history'
    arg_contributors = [
        ReindexerFactory.resumable_reindexer_args,
    ]

    def build(self):
        iteration_key = "UpdateUserSyncHistoryPillow_reindexer"
        doc_provider = CouchDocumentProvider(iteration_key, doc_type_tuples=[
            CommCareUser,
            WebUser
        ])
        return UserSyncHistoryReindexer(doc_provider, **self.options)
<|endoftext|>


<commit_before>"""Naming conventions for abd."""
# =============================================================================
# CONTENTS
# -----------------------------------------------------------------------------
# abdt_naming
#
# Public Classes:
#   Error
#   ClassicNaming
#    .make_tracker_branch_from_name
#    .make_tracker_branch_name
#    .make_review_branch_from_name
#    .make_review_branch_name
#
# Public Functions:
#   getReviewBranchPrefix
#   getWorkingBranchPrefix
#   getReservedBranchPrefix
#   isStatusBad
#   isStatusBadPreReview
#   isStatusBadLand
#   isReviewBranchPrefixed
#   makeReviewBranchNameFromWorkingBranch
#   getWorkingBranches
#
# Public Assignments:
#   WB_STATUS_OK
#   WB_STATUS_PREFIX_BAD
#   WB_STATUS_BAD_NAME
#   WB_STATUS_BAD_PREREVIEW
#   WB_STATUS_BAD_INREVIEW
#   WB_STATUS_BAD_LAND
#   WB_DICT_STATUS_DESC
#   WorkingBranch
#   ReviewBranch
#
# -----------------------------------------------------------------------------
# (this contents block is generated, edits will be lost)
# =============================================================================

from __future__ import absolute_import

import collections

import phlsys_string

WB_STATUS_OK = "ok"
WB_STATUS_PREFIX_BAD = "bad_"
WB_STATUS_BAD_NAME = WB_STATUS_PREFIX_BAD + "name"
WB_STATUS_BAD_PREREVIEW = WB_STATUS_PREFIX_BAD + "prerev"
WB_STATUS_BAD_INREVIEW = WB_STATUS_PREFIX_BAD + "inrev"
WB_STATUS_BAD_LAND = WB_STATUS_PREFIX_BAD + "land"

WB_DICT_STATUS_DESC = {
    WB_STATUS_OK:            "ok",
    WB_STATUS_BAD_NAME:      "branch name is invalid",
    WB_STATUS_BAD_PREREVIEW: "didn't manage to create a review",
    WB_STATUS_BAD_INREVIEW:  "bad update during review",
    WB_STATUS_BAD_LAND:      "didn't manage to land the change",
}

class Error(Exception):
    pass

def getReviewBranchPrefix():
    # don't see a reason to change this atm
    return "arcyd-review/"

def getWorkingBranchPrefix():
    # this may want to be configurable from the command-line, we'll probably
    # want to wrap the functions in this module in a class and store the
    # convention variables per-instance; will get the whole thing working
    # first and worry about that later.
    return "dev/arcyd/"

def getReservedBranchPrefix():
    # this may want to be configurable from the command-line, we'll probably
    # want to wrap the functions in this module in a class and store the
    # convention variables per-instance; will get the whole thing working
    # first and worry about that later.
    return "dev/arcyd/reserve"

def isStatusBad(working_branch):
    """Return True if the status of 'working_branch' is bad.

    :working_branch: a WorkingBranch
    :returns: True if the branch is bad

    """
    return working_branch.status.startswith(WB_STATUS_PREFIX_BAD)

def isStatusBadPreReview(working_branch):
    """Return True if 'working_branch' status is WB_STATUS_BAD_PREREVIEW.

    :working_branch: a WorkingBranch
    :returns: True if the branch is in WB_STATUS_BAD_PREREVIEW

    """
    return working_branch.status == WB_STATUS_BAD_PREREVIEW

def isStatusBadLand(working_branch):
    """Return True if 'working_branch' status is WB_STATUS_BAD_PREREVIEW.

    :working_branch: a WorkingBranch
    :returns: True if the branch is in WB_STATUS_BAD_PREREVIEW

    """
    return working_branch.status == WB_STATUS_BAD_LAND

def isReviewBranchPrefixed(name):
    prefix = getReviewBranchPrefix()
    return (len(name) > len(prefix)) and name.startswith(prefix)

WorkingBranch = collections.namedtuple(
    "abdt_naming__WorkingBranch", [
        "branch",
        "status",
        "description",
        "base",
        "id"])

ReviewBranch = collections.namedtuple(
    "abdt_naming__ReviewBranch", [
        "branch",
        "description",
        "base"])

def makeReviewBranchNameFromWorkingBranch(working_branch):
    """Return the string review branch name for 'working_branch'.

    :working_branch: a WorkingBranch
    :returns: the string review branch name for 'working_branch'

    """
    branch_name = getReviewBranchPrefix()
    branch_name += working_branch.description
    branch_name += "/" + working_branch.base
    return branch_name

class ClassicNaming(object):

    def __init__(self):
        super(ClassicNaming, self).__init__()
        self._tracking_branch_prefix = 'dev/arcyd/'
        self._reserve_branch_prefix = 'dev/arcyd/reserve'
        self._review_branch_prefix = 'arcyd-review/'

    def make_tracker_branch_from_name(self, branch_name):
        """Return the WorkingBranch for 'branch_name' or None if invalid.

        Usage example:
            >>> naming = ClassicNaming()
            >>> make_branch = naming.make_tracker_branch_from_name
            >>> make_branch('dev/arcyd/ok/mywork/master/99')
            ... # doctest: +NORMALIZE_WHITESPACE
            abdt_naming__WorkingBranch(branch='dev/arcyd/ok/mywork/master/99',
                                    status='ok',
                                    description='mywork',
                                    base='master',
                                    id='99')

            >>> make_branch('invalid/mywork/master')
            Traceback (most recent call last):
                ...
            Error

        :branch_name: string name of the working branch
        :returns: WorkingBranch or None if invalid

        """
        if branch_name == self._reserve_branch_prefix:
            raise Error()  # ignore the reserved branch

        suffix = phlsys_string.after_prefix(
            branch_name, self._tracking_branch_prefix)

        if not suffix:
            raise Error()  # review branches must start with the prefix

        parts = suffix.split("/")
        if len(parts) < 4:
            raise Error()  # suffix should be status/description/base(/...)/id

        return WorkingBranch(
            branch=branch_name,
            status=parts[0],
            description=parts[1],
            base='/'.join(parts[2:-1]),
            id=parts[-1])

    def make_tracker_branch_name(self, status, description, base, review_id):
        """Return the unique string name of the tracker branch for params.

        Working branches are of the form:
            <working branch prefix>/description/base

        Usage example:
            >>> naming = ClassicNaming()
            >>> make_name = naming.make_tracker_branch_name
            >>> make_name('ok', 'mywork', 'master',  99)
            'dev/arcyd/ok/mywork/master/99'

        :description: string descriptive name of the branch
        :base: string name of the branch to diff against and land on
        :id: identifier for the review, converted to str() for convenience
        :returns: string name of the working branch

        """
        tracker_branch = ""
        tracker_branch += self._tracking_branch_prefix
        tracker_branch += status
        tracker_branch += "/" + description
        tracker_branch += "/" + base
        tracker_branch += "/" + str(review_id)
        return tracker_branch

    def make_review_branch_from_name(self, branch_name):
        """Return the ReviewBranch for 'branch_name' or None if invalid.

        Usage example:
            >>> naming = ClassicNaming()
            >>> make_branch = naming.make_review_branch_from_name
            >>> make_branch('arcyd-review/mywork/master')
            ... # doctest: +NORMALIZE_WHITESPACE
            abdt_naming__ReviewBranch(branch='arcyd-review/mywork/master',
                                    description='mywork',
                                    base='master')

            >>> make_branch('invalid/mywork/master')
            Traceback (most recent call last):
                ...
            Error

        :branch_name: string name of the review branch
        :returns: ReviewBranch or None if invalid

        """
        suffix = phlsys_string.after_prefix(
            branch_name, getReviewBranchPrefix())
        if not suffix:
            raise Error()  # review branches must start with the prefix

        parts = suffix.split("/")
        if len(parts) < 2:
            raise Error()  # suffix should be description/base(/...)

        return ReviewBranch(
            branch=branch_name,
            description=parts[0],
            base='/'.join(parts[1:]))

    def make_review_branch_name(self, description, base):
        """Return the unique string name of the review branch for these params.

        Review branches are of the form:
            <review branch prefix>/description/base

        Usage Example:
            >>> naming = ClassicNaming()
            >>> make_name = naming.make_review_branch_name
            >>> make_name('mywork', 'master')
            'arcyd-review/mywork/master'

        :description: string descriptive name of the branch
        :base: string name of the branch to diff against and land on
        :returns: string name of the review branch

        """
        branch_name = self._review_branch_prefix
        branch_name += description
        branch_name += "/" + base
        return branch_name

def getWorkingBranches(branch_list):
    """Return a list of WorkingBranch made from strings in 'branch_list'.

    Strings that aren't valid working branch names are ignored.

    Usage example:
        >>> getWorkingBranches(['dev/arcyd/ok/mywork/master/99'])
        ... # doctest: +NORMALIZE_WHITESPACE
        [abdt_naming__WorkingBranch(branch='dev/arcyd/ok/mywork/master/99',
                                   status='ok',
                                   description='mywork',
                                   base='master',
                                   id='99')]

        >>> getWorkingBranches([])
        []

        >>> getWorkingBranches(['invalid'])
        []

    :branch_list: list of branch name strings
    :returns: list of WorkingBranch

    """
    working_branch_list = []
    naming = ClassicNaming()
    for branch in branch_list:
        try:
            working_branch_list.append(
                naming.make_tracker_branch_from_name(branch))
        except Error:
            pass  # ignore naming errors, we only want the valid branches
    return working_branch_list

#------------------------------------------------------------------------------
# Copyright (C) 2012 Bloomberg L.P.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#------------------------------- END-OF-FILE ----------------------------------
<commit_msg>abdt_naming: refactor, use _review_branch_prefix<commit_after>"""Naming conventions for abd."""
# =============================================================================
# CONTENTS
# -----------------------------------------------------------------------------
# abdt_naming
#
# Public Classes:
#   Error
#   ClassicNaming
#    .make_tracker_branch_from_name
#    .make_tracker_branch_name
#    .make_review_branch_from_name
#    .make_review_branch_name
#
# Public Functions:
#   getReviewBranchPrefix
#   getWorkingBranchPrefix
#   getReservedBranchPrefix
#   isStatusBad
#   isStatusBadPreReview
#   isStatusBadLand
#   isReviewBranchPrefixed
#   makeReviewBranchNameFromWorkingBranch
#   getWorkingBranches
#
# Public Assignments:
#   WB_STATUS_OK
#   WB_STATUS_PREFIX_BAD
#   WB_STATUS_BAD_NAME
#   WB_STATUS_BAD_PREREVIEW
#   WB_STATUS_BAD_INREVIEW
#   WB_STATUS_BAD_LAND
#   WB_DICT_STATUS_DESC
#   WorkingBranch
#   ReviewBranch
#
# -----------------------------------------------------------------------------
# (this contents block is generated, edits will be lost)
# =============================================================================

from __future__ import absolute_import

import collections

import phlsys_string

WB_STATUS_OK = "ok"
WB_STATUS_PREFIX_BAD = "bad_"
WB_STATUS_BAD_NAME = WB_STATUS_PREFIX_BAD + "name"
WB_STATUS_BAD_PREREVIEW = WB_STATUS_PREFIX_BAD + "prerev"
WB_STATUS_BAD_INREVIEW = WB_STATUS_PREFIX_BAD + "inrev"
WB_STATUS_BAD_LAND = WB_STATUS_PREFIX_BAD + "land"

WB_DICT_STATUS_DESC = {
    WB_STATUS_OK:            "ok",
    WB_STATUS_BAD_NAME:      "branch name is invalid",
    WB_STATUS_BAD_PREREVIEW: "didn't manage to create a review",
    WB_STATUS_BAD_INREVIEW:  "bad update during review",
    WB_STATUS_BAD_LAND:      "didn't manage to land the change",
}

class Error(Exception):
    pass

def getReviewBranchPrefix():
    # don't see a reason to change this atm
    return "arcyd-review/"

def getWorkingBranchPrefix():
    # this may want to be configurable from the command-line, we'll probably
    # want to wrap the functions in this module in a class and store the
    # convention variables per-instance; will get the whole thing working
    # first and worry about that later.
    return "dev/arcyd/"

def getReservedBranchPrefix():
    # this may want to be configurable from the command-line, we'll probably
    # want to wrap the functions in this module in a class and store the
    # convention variables per-instance; will get the whole thing working
    # first and worry about that later.
    return "dev/arcyd/reserve"

def isStatusBad(working_branch):
    """Return True if the status of 'working_branch' is bad.

    :working_branch: a WorkingBranch
    :returns: True if the branch is bad

    """
    return working_branch.status.startswith(WB_STATUS_PREFIX_BAD)

def isStatusBadPreReview(working_branch):
    """Return True if 'working_branch' status is WB_STATUS_BAD_PREREVIEW.

    :working_branch: a WorkingBranch
    :returns: True if the branch is in WB_STATUS_BAD_PREREVIEW

    """
    return working_branch.status == WB_STATUS_BAD_PREREVIEW

def isStatusBadLand(working_branch):
    """Return True if 'working_branch' status is WB_STATUS_BAD_PREREVIEW.

    :working_branch: a WorkingBranch
    :returns: True if the branch is in WB_STATUS_BAD_PREREVIEW

    """
    return working_branch.status == WB_STATUS_BAD_LAND

def isReviewBranchPrefixed(name):
    prefix = getReviewBranchPrefix()
    return (len(name) > len(prefix)) and name.startswith(prefix)

WorkingBranch = collections.namedtuple(
    "abdt_naming__WorkingBranch", [
        "branch",
        "status",
        "description",
        "base",
        "id"])

ReviewBranch = collections.namedtuple(
    "abdt_naming__ReviewBranch", [
        "branch",
        "description",
        "base"])

def makeReviewBranchNameFromWorkingBranch(working_branch):
    """Return the string review branch name for 'working_branch'.

    :working_branch: a WorkingBranch
    :returns: the string review branch name for 'working_branch'

    """
    branch_name = getReviewBranchPrefix()
    branch_name += working_branch.description
    branch_name += "/" + working_branch.base
    return branch_name

class ClassicNaming(object):

    def __init__(self):
        super(ClassicNaming, self).__init__()
        self._tracking_branch_prefix = 'dev/arcyd/'
        self._reserve_branch_prefix = 'dev/arcyd/reserve'
        self._review_branch_prefix = 'arcyd-review/'

    def make_tracker_branch_from_name(self, branch_name):
        """Return the WorkingBranch for 'branch_name' or None if invalid.

        Usage example:
            >>> naming = ClassicNaming()
            >>> make_branch = naming.make_tracker_branch_from_name
            >>> make_branch('dev/arcyd/ok/mywork/master/99')
            ... # doctest: +NORMALIZE_WHITESPACE
            abdt_naming__WorkingBranch(branch='dev/arcyd/ok/mywork/master/99',
                                    status='ok',
                                    description='mywork',
                                    base='master',
                                    id='99')

            >>> make_branch('invalid/mywork/master')
            Traceback (most recent call last):
                ...
            Error

        :branch_name: string name of the working branch
        :returns: WorkingBranch or None if invalid

        """
        if branch_name == self._reserve_branch_prefix:
            raise Error()  # ignore the reserved branch

        suffix = phlsys_string.after_prefix(
            branch_name, self._tracking_branch_prefix)

        if not suffix:
            raise Error()  # review branches must start with the prefix

        parts = suffix.split("/")
        if len(parts) < 4:
            raise Error()  # suffix should be status/description/base(/...)/id

        return WorkingBranch(
            branch=branch_name,
            status=parts[0],
            description=parts[1],
            base='/'.join(parts[2:-1]),
            id=parts[-1])

    def make_tracker_branch_name(self, status, description, base, review_id):
        """Return the unique string name of the tracker branch for params.

        Working branches are of the form:
            <working branch prefix>/description/base

        Usage example:
            >>> naming = ClassicNaming()
            >>> make_name = naming.make_tracker_branch_name
            >>> make_name('ok', 'mywork', 'master',  99)
            'dev/arcyd/ok/mywork/master/99'

        :description: string descriptive name of the branch
        :base: string name of the branch to diff against and land on
        :id: identifier for the review, converted to str() for convenience
        :returns: string name of the working branch

        """
        tracker_branch = ""
        tracker_branch += self._tracking_branch_prefix
        tracker_branch += status
        tracker_branch += "/" + description
        tracker_branch += "/" + base
        tracker_branch += "/" + str(review_id)
        return tracker_branch

    def make_review_branch_from_name(self, branch_name):
        """Return the ReviewBranch for 'branch_name' or None if invalid.

        Usage example:
            >>> naming = ClassicNaming()
            >>> make_branch = naming.make_review_branch_from_name
            >>> make_branch('arcyd-review/mywork/master')
            ... # doctest: +NORMALIZE_WHITESPACE
            abdt_naming__ReviewBranch(branch='arcyd-review/mywork/master',
                                    description='mywork',
                                    base='master')

            >>> make_branch('invalid/mywork/master')
            Traceback (most recent call last):
                ...
            Error

        :branch_name: string name of the review branch
        :returns: ReviewBranch or None if invalid

        """
        suffix = phlsys_string.after_prefix(
            branch_name, self._review_branch_prefix)
        if not suffix:
            raise Error()  # review branches must start with the prefix

        parts = suffix.split("/")
        if len(parts) < 2:
            raise Error()  # suffix should be description/base(/...)

        return ReviewBranch(
            branch=branch_name,
            description=parts[0],
            base='/'.join(parts[1:]))

    def make_review_branch_name(self, description, base):
        """Return the unique string name of the review branch for these params.

        Review branches are of the form:
            <review branch prefix>/description/base

        Usage Example:
            >>> naming = ClassicNaming()
            >>> make_name = naming.make_review_branch_name
            >>> make_name('mywork', 'master')
            'arcyd-review/mywork/master'

        :description: string descriptive name of the branch
        :base: string name of the branch to diff against and land on
        :returns: string name of the review branch

        """
        branch_name = self._review_branch_prefix
        branch_name += description
        branch_name += "/" + base
        return branch_name

def getWorkingBranches(branch_list):
    """Return a list of WorkingBranch made from strings in 'branch_list'.

    Strings that aren't valid working branch names are ignored.

    Usage example:
        >>> getWorkingBranches(['dev/arcyd/ok/mywork/master/99'])
        ... # doctest: +NORMALIZE_WHITESPACE
        [abdt_naming__WorkingBranch(branch='dev/arcyd/ok/mywork/master/99',
                                   status='ok',
                                   description='mywork',
                                   base='master',
                                   id='99')]

        >>> getWorkingBranches([])
        []

        >>> getWorkingBranches(['invalid'])
        []

    :branch_list: list of branch name strings
    :returns: list of WorkingBranch

    """
    working_branch_list = []
    naming = ClassicNaming()
    for branch in branch_list:
        try:
            working_branch_list.append(
                naming.make_tracker_branch_from_name(branch))
        except Error:
            pass  # ignore naming errors, we only want the valid branches
    return working_branch_list

#------------------------------------------------------------------------------
# Copyright (C) 2012 Bloomberg L.P.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#------------------------------- END-OF-FILE ----------------------------------
<|endoftext|>


<commit_before>import gym
import time
import torch
from torch import nn, autograd, optim
import torch.nn.functional as F
import numpy as np

num_hidden = 4
print_every = 1
num_theta_samples = 10
num_samples_per_theta = 10
initial_variance = 10
num_elite = 3

class Policy(nn.Module):
    def __init__(self, num_inputs, num_hidden, num_actions):
        super().__init__()
        self.num_actions = num_actions
        self.num_inputs = num_inputs
        self.num_hidden = num_hidden
        # self.h1 = nn.Linear(num_inputs, num_hidden)
        # self.h2 = nn.Linear(num_hidden, num_actions)
        self.h1 = nn.Linear(num_inputs, num_actions)

    def forward(self, x):
        x = self.h1(x)
        x = F.softmax(x)
        multinomial_res = torch.multinomial(x, num_samples=1)
        return multinomial_res

def run_episode(policy, render):
    x = env.reset()
    reward = 0
    # multinomial_res_nodes = []
    for _ in range(1000):
        if render:
            env.render()
        a_idx = policy(autograd.Variable(torch.from_numpy(x.astype(np.float32)).view(1, -1)))
        # multinomial_res_nodes.append(a_idx)
        # a_idx = a_idx.data[0]
        # print('a_idx', a_idx.data[0])
        # a = env.action_space.sample()
        x, r, done, info = env.step(a_idx.data[0][0])
        # print('a', a, 'x', x, 'r', r, 'done', done)
        reward += r
        if done:
            break
        # time.sleep(0.02)

    # avg_reward = reward / len(multinomial_res_nodes)
    # reward = reward - 18
    # reward = reward * reward
    # for node in multinomial_res_nodes:
        # node.reinforce(reward)
    return reward

env = gym.make('CartPole-v0')
# print('action_space', env.action_space)
# print(dir(env.action_space))
# print('num_actions', env.action_space.n)
# print('num_inputs', env.observation_space.shape[0])
policy = Policy(
    num_inputs=env.observation_space.shape[0],
    num_hidden=num_hidden,
    num_actions=env.action_space.n)
# opt = optim.Adam(params=policy.parameters(), lr=0.001)
# print('observation_space', env.observation_space)

def tensor_list_to_tensor(l):
    shape = l[0].size()
    N = len(l)
    t = torch.zeros(N, *shape)
    # print('t.size()', t.size())
    for i in range(N):
        t[i] = l[i]
    return t
    # asdfsdf

def draw_theta(theta_mu, theta_var):
    """
    assumes theta_mu and theta_var are both lists of torch tensors
    thus, returns also a list, of torch tensors
    """
    theta = []
    # for i, p in enumerate(policy.parameters()):
    for i, mu in enumerate(theta_mu):
        var = theta_var[i]
        _theta = torch.normal(mu, var.sqrt())
        # p.data = this_theta
        theta.append(_theta)
    return theta

def assign_theta_to_params(params, theta):
    for i, p in enumerate(params):
        p.data.copy_(theta[i])
# def render_example():

# episode = 0
iteration = 0
sum_reward = 0
# theta = [p for p in policy.parameters()]
theta_shapes = [p.data.shape for p in policy.parameters()]
theta_mu = []
theta_var = []
for shape in theta_shapes:
    theta_mu.append(torch.zeros(shape))
    var = torch.ones(shape)
    var *= initial_variance
    print('var', var)
    theta_var.append(var)
# print('theta_shapes', theta_shapes)
# print(theta_mu)
# print(theta_var)
# asdf
while True:
    # opt.zero_grad()
    res = []
    for n in range(num_theta_samples):
        theta = draw_theta(theta_mu=theta_mu, theta_var=theta_var)
        assign_theta_to_params(params=policy.parameters(), theta=theta)
        sum_reward_this_theta = 0
        for j in range(num_samples_per_theta):
            # render = episode % print_every == 0 and n % (num_samples // 3) == 0
            reward = run_episode(policy=policy, render=False)
            sum_reward_this_theta += reward
        avg_reward_this_theta = sum_reward_this_theta / num_samples_per_theta
        res.append({'theta': theta, 'reward': avg_reward_this_theta})
    sum_reward += avg_reward_this_theta
    # print(res)
    res.sort(key=lambda x: x['reward'], reverse=True)
    elite_samples = res[:num_elite]
    # print('elite_samples', elite_samples)
    for i, p in enumerate(policy.parameters()):
        this_elite_thetas = tensor_list_to_tensor([result['theta'][i] for result in elite_samples])
        # print('this_elite_thetas', this_elite_thetas)
        this_new_mu = this_elite_thetas.mean(dim=0)
        this_new_var = this_elite_thetas.var(dim=0)
        theta_mu[i] = this_new_mu
        theta_var[i] = this_new_var

    if iteration % print_every == 0:
        print('theta_mu', theta_mu)
        print('theta_var', theta_var)
        print('iteration %s avg_reward %s' % (iteration, sum_reward / print_every))
        # render_example()
        assign_theta_to_params(params=policy.parameters(), theta=theta_mu)
        run_episode(policy=policy, render=True)
        sum_reward = 0
    iteration += 1
    # asdfasd
<commit_msg>save current cem1<commit_after>import gym
import time
import torch
from torch import nn, autograd, optim
import torch.nn.functional as F
import numpy as np

num_hidden = 4
print_every = 1
num_theta_samples = 10
num_samples_per_theta = 10
initial_variance = 10
num_elite = 3
max_steps_per_episode = 2000

class Policy(nn.Module):
    def __init__(self, num_inputs, num_hidden, num_actions):
        super().__init__()
        self.num_actions = num_actions
        self.num_inputs = num_inputs
        self.num_hidden = num_hidden
        # self.h1 = nn.Linear(num_inputs, num_hidden)
        # self.h2 = nn.Linear(num_hidden, num_actions)
        self.h1 = nn.Linear(num_inputs, num_actions)

    def forward(self, x):
        x = self.h1(x)
        x = F.softmax(x)
        multinomial_res = torch.multinomial(x, num_samples=1)
        return multinomial_res

def run_episode(policy, render):
    x = env.reset()
    reward = 0
    # multinomial_res_nodes = []
    for _ in range(max_steps_per_episode):
        if render:
            env.render()
        a_idx = policy(autograd.Variable(torch.from_numpy(x.astype(np.float32)).view(1, -1)))
        # multinomial_res_nodes.append(a_idx)
        # a_idx = a_idx.data[0]
        # print('a_idx', a_idx.data[0])
        # a = env.action_space.sample()
        x, r, done, info = env.step(a_idx.data[0][0])
        # print('a', a, 'x', x, 'r', r, 'done', done)
        reward += r
        if done:
            break
        # time.sleep(0.02)

    # avg_reward = reward / len(multinomial_res_nodes)
    # reward = reward - 18
    # reward = reward * reward
    # for node in multinomial_res_nodes:
        # node.reinforce(reward)
    return reward

env = gym.make('CartPole-v0')
# print('action_space', env.action_space)
# print(dir(env.action_space))
# print('num_actions', env.action_space.n)
# print('num_inputs', env.observation_space.shape[0])
policy = Policy(
    num_inputs=env.observation_space.shape[0],
    num_hidden=num_hidden,
    num_actions=env.action_space.n)
# opt = optim.Adam(params=policy.parameters(), lr=0.001)
# print('observation_space', env.observation_space)

def tensor_list_to_tensor(l):
    shape = l[0].size()
    N = len(l)
    t = torch.zeros(N, *shape)
    # print('t.size()', t.size())
    for i in range(N):
        t[i] = l[i]
    return t
    # asdfsdf

def draw_theta(theta_mu, theta_var):
    """
    assumes theta_mu and theta_var are both lists of torch tensors
    thus, returns also a list, of torch tensors
    """
    theta = []
    # for i, p in enumerate(policy.parameters()):
    for i, mu in enumerate(theta_mu):
        var = theta_var[i]
        _theta = torch.normal(mu, var.sqrt())
        # p.data = this_theta
        theta.append(_theta)
    return theta

def assign_theta_to_params(params, theta):
    for i, p in enumerate(params):
        p.data.copy_(theta[i])
# def render_example():

# episode = 0
iteration = 0
sum_reward = 0
# theta = [p for p in policy.parameters()]
theta_shapes = [p.data.shape for p in policy.parameters()]
theta_mu = []
theta_var = []
for shape in theta_shapes:
    theta_mu.append(torch.zeros(shape))
    var = torch.ones(shape)
    var *= initial_variance
    print('var', var)
    theta_var.append(var)
# print('theta_shapes', theta_shapes)
# print(theta_mu)
# print(theta_var)
# asdf
while True:
    # opt.zero_grad()
    res = []
    for n in range(num_theta_samples):
        theta = draw_theta(theta_mu=theta_mu, theta_var=theta_var)
        assign_theta_to_params(params=policy.parameters(), theta=theta)
        sum_reward_this_theta = 0
        for j in range(num_samples_per_theta):
            # render = episode % print_every == 0 and n % (num_samples // 3) == 0
            reward = run_episode(policy=policy, render=False)
            sum_reward_this_theta += reward
        avg_reward_this_theta = sum_reward_this_theta / num_samples_per_theta
        res.append({'theta': theta, 'reward': avg_reward_this_theta})
        sum_reward += avg_reward_this_theta
    # print(res)
    res.sort(key=lambda x: x['reward'], reverse=True)
    elite_samples = res[:num_elite]
    # print('elite_samples', elite_samples)
    for i, p in enumerate(policy.parameters()):
        this_elite_thetas = tensor_list_to_tensor([result['theta'][i] for result in elite_samples])
        # print('this_elite_thetas', this_elite_thetas)
        this_new_mu = this_elite_thetas.mean(dim=0)
        this_new_var = this_elite_thetas.var(dim=0)
        theta_mu[i] = this_new_mu
        theta_var[i] = this_new_var

    if iteration % print_every == 0:
        print('theta_mu', theta_mu)
        print('theta_var', theta_var)
        print('iteration %s avg_reward %s' % (iteration, sum_reward / print_every))
        # render_example()
        assign_theta_to_params(params=policy.parameters(), theta=theta_mu)
        run_episode(policy=policy, render=True)
        sum_reward = 0
    iteration += 1
    # asdfasd
<|endoftext|>


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

import binascii
import logging

from dcerpc import MSRPCRequestHeader, MSRPCRespHeader
import kmsBase
import rpcBase
from formatText import justify, shell_message, byterize

logger = logging.getLogger('root')

class handler(rpcBase.rpcBase):
        def parseRequest(self):
                request = MSRPCRequestHeader(self.data)
                shell_message(nshell = 14)
                request = byterize(request)
                logger.debug("RPC Message Request Bytes: \n%s\n" % justify(binascii.b2a_hex(self.data).decode('utf-8')))
                logger.debug("RPC Message Request: \n%s\n" % justify(request.dump(print_to_stdout = False)))

                return request

        def generateResponse(self, request):
                responseData = kmsBase.generateKmsResponseData(request['pduData'], self.config)
                envelopeLength = len(responseData)

                response = MSRPCRespHeader()
                response['ver_major'] = request['ver_major']
                response['ver_minor'] = request['ver_minor']
                response['type'] = self.packetType['response']
                response['flags'] = self.packetFlags['firstFrag'] | self.packetFlags['lastFrag']
                response['representation'] = request['representation']
                response['call_id'] = request['call_id']

                response['alloc_hint'] = envelopeLength
                response['ctx_id'] = request['ctx_id']
                response['cancel_count'] = 0

                response['pduData'] = responseData

                shell_message(nshell = 17)
                response = byterize(response)
                logger.debug("RPC Message Response: \n%s\n" % justify(response.dump(print_to_stdout = False)))
                logger.debug("RPC Message Response Bytes: \n%s\n" % justify(binascii.b2a_hex(str(response).encode('latin-1')).decode('utf-8')))

                return response

        def generateRequest(self):
                request = MSRPCRequestHeader()

                request['ver_major'] = 5
                request['ver_minor'] = 0
                request['type'] = self.packetType['request']
                request['flags'] = self.packetFlags['firstFrag'] | self.packetFlags['lastFrag']
                request['representation'] = 0x10
                request['call_id'] = self.config['call_id']
                request['alloc_hint'] = len(self.data)
                request['pduData'] = str(self.data)

                shell_message(nshell = 11)
                request = byterize(request)
                logger.debug("RPC Message Request: \n%s\n" % justify(request.dump(print_to_stdout = False)))
                logger.debug("RPC Message Request Bytes: \n%s\n" % justify(binascii.b2a_hex(str(request).encode('latin-1')).decode('utf-8')))

                return request

        def parseResponse(self):
                return response
<commit_msg>Delete rpcRequest.py<commit_after><|endoftext|>


<commit_before>__author__ = 'Ian S. Evans'

import json
from collections import UserDict
from math import ceil

MIMETYPE = 'application/vnd.collection+json'

class CollectionPlusJSON(UserDict):

    mimetype = MIMETYPE

    class BaseCollectionItem(object):

        def __init__(self, *args, **kwargs):
            pass

        def __str__(self):
            return json.dumps(self.__dict__)

        def __repr__(self):
            return self.__str__()

    class Error(BaseCollectionItem):

        def __init__(self, title=None, code=None, message=None):
            if title is not None:
                self.title = title
            if code is not None:
                self.code = code
            if message is not None:
                self.message = message
            try:
                # Python 3
                super().__init__()
            except ValueError:
                super(CollectionPlusJSON.Error, self).__init__()

    class Item(UserDict, BaseCollectionItem):

        def __init__(self, uri=None, links=None, **kwargs):
            if uri is not None:
                self.href = uri
            if links is not None:
                self.links = links
            try:
                # Python 3
                super().__init__(**kwargs)
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Item, self).__init__(**kwargs)

        def __str__(self):
            item_dict = dict(self.__dict__)
            if self.data:
                data_list = []
                for k, v in self.data.items():
                    to_append = {"name": k}
                    if not isinstance(v, str):
                        to_append["value"] = v[0]
                        to_append["prompt"] = v[1]
                    else:
                        to_append['value'] = v
                    data_list.append(to_append)
                item_dict['data'] = data_list
            return json.dumps(item_dict)

    class Link(BaseCollectionItem):

        def __init__(self, uri, rel, name=None, render=None, prompt=None):
            self.uri = uri
            self.rel = rel
            if name is not None:
                self.name = name
            if render is not None:
                self.render = render
            if prompt is not None:
                self.prompt = prompt
            try:
                # Python 3
                super().__init__()
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Link, self).__init__()

    class Query(Item):

        def __init__(self, uri, rel=None, prompt=None, **kwargs):
            self.uri = uri
            if rel is not None:
                self.rel = rel
            if prompt is not None:
                self.prompt = prompt
            try:
                # Python 3
                super().__init__(**kwargs)
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Query, self).__init__(**kwargs)

    class Template(Item):

        def __init__(self, **kwargs):
            try:
                # Python 3
                super().__init__(**kwargs)
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Template, self).__init__(**kwargs)

    def __init__(self, version=1.0, href='/api/', **kwargs):
        collection = {
            'version': str(version),
            'href': href,
        }
        collection = dict(collection, **kwargs)  # zip together standard collection + extended properties
        try:
            # Python 3
            super().__init__(**collection)
        except TypeError:
            # Python 2
            super(CollectionPlusJSON, self).__init__(**collection)

    def __str__(self):
        return json.dumps({'collection': self.data})

    def __repr__(self):
        return self.__str__()

    def __setitem__(self, key, value):
        type_map = {
            'links': list,
            'items': list,
            'queries': list,
            'template': self.Template,
            'error': self.Error
        }
        if key in type_map.keys():
            if not isinstance(value, type_map[key]):
                raise TypeError(
                    '{key} must be an instance of {classname}'.format(key=key, classname=str(type_map[key]))
                )
        self.data[key] = value

    def append_item(self, item):
        """
        Append an item to this collection's 'items' property.
        :param item: The Item to append.
        :return:
        """
        if not isinstance(item, self.Item):
            raise TypeError('item must be an instance of {classname}'.format(classname=str(self.Item)))
        else:
            if self.data.get('items') is not None:
                self.data.get('items').append(item)
            else:
                self.data['items'] = [item]

    def append_link(self, link):
        """
        Append a link to this collection's 'links' property.
        :param link: The Link to append.
        :return:
        """
        if not isinstance(link, self.Link):
            raise TypeError('item must be an instance of {classname}'.format(classname=str(self.Link)))
        else:
            if self.data.get('links') is not None:
                self.data.get('links').append(link)
            else:
                self.data['links'] = [link]

    def append_query(self, query):
        """
        Append a query to this collection's 'queries' property.
        :param query: The Query to append.
        :return:
        """
        if not isinstance(query, self.Query):
            raise TypeError('item must be an instance of {classname}'.format(classname=str(self.Query)))
        else:
            if self.data.get('queries') is not None:
                self.data.get('queries').append(query)
            else:
                self.data['queries'] = [query]

    def paginate(self, endpoint='', uri_template='{endpoint_uri}?page={page}&per_page={per_page}', page=1, per_page=5,
                 leading=2, trailing=2):
        """
        Paginate this collection, automatically trimming items and adding appropriate links for navigation.
        :param endpoint: The URI for this resource.
        :param uri_template: A string providing a template for paginated URI structure. May include the following keys:
        "{endpoint_uri}" - This will evaluate to the value of the 'endpoint' param.
        "{page}" - The page number will be inserted here.
        "{per_page}" - The number of items to display per page will be inserted here.
        :param page: The desired page for this representation.
        :param per_page: The number of items per page for this representation.
        :param leading: The number of leading pages before "this" page to add to this collection's "links".
        :param trailing: The number of trailing pages after "this" page to add to this collection's "links".
        :return:
        """
        if (type(page) is not int) or (type(per_page) is not int):

            try:
                page = abs(int(page))
            except (ValueError, TypeError):
                page = 1

            try:
                per_page = abs(int(per_page))
            except (ValueError, TypeError):
                per_page = 5

            if not page:
                page = 1
            if not per_page:
                per_page = 5

        number_of_pages = int(ceil(len(self.data.get('items')) / per_page))

        if page > number_of_pages:
            page = number_of_pages

        page_index_begin = ((page * per_page) - per_page)
        page_index_end = (page * per_page)

        self.data['items'] = self.data.get('items')[page_index_begin:page_index_end]

        if page > 1:
            self.append_link(self.Link(
                uri_template.format(endpoint_uri=endpoint, page=1, per_page=per_page),
                'first',
                prompt='First'
            ))

            self.append_link(self.Link(
                uri_template.format(endpoint_uri=endpoint, page=(page - 1), per_page=per_page),
                'prev',
                prompt='Previous'
            ))

            if page - leading > 0:
                self.append_link(self.Link(
                    '',
                    'skip',
                    prompt='…'
                ))

            for lead_page in range(leading):
                page_num = page - lead_page
                if page_num > 0:
                    self.append_link(self.Link(
                        uri_template.format(endpoint_uri=endpoint, page=page_num, per_page=per_page),
                        'more',
                        prompt=str(page_num)
                    ))

        self.append_link(self.Link(
            uri_template.format(endpoint_uri=endpoint, page=page, per_page=per_page),
            'self',
            prompt=str(page)
        ))

        if page < number_of_pages:
            for trail_page in range(1, trailing + 1):
                page_num = page + trail_page
                if page_num < number_of_pages:
                    self.append_link(self.Link(
                        uri_template.format(endpoint_uri=endpoint, page=page_num, per_page=per_page),
                        'more',
                        prompt=str(page_num)
                    ))

            if page + leading < number_of_pages:
                self.append_link(self.Link(
                    '',
                    'skip',
                    prompt='…'
                ))

            if page < number_of_pages:
                self.append_link(self.Link(
                    uri_template.format(endpoint_uri=endpoint, page=page + 1, per_page=per_page),
                    'next',
                    prompt='Next'
                ))

                self.append_link(self.Link(
                    uri_template.format(endpoint_uri=endpoint, page=number_of_pages, per_page=per_page),
                    'last',
                    prompt='Last'
                ))
<commit_msg>* small fixes to please PyCharm's linter<commit_after>__author__ = 'Ian S. Evans'

import json
from collections import UserDict
from math import ceil

MIMETYPE = 'application/vnd.collection+json'

class CollectionPlusJSON(UserDict):

    mimetype = MIMETYPE

    class BaseCollectionItem(object):

        def __init__(self, *args, **kwargs):
            pass

        def __str__(self):
            return json.dumps(self.__dict__)

        def __repr__(self):
            return self.__str__()

    class Error(BaseCollectionItem):

        def __init__(self, title=None, code=None, message=None):
            if title is not None:
                self.title = title
            if code is not None:
                self.code = code
            if message is not None:
                self.message = message
            try:
                # Python 3
                super().__init__()
            except ValueError:
                # Python 2
                super(CollectionPlusJSON.Error, self).__init__()

    class Item(UserDict, BaseCollectionItem):

        def __init__(self, uri=None, links=None, **kwargs):
            if uri is not None:
                self.href = uri
            if links is not None:
                self.links = links
            try:
                # Python 3
                super().__init__(**kwargs)
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Item, self).__init__(**kwargs)

        def __str__(self):
            item_dict = dict(self.__dict__)
            if self.data:
                data_list = []
                for k, v in self.data.items():
                    to_append = {"name": k}
                    if not isinstance(v, str):
                        to_append["value"] = v[0]
                        to_append["prompt"] = v[1]
                    else:
                        to_append['value'] = v
                    data_list.append(to_append)
                item_dict['data'] = data_list
            return json.dumps(item_dict)

    class Link(BaseCollectionItem):

        def __init__(self, uri, rel, name=None, render=None, prompt=None):
            self.uri = uri
            self.rel = rel
            if name is not None:
                self.name = name
            if render is not None:
                self.render = render
            if prompt is not None:
                self.prompt = prompt
            try:
                # Python 3
                super().__init__()
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Link, self).__init__()

    class Query(Item):

        def __init__(self, uri, rel=None, prompt=None, **kwargs):
            self.uri = uri
            if rel is not None:
                self.rel = rel
            if prompt is not None:
                self.prompt = prompt
            try:
                # Python 3
                super().__init__(**kwargs)
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Query, self).__init__(**kwargs)

    class Template(Item):

        def __init__(self, **kwargs):
            try:
                # Python 3
                super().__init__(**kwargs)
            except TypeError:
                # Python 2
                super(CollectionPlusJSON.Template, self).__init__(**kwargs)

    def __init__(self, version=1.0, href='/api/', **kwargs):
        collection = {
            'version': str(version),
            'href': href,
        }
        # Using this method to invoke custom __setitem__ to verify types
        for k, v in kwargs.items():
            self[k] = v
        try:
            # Python 3
            super().__init__(**collection)
        except TypeError:
            # Python 2
            super(CollectionPlusJSON, self).__init__(**collection)

    def __str__(self):
        return json.dumps({'collection': self.data})

    def __repr__(self):
        return self.__str__()

    def __setitem__(self, key, value):
        type_map = {
            'links': list,
            'items': list,
            'queries': list,
            'template': self.Template,
            'error': self.Error
        }
        if key in type_map.keys():
            if not isinstance(value, type_map[key]):
                raise TypeError(
                    '{key} must be an instance of {classname}'.format(key=key, classname=str(type_map[key]))
                )
        self.data[key] = value

    def append_item(self, item):
        """
        Append an item to this collection's 'items' property.
        :param item: The Item to append.
        :return:
        """
        if not isinstance(item, self.Item):
            raise TypeError('item must be an instance of {classname}'.format(classname=str(self.Item)))
        else:
            if self.data.get('items') is not None:
                self.data.get('items').append(item)
            else:
                self.data['items'] = [item]

    def append_link(self, link):
        """
        Append a link to this collection's 'links' property.
        :param link: The Link to append.
        :return:
        """
        if not isinstance(link, self.Link):
            raise TypeError('item must be an instance of {classname}'.format(classname=str(self.Link)))
        else:
            if self.data.get('links') is not None:
                self.data.get('links').append(link)
            else:
                self.data['links'] = [link]

    def append_query(self, query):
        """
        Append a query to this collection's 'queries' property.
        :param query: The Query to append.
        :return:
        """
        if not isinstance(query, self.Query):
            raise TypeError('item must be an instance of {classname}'.format(classname=str(self.Query)))
        else:
            if self.data.get('queries') is not None:
                self.data.get('queries').append(query)
            else:
                self.data['queries'] = [query]

    def paginate(self, endpoint='', uri_template='{endpoint_uri}?page={page}&per_page={per_page}', page=1, per_page=5,
                 leading=2, trailing=2):
        """
        Paginate this collection, automatically trimming items and adding appropriate links for navigation.
        :param endpoint: The URI for this resource.
        :param uri_template: A string providing a template for paginated URI structure. May include the following keys:
        "{endpoint_uri}" - This will evaluate to the value of the 'endpoint' param.
        "{page}" - The page number will be inserted here.
        "{per_page}" - The number of items to display per page will be inserted here.
        :param page: The desired page for this representation.
        :param per_page: The number of items per page for this representation.
        :param leading: The number of leading pages before "this" page to add to this collection's "links".
        :param trailing: The number of trailing pages after "this" page to add to this collection's "links".
        :return:
        """
        if (type(page) is not int) or (type(per_page) is not int):

            try:
                page = abs(int(page))
            except (ValueError, TypeError):
                page = 1

            try:
                per_page = abs(int(per_page))
            except (ValueError, TypeError):
                per_page = 5

            if not page:
                page = 1
            if not per_page:
                per_page = 5

        number_of_pages = int(ceil(len(self.data.get('items')) / per_page))

        if page > number_of_pages:
            page = number_of_pages

        page_index_begin = ((page * per_page) - per_page)
        page_index_end = (page * per_page)

        self.data['items'] = self.data.get('items')[page_index_begin:page_index_end]

        if page > 1:
            self.append_link(self.Link(
                uri_template.format(endpoint_uri=endpoint, page=1, per_page=per_page),
                'first',
                prompt='First'
            ))

            self.append_link(self.Link(
                uri_template.format(endpoint_uri=endpoint, page=(page - 1), per_page=per_page),
                'prev',
                prompt='Previous'
            ))

            if page - leading > 0:
                self.append_link(self.Link(
                    '',
                    'skip',
                    prompt='…'
                ))

            for lead_page in range(leading):
                page_num = page - lead_page
                if page_num > 0:
                    self.append_link(self.Link(
                        uri_template.format(endpoint_uri=endpoint, page=page_num, per_page=per_page),
                        'more',
                        prompt=str(page_num)
                    ))

        self.append_link(self.Link(
            uri_template.format(endpoint_uri=endpoint, page=page, per_page=per_page),
            'self',
            prompt=str(page)
        ))

        if page < number_of_pages:
            for trail_page in range(1, trailing + 1):
                page_num = page + trail_page
                if page_num < number_of_pages:
                    self.append_link(self.Link(
                        uri_template.format(endpoint_uri=endpoint, page=page_num, per_page=per_page),
                        'more',
                        prompt=str(page_num)
                    ))

            if page + leading < number_of_pages:
                self.append_link(self.Link(
                    '',
                    'skip',
                    prompt='…'
                ))

            if page < number_of_pages:
                self.append_link(self.Link(
                    uri_template.format(endpoint_uri=endpoint, page=page + 1, per_page=per_page),
                    'next',
                    prompt='Next'
                ))

                self.append_link(self.Link(
                    uri_template.format(endpoint_uri=endpoint, page=number_of_pages, per_page=per_page),
                    'last',
                    prompt='Last'
                ))
<|endoftext|>


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

"""
Patch of the original code:

- Bugfix: pass transactional when using the datastore for the pickled payload
- Bugfix: delete the entity if the task cannot be added
- New: defer_async(...)
- New: defer_multi_async(tasks) / defer_multi / defer_task
- New: Allow huge payloads using the blobstore
- New: Pass in `_urlsuffix` e.g. ('foo', 'bar') for easier readable weblogs.
         Task url will be set to "/_ah/queue/deferred/foo/bar"
- New: Pass in anything as `_name` and it will get sha256 hashed if its not a
        valid name already

"""

from google.appengine.ext import deferred as old_deferred

from collections import namedtuple, defaultdict
import hashlib

from google.appengine.api import taskqueue
from google.appengine.ext import ndb
from google.appengine.ext import blobstore
from google.appengine.api import files

_TASKQUEUE_HEADERS = {"Content-Type": "application/octet-stream"}
_DEFAULT_URL = "/_ah/queue/deferred"
_DEFAULT_QUEUE = "default"
_TRANSACTIONAL_DEFAULT = 'auto'
SMALL_PAYLOAD = 100000
LARGE_PAYLOAD = 1000000

PermanentTaskFailure = old_deferred.PermanentTaskFailure
run = old_deferred.run

class _DeferredTaskEntity(ndb.Model):
    """Datastore representation of a deferred task.

    This is used in cases when the deferred task is too big to be included as
    payload with the task queue entry.
    """
    _use_cache = False

    large = ndb.BlobProperty(indexed=False)
    huge = ndb.BlobKeyProperty(indexed=False)

    @classmethod
    def _get_kind(cls):
        return '_Deferred2_Payload'

    @property
    def payload(self):
        if self.large:
            return self.large
        elif self.huge:
            return blobstore.BlobReader(self.huge).read()

    @payload.setter
    def payload(self, value):
        if len(value) < LARGE_PAYLOAD:
            self.large = value
            return

        filename = files.blobstore.create(mime_type='application/octet-stream')
        with files.open(filename, 'a') as f:
            f.write(value)
        files.finalize(filename)

        self.huge = files.blobstore.get_blob_key(filename)

    @ndb.tasklet
    def delete_async(self):
        if self.huge:
            blobstore.BlobInfo(self.huge).delete()

        yield self.key.delete_async()

    def delete(self):
        return self.delete_async().get_result()

def run_from_datastore(key):
    """Retrieves a task from the datastore and executes it.

    Args:
        key: The datastore key of a _DeferredTaskEntity storing the task.
    Returns:
        The return value of the function invocation.
    """
    entity = key.get()
    if not entity:
        raise PermanentTaskFailure()
    try:
        run(entity.payload)
        entity.delete()
    except PermanentTaskFailure:
        entity.delete()
        raise

serialize = old_deferred.serialize

DeferredTask = namedtuple(
    'DeferredTask', ['queue_name', 'transactional', 'task_def'])  \
    # type: Tuple[str, bool, dict]

def convert_task_to_api_task(task_def):
    # type: (dict) -> taskqueue.Task
    return taskqueue.Task(**task_def)

def task(obj, *args, **kwargs):
    # type: (...) -> DeferredTask
    taskargs = {x[1:]: kwargs.pop(x)
                for x in ("_countdown", "_eta", "_name", "_target",
                          "_url", "_retry_options") if x in kwargs}

    taskargs.setdefault('url', _DEFAULT_URL)
    taskargs["headers"] = dict(_TASKQUEUE_HEADERS)
    taskargs["headers"].update(kwargs.pop("_headers", {}))

    urlsuffix = kwargs.pop("_urlsuffix", None)
    if urlsuffix:
        if not isinstance(urlsuffix, basestring):
            urlsuffix = "/".join(map(str, urlsuffix))
        taskargs["url"] += "/{}".format(urlsuffix)

    if 'name' in taskargs:
        name = taskargs['name']
        if not isinstance(name, basestring):
            name = str(name)

        if not taskqueue.taskqueue._TASK_NAME_RE.match(name):
            name = hashlib.sha256(name).hexdigest()

        taskargs['name'] = name

    transactional = kwargs.pop("_transactional", _TRANSACTIONAL_DEFAULT)
    if transactional == 'auto':
        transactional = False if 'name' in taskargs else ndb.in_transaction()
    queue = kwargs.pop("_queue", _DEFAULT_QUEUE)

    taskargs['payload'] = serialize(obj, *args, **kwargs)

    return DeferredTask(queue_name=queue,
                        transactional=transactional,
                        task_def=taskargs)

def defer_async(obj, *args, **kwargs):
    # type: (...) -> Future[taskqueue.Task]
    """Defers a callable for execution later.

    The default deferred URL of /_ah/queue/deferred will be used unless an
    alternate URL is explicitly specified. If you want to use the default URL
    for a queue, specify _url=None. If you specify a different URL, you will
    need to install the handler on that URL (see the module docstring for
    details).

    Args:
        obj: The callable to execute. See module docstring for restrictions.
        args: Positional arguments to call the callable with.
        kwargs: _countdown, _eta, _headers, _name, _target, _transactional,
                _url, _retry_options, _queue: Passed through to the task queue
                - see the task queue documentation for details.
                Any other keyword arguments are passed through to the callable.
    Returns:
        A taskqueue.Task object which represents an enqueued callable.
    """

    deferred = task(obj, *args, **kwargs)
    return enqueue_async(deferred)

def defer(obj, *args, **kwargs):
    # type: (...) -> taskqueue.Task
    return defer_async(obj, *args, **kwargs).get_result()

@ndb.tasklet
def enqueue_async(task):
    # type: (DeferredTask) -> Future[taskqueue.Task]
    (queue, transactional, task_def) = task

    if len(task_def['payload']) < SMALL_PAYLOAD:
        task = convert_task_to_api_task(task_def)
        rv = yield task.add_async(queue, transactional=transactional)
        raise ndb.Return(rv)
    else:
        entity = _DeferredTaskEntity()
        entity.payload = task_def['payload']
        yield entity.put_async()

        try:
            new_def = task_def.copy()
            new_def['payload'] = serialize(run_from_datastore, entity.key)
            task = convert_task_to_api_task(new_def)
            rv = yield task.add_async(queue, transactional=transactional)
        except Exception as e:
            if not ndb.in_transaction():
                yield entity.delete_async()

            raise e

        raise ndb.Return(rv)

@ndb.tasklet
def handle_big_payloads(task):
    # type: (DeferredTask) -> DeferredTask
    task_def = task.task_def

    if len(task_def['payload']) < SMALL_PAYLOAD:
        raise ndb.Return(task)
    else:
        if not ndb.in_transaction():
            raise taskqueue.BadTransactionState(
                'Handling BIG payloads touches the db and thus requires '
                'a transaction.')

        entity = _DeferredTaskEntity()
        entity.payload = task_def['payload']
        key = yield entity.put_async()

        new_def = task_def.copy()
        new_def['payload'] = serialize(run_from_datastore, key)
        new_task = task._replace(task_def=new_def)
        # new_task = DeferredTask(queue, transactional, new_def)
        raise ndb.Return(new_task)

@ndb.tasklet
def final_transformation(task):
    (queue, transactional, task_def) = task
    return (queue, transactional, taskqueue.Task(**task_def))

class Batcher(object):
    def __init__(self, tasks):
        self._tasks = defaultdict(list)
        for t in tasks:
            self.add(t)

    def add(self, task):
        (queue, transactional, task) = task
        self._tasks[(queue, transactional)].append(task)

    def run_async(self):
        return [
            queue_multiple_tasks(queue, transactional, tasks)
            for (queue, transactional), tasks in self._tasks.iteritems()
        ]

@ndb.tasklet
def defer_multi_async(*tasks, **kwargs):
    # type: (List[DeferredTask]) -> Future[List[taskqueue.Task]]

    transformers = kwargs.pop(
        'transformers', [handle_big_payloads])  \
        # type: List[Callable[[DeferredTask], Optional[DeferredTask]]]

    transformers = transformers + [final_transformation]
    for fn in transformers:
        tasks = yield map(fn, tasks)
        tasks = filter(None, tasks)

    yield Batcher(tasks).run_async()

    raise ndb.Return(task for (_, _, task) in tasks)

def defer_multi(*tasks):
    # type: (List[DeferredTask]) -> List[taskqueue.Task]
    return defer_multi_async(*tasks).get_result()

@ndb.tasklet
def queue_multiple_tasks(queue, transactional, tasks):
    # Wrapper b/c Queue().add_async returns a UserRPC but a MultiFuture only
    # accepts Futures as dependents.
    raise ndb.Return(
        (yield taskqueue.Queue(queue).add_async(tasks,
                                                transactional=transactional))
    )

def one_shot_async(*tasks):
    return defer_async(defer_multi, *tasks)

def one_shot(*tasks):
    return one_shot_async(*tasks).get_result()
<commit_msg>Pass kwargs from defer_multi to its async pendant<commit_after>#!/usr/bin/env python
#
# Copyright 2007 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Patch of the original code:

- Bugfix: pass transactional when using the datastore for the pickled payload
- Bugfix: delete the entity if the task cannot be added
- New: defer_async(...)
- New: defer_multi_async(tasks) / defer_multi / defer_task
- New: Allow huge payloads using the blobstore
- New: Pass in `_urlsuffix` e.g. ('foo', 'bar') for easier readable weblogs.
         Task url will be set to "/_ah/queue/deferred/foo/bar"
- New: Pass in anything as `_name` and it will get sha256 hashed if its not a
        valid name already

"""

from google.appengine.ext import deferred as old_deferred

from collections import namedtuple, defaultdict
import hashlib

from google.appengine.api import taskqueue
from google.appengine.ext import ndb
from google.appengine.ext import blobstore
from google.appengine.api import files

_TASKQUEUE_HEADERS = {"Content-Type": "application/octet-stream"}
_DEFAULT_URL = "/_ah/queue/deferred"
_DEFAULT_QUEUE = "default"
_TRANSACTIONAL_DEFAULT = 'auto'
SMALL_PAYLOAD = 100000
LARGE_PAYLOAD = 1000000

PermanentTaskFailure = old_deferred.PermanentTaskFailure
run = old_deferred.run

class _DeferredTaskEntity(ndb.Model):
    """Datastore representation of a deferred task.

    This is used in cases when the deferred task is too big to be included as
    payload with the task queue entry.
    """
    _use_cache = False

    large = ndb.BlobProperty(indexed=False)
    huge = ndb.BlobKeyProperty(indexed=False)

    @classmethod
    def _get_kind(cls):
        return '_Deferred2_Payload'

    @property
    def payload(self):
        if self.large:
            return self.large
        elif self.huge:
            return blobstore.BlobReader(self.huge).read()

    @payload.setter
    def payload(self, value):
        if len(value) < LARGE_PAYLOAD:
            self.large = value
            return

        filename = files.blobstore.create(mime_type='application/octet-stream')
        with files.open(filename, 'a') as f:
            f.write(value)
        files.finalize(filename)

        self.huge = files.blobstore.get_blob_key(filename)

    @ndb.tasklet
    def delete_async(self):
        if self.huge:
            blobstore.BlobInfo(self.huge).delete()

        yield self.key.delete_async()

    def delete(self):
        return self.delete_async().get_result()

def run_from_datastore(key):
    """Retrieves a task from the datastore and executes it.

    Args:
        key: The datastore key of a _DeferredTaskEntity storing the task.
    Returns:
        The return value of the function invocation.
    """
    entity = key.get()
    if not entity:
        raise PermanentTaskFailure()
    try:
        run(entity.payload)
        entity.delete()
    except PermanentTaskFailure:
        entity.delete()
        raise

serialize = old_deferred.serialize

DeferredTask = namedtuple(
    'DeferredTask', ['queue_name', 'transactional', 'task_def'])  \
    # type: Tuple[str, bool, dict]

def convert_task_to_api_task(task_def):
    # type: (dict) -> taskqueue.Task
    return taskqueue.Task(**task_def)

def task(obj, *args, **kwargs):
    # type: (...) -> DeferredTask
    taskargs = {x[1:]: kwargs.pop(x)
                for x in ("_countdown", "_eta", "_name", "_target",
                          "_url", "_retry_options") if x in kwargs}

    taskargs.setdefault('url', _DEFAULT_URL)
    taskargs["headers"] = dict(_TASKQUEUE_HEADERS)
    taskargs["headers"].update(kwargs.pop("_headers", {}))

    urlsuffix = kwargs.pop("_urlsuffix", None)
    if urlsuffix:
        if not isinstance(urlsuffix, basestring):
            urlsuffix = "/".join(map(str, urlsuffix))
        taskargs["url"] += "/{}".format(urlsuffix)

    if 'name' in taskargs:
        name = taskargs['name']
        if not isinstance(name, basestring):
            name = str(name)

        if not taskqueue.taskqueue._TASK_NAME_RE.match(name):
            name = hashlib.sha256(name).hexdigest()

        taskargs['name'] = name

    transactional = kwargs.pop("_transactional", _TRANSACTIONAL_DEFAULT)
    if transactional == 'auto':
        transactional = False if 'name' in taskargs else ndb.in_transaction()
    queue = kwargs.pop("_queue", _DEFAULT_QUEUE)

    taskargs['payload'] = serialize(obj, *args, **kwargs)

    return DeferredTask(queue_name=queue,
                        transactional=transactional,
                        task_def=taskargs)

def defer_async(obj, *args, **kwargs):
    # type: (...) -> Future[taskqueue.Task]
    """Defers a callable for execution later.

    The default deferred URL of /_ah/queue/deferred will be used unless an
    alternate URL is explicitly specified. If you want to use the default URL
    for a queue, specify _url=None. If you specify a different URL, you will
    need to install the handler on that URL (see the module docstring for
    details).

    Args:
        obj: The callable to execute. See module docstring for restrictions.
        args: Positional arguments to call the callable with.
        kwargs: _countdown, _eta, _headers, _name, _target, _transactional,
                _url, _retry_options, _queue: Passed through to the task queue
                - see the task queue documentation for details.
                Any other keyword arguments are passed through to the callable.
    Returns:
        A taskqueue.Task object which represents an enqueued callable.
    """

    deferred = task(obj, *args, **kwargs)
    return enqueue_async(deferred)

def defer(obj, *args, **kwargs):
    # type: (...) -> taskqueue.Task
    return defer_async(obj, *args, **kwargs).get_result()

@ndb.tasklet
def enqueue_async(task):
    # type: (DeferredTask) -> Future[taskqueue.Task]
    (queue, transactional, task_def) = task

    if len(task_def['payload']) < SMALL_PAYLOAD:
        task = convert_task_to_api_task(task_def)
        rv = yield task.add_async(queue, transactional=transactional)
        raise ndb.Return(rv)
    else:
        entity = _DeferredTaskEntity()
        entity.payload = task_def['payload']
        yield entity.put_async()

        try:
            new_def = task_def.copy()
            new_def['payload'] = serialize(run_from_datastore, entity.key)
            task = convert_task_to_api_task(new_def)
            rv = yield task.add_async(queue, transactional=transactional)
        except Exception as e:
            if not ndb.in_transaction():
                yield entity.delete_async()

            raise e

        raise ndb.Return(rv)

@ndb.tasklet
def handle_big_payloads(task):
    # type: (DeferredTask) -> DeferredTask
    task_def = task.task_def

    if len(task_def['payload']) < SMALL_PAYLOAD:
        raise ndb.Return(task)
    else:
        if not ndb.in_transaction():
            raise taskqueue.BadTransactionState(
                'Handling BIG payloads touches the db and thus requires '
                'a transaction.')

        entity = _DeferredTaskEntity()
        entity.payload = task_def['payload']
        key = yield entity.put_async()

        new_def = task_def.copy()
        new_def['payload'] = serialize(run_from_datastore, key)
        new_task = task._replace(task_def=new_def)
        # new_task = DeferredTask(queue, transactional, new_def)
        raise ndb.Return(new_task)

@ndb.tasklet
def final_transformation(task):
    (queue, transactional, task_def) = task
    return (queue, transactional, taskqueue.Task(**task_def))

class Batcher(object):
    def __init__(self, tasks):
        self._tasks = defaultdict(list)
        for t in tasks:
            self.add(t)

    def add(self, task):
        (queue, transactional, task) = task
        self._tasks[(queue, transactional)].append(task)

    def run_async(self):
        return [
            queue_multiple_tasks(queue, transactional, tasks)
            for (queue, transactional), tasks in self._tasks.iteritems()
        ]

@ndb.tasklet
def defer_multi_async(*tasks, **kwargs):
    # type: (List[DeferredTask]) -> Future[List[taskqueue.Task]]

    transformers = kwargs.pop(
        'transformers', [handle_big_payloads])  \
        # type: List[Callable[[DeferredTask], Optional[DeferredTask]]]

    transformers = transformers + [final_transformation]
    for fn in transformers:
        tasks = yield map(fn, tasks)
        tasks = filter(None, tasks)

    yield Batcher(tasks).run_async()

    raise ndb.Return(task for (_, _, task) in tasks)

def defer_multi(*tasks, **kwargs):
    # type: (List[DeferredTask]) -> List[taskqueue.Task]
    return defer_multi_async(*tasks, **kwargs).get_result()

@ndb.tasklet
def queue_multiple_tasks(queue, transactional, tasks):
    # Wrapper b/c Queue().add_async returns a UserRPC but a MultiFuture only
    # accepts Futures as dependents.
    raise ndb.Return(
        (yield taskqueue.Queue(queue).add_async(tasks,
                                                transactional=transactional))
    )

def one_shot_async(*tasks):
    return defer_async(defer_multi, *tasks)

def one_shot(*tasks):
    return one_shot_async(*tasks).get_result()
<|endoftext|>