<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2010-Today OpenERP SA (<http://www.openerp.com>)
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>
#
##############################################################################

import base64
import re
from openerp import tools

from openerp.osv import osv
from openerp.osv import fields
from openerp.tools.safe_eval import safe_eval as eval
from openerp.tools.translate import _

# main mako-like expression pattern
EXPRESSION_PATTERN = re.compile('(\$\{.+?\})')

class mail_compose_message(osv.TransientModel):
    """ Generic message composition wizard. You may inherit from this wizard
        at model and view levels to provide specific features.

        The behavior of the wizard depends on the composition_mode field:
        - 'reply': reply to a previous message. The wizard is pre-populated
            via ``get_message_data``.
        - 'comment': new post on a record. The wizard is pre-populated via
            ``get_record_data``
        - 'mass_mail': wizard in mass mailing mode where the mail details can
            contain template placeholders that will be merged with actual data
            before being sent to each recipient.
    """
    _name = 'mail.compose.message'
    _inherit = 'mail.message'
    _description = 'Email composition wizard'
    _log_access = True

    def default_get(self, cr, uid, fields, context=None):
        """ Handle composition mode. Some details about context keys:
            - comment: default mode, model and ID of a record the user comments
                - default_model or active_model
                - default_res_id or active_id
            - reply: active_id of a message the user replies to
                - default_parent_id or message_id or active_id: ID of the
                    mail.message we reply to
                - message.res_model or default_model
                - message.res_id or default_res_id
            - mass_mail: model and IDs of records the user mass-mails
                - active_ids: record IDs
                - default_model or active_model
        """
        if context is None:
            context = {}
        result = super(mail_compose_message, self).default_get(cr, uid, fields, context=context)

        # get some important values from context
        composition_mode = context.get('default_composition_mode', context.get('mail.compose.message.mode'))
        model = context.get('default_model', context.get('active_model'))
        res_id = context.get('default_res_id', context.get('active_id'))
        message_id = context.get('default_parent_id', context.get('message_id', context.get('active_id')))
        active_ids = context.get('active_ids')

        # get default values according to the composition mode
        if composition_mode == 'reply':
            vals = self.get_message_data(cr, uid, message_id, context=context)
        elif composition_mode == 'comment' and model and res_id:
            vals = self.get_record_data(cr, uid, model, res_id, context=context)
        elif composition_mode == 'mass_mail' and model and active_ids:
            vals = {'model': model, 'res_id': res_id}
        else:
            vals = {'model': model, 'res_id': res_id}
        if composition_mode:
            vals['composition_mode'] = composition_mode

        for field in vals:
            if field in fields:
                result[field] = vals[field]

        # TDE HACK: as mailboxes used default_model='res.users' and default_res_id=uid
        # (because of lack of an accessible pid), creating a message on its own
        # profile may crash (res_users does not allow writing on it)
        # Posting on its own profile works (res_users redirect to res_partner)
        # but when creating the mail.message to create the mail.compose.message
        # access rights issues may rise
        # We therefore directly change the model and res_id
        if result.get('model') == 'res.users' and result.get('res_id') == uid:
            result['model'] = 'res.partner'
            result['res_id'] = self.pool.get('res.users').browse(cr, uid, uid).partner_id.id
        return result

    def _get_composition_mode_selection(self, cr, uid, context=None):
        return [('comment', 'Comment a document'), ('reply', 'Reply to a message'), ('mass_mail', 'Mass mailing')]

    _columns = {
        'composition_mode': fields.selection(
            lambda s, *a, **k: s._get_composition_mode_selection(*a, **k),
            string='Composition mode'),
        'partner_ids': fields.many2many('res.partner',
            'mail_compose_message_res_partner_rel',
            'wizard_id', 'partner_id', 'Additional contacts'),
        'post': fields.boolean('Post a copy in the document',
            help='Post a copy of the message on the document communication history.'),
        'notify': fields.boolean('Notify followers',
            help='Notify followers of the document'),
        'same_thread': fields.boolean('Replies in the document',
            help='Replies to the messages will go into the selected document.'),
        'attachment_ids': fields.many2many('ir.attachment',
            'mail_compose_message_ir_attachments_rel',
            'wizard_id', 'attachment_id', 'Attachments'),
        'filter_id': fields.many2one('ir.filters', 'Filters'),
    }

    _defaults = {
        'composition_mode': 'comment',
        'email_from': lambda self, cr, uid, ctx={}: self.pool.get('mail.mail')._get_default_from(cr, uid, context=ctx),
        'body': lambda self, cr, uid, ctx={}: '',
        'subject': lambda self, cr, uid, ctx={}: False,
        'partner_ids': lambda self, cr, uid, ctx={}: [],
        'notify': lambda self, cr, uid, ctx={}: False,
        'post': lambda self, cr, uid, ctx={}: True,
        'same_thread': lambda self, cr, uid, ctx={}: True,
    }

    def _notify(self, cr, uid, newid, context=None):
        """ Override specific notify method of mail.message, because we do
            not want that feature in the wizard. """
        return

    def get_record_data(self, cr, uid, model, res_id, context=None):
        """ Returns a defaults-like dict with initial values for the composition
            wizard when sending an email related to the document record
            identified by ``model`` and ``res_id``.

            :param str model: model name of the document record this mail is
                related to.
            :param int res_id: id of the document record this mail is related to
        """
        doc_name_get = self.pool.get(model).name_get(cr, uid, [res_id], context=context)
        if doc_name_get:
            record_name = doc_name_get[0][1]
        else:
            record_name = False
        return {'model': model, 'res_id': res_id, 'record_name': record_name}

    def get_message_data(self, cr, uid, message_id, context=None):
        """ Returns a defaults-like dict with initial values for the composition
            wizard when replying to the given message (e.g. including the quote
            of the initial message, and the correct recipients).

            :param int message_id: id of the mail.message to which the user
                is replying.
        """
        if not message_id:
            return {}
        if context is None:
            context = {}
        message_data = self.pool.get('mail.message').browse(cr, uid, message_id, context=context)

        # create subject
        re_prefix = _('Re:')
        reply_subject = tools.ustr(message_data.subject or '')
        if not (reply_subject.startswith('Re:') or reply_subject.startswith(re_prefix)) and message_data.subject:
            reply_subject = "%s %s" % (re_prefix, reply_subject)
        # get partner_ids from original message
        partner_ids = [partner.id for partner in message_data.partner_ids] if message_data.partner_ids else []
        partner_ids += context.get('default_partner_ids', [])

        # update the result
        result = {
            'record_name': message_data.record_name,
            'model': message_data.model,
            'res_id': message_data.res_id,
            'parent_id': message_data.id,
            'subject': reply_subject,
            'partner_ids': partner_ids,
        }
        return result

    #------------------------------------------------------
    # Wizard validation and send
    #------------------------------------------------------

    def send_mail(self, cr, uid, ids, context=None):
        """ Process the wizard content and proceed with sending the related
            email(s), rendering any template patterns on the fly if needed. """
        if context is None:
            context = {}
        active_ids = context.get('active_ids')

        for wizard in self.browse(cr, uid, ids, context=context):
            mass_mail_mode = wizard.composition_mode == 'mass_mail'
            active_model_pool = self.pool.get(wizard.model if wizard.model else 'mail.thread')

            # wizard works in batch mode: [res_id] or active_ids
            res_ids = active_ids if mass_mail_mode and wizard.model and active_ids else [wizard.res_id]
            for res_id in res_ids:
                # default values, according to the wizard options
                post_values = {
                    'subject': wizard.subject,
                    'body': wizard.body,
                    'parent_id': wizard.parent_id and wizard.parent_id.id,
                    'partner_ids': [(4, partner.id) for partner in wizard.partner_ids],
                    'attachments': [(attach.datas_fname or attach.name, base64.b64decode(attach.datas)) for attach in wizard.attachment_ids],
                }
                # mass mailing: render and override default values
                if mass_mail_mode and wizard.model:
                    email_dict = self.render_message(cr, uid, wizard, res_id, context=context)
                    new_partner_ids = email_dict.pop('partner_ids', [])
                    post_values['partner_ids'] += [(4, partner_id) for partner_id in new_partner_ids]
                    new_attachments = email_dict.pop('attachments', [])
                    post_values['attachments'] += new_attachments
                    post_values.update(email_dict)
                    # email_from: mass mailing only can specify another email_from
                    if email_dict.get('email_from'):
                        post_values['email_from'] = email_dict.pop('email_from')
                    # replies redirection: mass mailing only
                    if not wizard.same_thread:
                        post_values['reply_to'] = email_dict.pop('reply_to')
                # automatically subscribe recipients if asked to
                if context.get('mail_post_autofollow') and wizard.model and post_values.get('partner_ids'):
                    active_model_pool.message_subscribe(cr, uid, [res_id], [item[1] for item in post_values.get('partner_ids')], context=context)
                # post the message
                if mass_mail_mode and not wizard.post:
                    post_values['recipient_ids'] = post_values.pop('partner_ids')
                    self.pool.get('mail.mail').create(cr, uid, post_values, context=context)
                else:
                    subtype = 'mail.mt_comment'
                    if mass_mail_mode and not wizard.notify:
                        subtype = False
                    msg_id = active_model_pool.message_post(cr, uid, [res_id], type='comment', subtype=subtype, context=context, **post_values)
                    # mass_mailing, post without notify: notify specific partners
                    if mass_mail_mode and not wizard.notify and post_values['partner_ids']:
                        self.pool.get('mail.notification')._notify(cr, uid, msg_id, [item[1] for item in post_values['partner_ids']], context=context)

        return {'type': 'ir.actions.act_window_close'}

    def render_message(self, cr, uid, wizard, res_id, context=None):
        """ Generate an email from the template for given (wizard.model, res_id)
            pair. This method is meant to be inherited by email_template that
            will produce a more complete dictionary. """
        return {
            'subject': self.render_template(cr, uid, wizard.subject, wizard.model, res_id, context),
            'body': self.render_template(cr, uid, wizard.body, wizard.model, res_id, context),
            'email_from': self.render_template(cr, uid, wizard.email_from, wizard.model, res_id, context),
            'reply_to': self.render_template(cr, uid, wizard.reply_to, wizard.model, res_id, context),
        }

    def render_template(self, cr, uid, template, model, res_id, context=None):
        """ Render the given template text, replace mako-like expressions ``${expr}``
            with the result of evaluating these expressions with an evaluation context
            containing:

                * ``user``: browse_record of the current user
                * ``object``: browse_record of the document record this mail is
                              related to
                * ``context``: the context passed to the mail composition wizard

            :param str template: the template text to render
            :param str model: model name of the document record this mail is related to.
            :param int res_id: id of the document record this mail is related to.
        """
        if context is None:
            context = {}

        def merge(match):
            exp = str(match.group()[2:-1]).strip()
            result = eval(exp, {
                'user': self.pool.get('res.users').browse(cr, uid, uid, context=context),
                'object': self.pool.get(model).browse(cr, uid, res_id, context=context),
                'context': dict(context),  # copy context to prevent side-effects of eval
                })
            return result and tools.ustr(result) or ''
        return template and EXPRESSION_PATTERN.sub(merge, template)
<commit_msg>[FIX] mail_compose_message: remove some keys from the context, as they conflict with mail_mail values.<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2010-Today OpenERP SA (<http://www.openerp.com>)
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>
#
##############################################################################

import base64
import re
from openerp import tools

from openerp.osv import osv
from openerp.osv import fields
from openerp.tools.safe_eval import safe_eval as eval
from openerp.tools.translate import _

# main mako-like expression pattern
EXPRESSION_PATTERN = re.compile('(\$\{.+?\})')

class mail_compose_message(osv.TransientModel):
    """ Generic message composition wizard. You may inherit from this wizard
        at model and view levels to provide specific features.

        The behavior of the wizard depends on the composition_mode field:
        - 'reply': reply to a previous message. The wizard is pre-populated
            via ``get_message_data``.
        - 'comment': new post on a record. The wizard is pre-populated via
            ``get_record_data``
        - 'mass_mail': wizard in mass mailing mode where the mail details can
            contain template placeholders that will be merged with actual data
            before being sent to each recipient.
    """
    _name = 'mail.compose.message'
    _inherit = 'mail.message'
    _description = 'Email composition wizard'
    _log_access = True

    def default_get(self, cr, uid, fields, context=None):
        """ Handle composition mode. Some details about context keys:
            - comment: default mode, model and ID of a record the user comments
                - default_model or active_model
                - default_res_id or active_id
            - reply: active_id of a message the user replies to
                - default_parent_id or message_id or active_id: ID of the
                    mail.message we reply to
                - message.res_model or default_model
                - message.res_id or default_res_id
            - mass_mail: model and IDs of records the user mass-mails
                - active_ids: record IDs
                - default_model or active_model
        """
        if context is None:
            context = {}
        result = super(mail_compose_message, self).default_get(cr, uid, fields, context=context)

        # get some important values from context
        composition_mode = context.get('default_composition_mode', context.get('mail.compose.message.mode'))
        model = context.get('default_model', context.get('active_model'))
        res_id = context.get('default_res_id', context.get('active_id'))
        message_id = context.get('default_parent_id', context.get('message_id', context.get('active_id')))
        active_ids = context.get('active_ids')

        # get default values according to the composition mode
        if composition_mode == 'reply':
            vals = self.get_message_data(cr, uid, message_id, context=context)
        elif composition_mode == 'comment' and model and res_id:
            vals = self.get_record_data(cr, uid, model, res_id, context=context)
        elif composition_mode == 'mass_mail' and model and active_ids:
            vals = {'model': model, 'res_id': res_id}
        else:
            vals = {'model': model, 'res_id': res_id}
        if composition_mode:
            vals['composition_mode'] = composition_mode

        for field in vals:
            if field in fields:
                result[field] = vals[field]

        # TDE HACK: as mailboxes used default_model='res.users' and default_res_id=uid
        # (because of lack of an accessible pid), creating a message on its own
        # profile may crash (res_users does not allow writing on it)
        # Posting on its own profile works (res_users redirect to res_partner)
        # but when creating the mail.message to create the mail.compose.message
        # access rights issues may rise
        # We therefore directly change the model and res_id
        if result.get('model') == 'res.users' and result.get('res_id') == uid:
            result['model'] = 'res.partner'
            result['res_id'] = self.pool.get('res.users').browse(cr, uid, uid).partner_id.id
        return result

    def _get_composition_mode_selection(self, cr, uid, context=None):
        return [('comment', 'Comment a document'), ('reply', 'Reply to a message'), ('mass_mail', 'Mass mailing')]

    _columns = {
        'composition_mode': fields.selection(
            lambda s, *a, **k: s._get_composition_mode_selection(*a, **k),
            string='Composition mode'),
        'partner_ids': fields.many2many('res.partner',
            'mail_compose_message_res_partner_rel',
            'wizard_id', 'partner_id', 'Additional contacts'),
        'post': fields.boolean('Post a copy in the document',
            help='Post a copy of the message on the document communication history.'),
        'notify': fields.boolean('Notify followers',
            help='Notify followers of the document'),
        'same_thread': fields.boolean('Replies in the document',
            help='Replies to the messages will go into the selected document.'),
        'attachment_ids': fields.many2many('ir.attachment',
            'mail_compose_message_ir_attachments_rel',
            'wizard_id', 'attachment_id', 'Attachments'),
        'filter_id': fields.many2one('ir.filters', 'Filters'),
    }

    _defaults = {
        'composition_mode': 'comment',
        'email_from': lambda self, cr, uid, ctx={}: self.pool.get('mail.mail')._get_default_from(cr, uid, context=ctx),
        'body': lambda self, cr, uid, ctx={}: '',
        'subject': lambda self, cr, uid, ctx={}: False,
        'partner_ids': lambda self, cr, uid, ctx={}: [],
        'notify': lambda self, cr, uid, ctx={}: False,
        'post': lambda self, cr, uid, ctx={}: True,
        'same_thread': lambda self, cr, uid, ctx={}: True,
    }

    def _notify(self, cr, uid, newid, context=None):
        """ Override specific notify method of mail.message, because we do
            not want that feature in the wizard. """
        return

    def get_record_data(self, cr, uid, model, res_id, context=None):
        """ Returns a defaults-like dict with initial values for the composition
            wizard when sending an email related to the document record
            identified by ``model`` and ``res_id``.

            :param str model: model name of the document record this mail is
                related to.
            :param int res_id: id of the document record this mail is related to
        """
        doc_name_get = self.pool.get(model).name_get(cr, uid, [res_id], context=context)
        if doc_name_get:
            record_name = doc_name_get[0][1]
        else:
            record_name = False
        return {'model': model, 'res_id': res_id, 'record_name': record_name}

    def get_message_data(self, cr, uid, message_id, context=None):
        """ Returns a defaults-like dict with initial values for the composition
            wizard when replying to the given message (e.g. including the quote
            of the initial message, and the correct recipients).

            :param int message_id: id of the mail.message to which the user
                is replying.
        """
        if not message_id:
            return {}
        if context is None:
            context = {}
        message_data = self.pool.get('mail.message').browse(cr, uid, message_id, context=context)

        # create subject
        re_prefix = _('Re:')
        reply_subject = tools.ustr(message_data.subject or '')
        if not (reply_subject.startswith('Re:') or reply_subject.startswith(re_prefix)) and message_data.subject:
            reply_subject = "%s %s" % (re_prefix, reply_subject)
        # get partner_ids from original message
        partner_ids = [partner.id for partner in message_data.partner_ids] if message_data.partner_ids else []
        partner_ids += context.get('default_partner_ids', [])

        # update the result
        result = {
            'record_name': message_data.record_name,
            'model': message_data.model,
            'res_id': message_data.res_id,
            'parent_id': message_data.id,
            'subject': reply_subject,
            'partner_ids': partner_ids,
        }
        return result

    #------------------------------------------------------
    # Wizard validation and send
    #------------------------------------------------------

    def send_mail(self, cr, uid, ids, context=None):
        """ Process the wizard content and proceed with sending the related
            email(s), rendering any template patterns on the fly if needed. """
        if context is None:
            context = {}
        active_ids = context.get('active_ids')

        for wizard in self.browse(cr, uid, ids, context=context):
            mass_mail_mode = wizard.composition_mode == 'mass_mail'
            active_model_pool = self.pool.get(wizard.model if wizard.model else 'mail.thread')

            # wizard works in batch mode: [res_id] or active_ids
            res_ids = active_ids if mass_mail_mode and wizard.model and active_ids else [wizard.res_id]
            for res_id in res_ids:
                # default values, according to the wizard options
                post_values = {
                    'subject': wizard.subject,
                    'body': wizard.body,
                    'parent_id': wizard.parent_id and wizard.parent_id.id,
                    'partner_ids': [(4, partner.id) for partner in wizard.partner_ids],
                    'attachments': [(attach.datas_fname or attach.name, base64.b64decode(attach.datas)) for attach in wizard.attachment_ids],
                }
                # mass mailing: render and override default values
                if mass_mail_mode and wizard.model:
                    email_dict = self.render_message(cr, uid, wizard, res_id, context=context)
                    new_partner_ids = email_dict.pop('partner_ids', [])
                    post_values['partner_ids'] += [(4, partner_id) for partner_id in new_partner_ids]
                    new_attachments = email_dict.pop('attachments', [])
                    post_values['attachments'] += new_attachments
                    post_values.update(email_dict)
                    # email_from: mass mailing only can specify another email_from
                    if email_dict.get('email_from'):
                        post_values['email_from'] = email_dict.pop('email_from')
                    # replies redirection: mass mailing only
                    if not wizard.same_thread:
                        post_values['reply_to'] = email_dict.pop('reply_to')
                # automatically subscribe recipients if asked to
                if context.get('mail_post_autofollow') and wizard.model and post_values.get('partner_ids'):
                    active_model_pool.message_subscribe(cr, uid, [res_id], [item[1] for item in post_values.get('partner_ids')], context=context)
                # clean the context (hint: mass mailing sets some default values that
                # could be wrongly interpreted by mail_mail)
                context.pop('default_email_to', None)
                context.pop('default_partner_ids', None)
                # post the message
                if mass_mail_mode and not wizard.post:
                    post_values['recipient_ids'] = post_values.pop('partner_ids')
                    self.pool.get('mail.mail').create(cr, uid, post_values, context=context)
                else:
                    subtype = 'mail.mt_comment'
                    if mass_mail_mode and not wizard.notify:
                        subtype = False
                    msg_id = active_model_pool.message_post(cr, uid, [res_id], type='comment', subtype=subtype, context=context, **post_values)
                    # mass_mailing, post without notify: notify specific partners
                    if mass_mail_mode and not wizard.notify and post_values['partner_ids']:
                        self.pool.get('mail.notification')._notify(cr, uid, msg_id, [item[1] for item in post_values['partner_ids']], context=context)

        return {'type': 'ir.actions.act_window_close'}

    def render_message(self, cr, uid, wizard, res_id, context=None):
        """ Generate an email from the template for given (wizard.model, res_id)
            pair. This method is meant to be inherited by email_template that
            will produce a more complete dictionary. """
        return {
            'subject': self.render_template(cr, uid, wizard.subject, wizard.model, res_id, context),
            'body': self.render_template(cr, uid, wizard.body, wizard.model, res_id, context),
            'email_from': self.render_template(cr, uid, wizard.email_from, wizard.model, res_id, context),
            'reply_to': self.render_template(cr, uid, wizard.reply_to, wizard.model, res_id, context),
        }

    def render_template(self, cr, uid, template, model, res_id, context=None):
        """ Render the given template text, replace mako-like expressions ``${expr}``
            with the result of evaluating these expressions with an evaluation context
            containing:

                * ``user``: browse_record of the current user
                * ``object``: browse_record of the document record this mail is
                              related to
                * ``context``: the context passed to the mail composition wizard

            :param str template: the template text to render
            :param str model: model name of the document record this mail is related to.
            :param int res_id: id of the document record this mail is related to.
        """
        if context is None:
            context = {}

        def merge(match):
            exp = str(match.group()[2:-1]).strip()
            result = eval(exp, {
                'user': self.pool.get('res.users').browse(cr, uid, uid, context=context),
                'object': self.pool.get(model).browse(cr, uid, res_id, context=context),
                'context': dict(context),  # copy context to prevent side-effects of eval
                })
            return result and tools.ustr(result) or ''
        return template and EXPRESSION_PATTERN.sub(merge, template)
<|endoftext|>


<commit_before>import fabric.state
from fabric.context_managers import settings
from fabric.operations import prompt
from fabric.tasks import requires_parallel
from fabric.utils import abort

from .colors import green
from .state import blueprints
from .utils import info

class Task(object):

    @staticmethod
    def get_hosts_and_effective_roles(self, arg_hosts, arg_roles, arg_exclude_hosts, env=None):
        all_hosts, effective_roles = Task.get_hosts_and_effective_roles.original(self,
                                                                                arg_hosts=arg_hosts,
                                                                                arg_roles=arg_roles,
                                                                                arg_exclude_hosts=arg_exclude_hosts,
                                                                                env=env)
        # Prompt hosts if more than 1
        if fabric.state.env.prompt_hosts and len(all_hosts) > 1 and not requires_parallel(self):
            print("0. All")
            for i, host in enumerate(all_hosts, start=1):
                print("{i}. {host}".format(i=i, host=host))

            def validate_hosts(host_input):
                check = lambda s: s.isdigit() and int(s) in range(len(all_hosts))
                if not all(map(check, host_input.split(','))):
                    raise Exception('Invalid host {!r}'.format(host_input))
                else:
                    return host_input

            host_choice = prompt(green('Select host(s)'), default='0', validate=validate_hosts)
            indices = map(int, host_choice.split(','))
            if len(indices) > 1 or indices[0] > 0:
                all_hosts = [all_hosts[i-1] for i in indices if i > 0]

        return all_hosts, effective_roles

def help_task(blueprint_name=None):
    """
    Display blueprint help

    :param blueprint_name: Blueprint name
    """
    if not blueprint_name:
        abort('No blueprint provided, example: $ fab help:python')

    blueprint = blueprints.get(blueprint_name)
    if not blueprint:
        abort('Unknown blueprint "{}", using correct role?'.format(blueprint_name))

    help(blueprint['__module__'])

def init_task(blueprint_name=None):
    """
    Copy template structure from blueprint defaults.

    :param blueprint_name: Blueprint name
    """
    if not blueprint_name:
        abort('No blueprint provided, example: $ fab init:memcached')

    blueprint = blueprints.get(blueprint_name)
    if not blueprint:
        abort('Unknown blueprint "{}", using correct role?'.format(blueprint_name))

    module = blueprint['__module__']
    if hasattr(module, 'blueprint'):
        module.blueprint.inherit_templates()

    if module.__doc__:
        info('Help: {}', module.blueprint.name)
        print module.__doc__

def _execute(task, host, my_env, args, kwargs, jobs, queue, multiprocessing):
    """
    Patched version of fabric.tasks._execute.
    Wraps original with `effective_roles` as `roles` in env to apply definitions.
    """
    with settings(roles=my_env['effective_roles']):
        return _execute.original(task, host, my_env, args, kwargs, jobs, queue, multiprocessing)
<commit_msg>Actually validate all hosts<commit_after>import fabric.state
from fabric.context_managers import settings
from fabric.operations import prompt
from fabric.tasks import requires_parallel
from fabric.utils import abort

from .colors import green
from .state import blueprints
from .utils import info

class Task(object):

    @staticmethod
    def get_hosts_and_effective_roles(self, arg_hosts, arg_roles, arg_exclude_hosts, env=None):
        all_hosts, effective_roles = Task.get_hosts_and_effective_roles.original(self,
                                                                                arg_hosts=arg_hosts,
                                                                                arg_roles=arg_roles,
                                                                                arg_exclude_hosts=arg_exclude_hosts,
                                                                                env=env)
        # Prompt hosts if more than 1
        if fabric.state.env.prompt_hosts and len(all_hosts) > 1 and not requires_parallel(self):
            print("0. All")
            for i, host in enumerate(all_hosts, start=1):
                print("{i}. {host}".format(i=i, host=host))

            def validate_hosts(host_input):
                check = lambda s: s.isdigit() and int(s) in range(len(all_hosts) + 1)
                if not all(map(check, host_input.split(','))):
                    raise Exception('Invalid host {!r}'.format(host_input))
                else:
                    return host_input

            host_choice = prompt(green('Select host(s)'), default='0', validate=validate_hosts)
            indices = map(int, host_choice.split(','))
            if len(indices) > 1 or indices[0] > 0:
                all_hosts = [all_hosts[i-1] for i in indices if i > 0]

        return all_hosts, effective_roles

def help_task(blueprint_name=None):
    """
    Display blueprint help

    :param blueprint_name: Blueprint name
    """
    if not blueprint_name:
        abort('No blueprint provided, example: $ fab help:python')

    blueprint = blueprints.get(blueprint_name)
    if not blueprint:
        abort('Unknown blueprint "{}", using correct role?'.format(blueprint_name))

    help(blueprint['__module__'])

def init_task(blueprint_name=None):
    """
    Copy template structure from blueprint defaults.

    :param blueprint_name: Blueprint name
    """
    if not blueprint_name:
        abort('No blueprint provided, example: $ fab init:memcached')

    blueprint = blueprints.get(blueprint_name)
    if not blueprint:
        abort('Unknown blueprint "{}", using correct role?'.format(blueprint_name))

    module = blueprint['__module__']
    if hasattr(module, 'blueprint'):
        module.blueprint.inherit_templates()

    if module.__doc__:
        info('Help: {}', module.blueprint.name)
        print module.__doc__

def _execute(task, host, my_env, args, kwargs, jobs, queue, multiprocessing):
    """
    Patched version of fabric.tasks._execute.
    Wraps original with `effective_roles` as `roles` in env to apply definitions.
    """
    with settings(roles=my_env['effective_roles']):
        return _execute.original(task, host, my_env, args, kwargs, jobs, queue, multiprocessing)
<|endoftext|>


<commit_before>#  Copyright 2008 Nokia Siemens Networks Oyj
#
#  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.

import os
import types

from robot import utils
from robot.errors import DataError
from robot.common import BaseLibrary
from robot.output import SystemLogger

from handlers import PythonHandler, JavaHandler, DynamicHandler

def TestLibrary(name, args=None, syslog=None):
    if syslog is None:
        syslog = SystemLogger()
    libcode, source = utils.import_(name)
    args = utils.to_list(args)
    if isinstance(libcode, types.ModuleType):
        if args:
            raise DataError('Libraries implemented as modules do not take '
                            'arguments, got: %s' % str(args))
        return ModuleLibrary(libcode, source, name, args, syslog)
    if _has_method(libcode, ['get_keyword_names', 'getKeywordNames']):
        if _has_method(libcode, ['run_keyword', 'runKeyword']):
            return DynamicLibrary(libcode, source, name, args, syslog)
        else:
            return HybridLibrary(libcode, source, name, args, syslog)
    # Using type check and not isinstance for ClassType, because it does not
    # match Java classes whose type is javaclass.
    if type(libcode) is types.ClassType or isinstance(libcode, types.TypeType):
        return PythonLibrary(libcode, source, name, args, syslog)
    return JavaLibrary(libcode, source, name, args, syslog)

def _has_method(code, names):
    for name in names:
        if hasattr(code, name) and _is_method(getattr(code, name)):
            return True
    return False

def _is_method(code):
    return isinstance(code, types.MethodType) or \
            (utils.is_jython and 'reflectedfunction' in str(type(code)))

class _BaseTestLibrary(BaseLibrary):

    def __init__(self, libcode, source, name, args, syslog):
        if os.path.exists(name):
            name = os.path.splitext(os.path.basename(name))[0]
        self.source = source
        self.version = self._get_version(libcode)
        self.name = name
        self.orig_name = name # Stores original name also after copying
        self.args = args
        self._instance_cache = []
        if libcode is not None:
            self.doc = utils.get_doc(libcode)
            self.scope = self._get_scope(libcode)
            self._libcode = libcode
            self._libinst = self.get_instance()
            self.handlers = self._create_handlers(syslog)
            self._init_scope_handling(self.scope)

    def _get_version(self, code):
        try:
            return str(code.ROBOT_LIBRARY_VERSION)
        except AttributeError:
            try:
                return str(code.__version__)
            except AttributeError:
                return '<unknown>'

    def _init_scope_handling(self, scope):
        if scope == 'GLOBAL':
            return
        self._libinst = None
        self.start_suite = self._caching_start
        self.end_suite = self._restoring_end
        if scope == 'TESTCASE':
            self.start_test = self._caching_start
            self.end_test = self._restoring_end

    def copy(self, name):
        lib = _BaseTestLibrary(None, self.version, name, self.args, None)
        lib.orig_name = self.name
        lib.doc = self.doc
        lib.scope = self.scope
        lib._libcode = self._libcode
        lib._libinst = self._libinst
        lib.handlers = utils.NormalizedDict(ignore=['_'])
        for name, handler in self.handlers.items():
            lib.handlers[name] = handler.copy(lib)
        return lib

    def start_suite(self):
        pass

    def end_suite(self):
        pass

    def start_test(self):
        pass

    def end_test(self):
        pass

    def _caching_start(self):
        self._instance_cache.append(self._libinst)
        self._libinst = None

    def _restoring_end(self):
        self._libinst = self._instance_cache.pop()

    def _get_scope(self, libcode):
        try:
            scope = libcode.ROBOT_LIBRARY_SCOPE
            scope = utils.normalize(scope, ignore=['_']).upper()
        except:
            scope = 'TESTCASE'
        return scope in ['GLOBAL','TESTSUITE'] and scope or 'TESTCASE'

    def get_instance(self):
        try:
            if self._libinst is None:
                self._libinst = self._get_instance()
        except AttributeError:
            self._libinst = self._get_instance()
        return self._libinst

    def _get_instance(self):
        try:
            return self._libcode(*self.args)
        except:
            self._raise_creating_instance_failed()

    def _create_handlers(self, syslog):
        success, failure, details = self._get_reporting_methods(syslog)
        handlers = utils.NormalizedDict(ignore=['_'])
        for name in self._get_handler_names(self._libinst):
            err_pre = "Adding keyword '%s' to library '%s' failed: " % (name, self.name)
            try:
                method = self._get_handler_method(self._libinst, name)
                success("Got handler method '%s'" % (name))
            except TypeError:
                failure(err_pre + 'Not a method or function')
                continue
            except:
                err_msg, err_details = utils.get_error_details()
                failure(err_pre + 'Getting handler method failed: ' + err_msg)
                details('Details:\n%s' % err_details)
                continue
            try:
                handlers[name] = self._create_handler(name, method)
                success("Created keyword '%s'" % handlers[name].name)
            except:
                err_msg, err_details = utils.get_error_details()
                failure(err_pre + 'Creating keyword failed: ' + err_msg)
                details('Details:\n%s' % err_details)
        return handlers

    def _get_reporting_methods(self, syslog):
        # success, failure, details
        return syslog.debug, syslog.info, syslog.debug

    def _get_handler_names(self, libcode):
        return [ name for name in dir(libcode)
                 if not name.startswith('_') and name != 'ROBOT_LIBRARY_SCOPE' ]

    def _get_handler_method(self, libcode, name):
        method = getattr(libcode, name)
        if not isinstance(method, (types.MethodType, types.FunctionType)):
            raise TypeError('Not a method or function')
        return method

    def _raise_creating_instance_failed(self):
        error_msg, error_details = utils.get_error_details()
        msg = "Creating an instance of the test library '%s' " % self.name
        if len(self.args) == 0:
            msg += "with no arguments "
        elif len(self.args) == 1:
            msg += "with argument '%s' " % self.args[0]
        else:
            msg += "with arguments %s " % utils.seq2str(self.args)
        msg += "failed: " + error_msg
        msg += "\n" + error_details
        raise DataError(msg)

class ModuleLibrary(_BaseTestLibrary):

    def _get_scope(self, libcode):
        return 'GLOBAL'

    def get_instance(self):
        return self._libcode

    def _create_handler(self, handler_name, handler_method):
        return PythonHandler(self, handler_name, handler_method)

class PythonLibrary(_BaseTestLibrary):

    def _create_handler(self, handler_name, handler_method):
        if self._is_java_method(handler_method):
            return JavaHandler(self, handler_name, handler_method)
        return PythonHandler(self, handler_name, handler_method)

    def _is_java_method(self, method):
        if not utils.is_jython:
            return False
        try:
            return 'reflectedfunction' in str(type(method.im_func))
        except AttributeError:
            return False

class HybridLibrary(PythonLibrary):

    def _get_handler_names(self, instance):
        try:
            return instance.get_keyword_names()
        except AttributeError:
            return instance.getKeywordNames()

    def _get_reporting_methods(self, syslog):
        # Use syslog.warn for reporting possible failures when creating kws
        # to make them visible. With hybrid libraries kw names are returned
        # explicitly so creating them should also pass.
        return syslog.debug, syslog.warn, syslog.info

class DynamicLibrary(_BaseTestLibrary):

    _get_kw_doc_names = ['get_keyword_documentation', 'getKeywordDocumentation']
    _get_kw_args_names = ['get_keyword_arguments', 'getKeywordArguments']

    def __init__(self, libcode, source, name, args, syslog):
        self._get_keyword_documentation = self._get_method(libcode, *self._get_kw_doc_names)
        self._get_keyword_arguments = self._get_method(libcode, *self._get_kw_args_names)
        _BaseTestLibrary.__init__(self, libcode, source, name, args, syslog)

    def _get_method(self, libcode, name, alternative):
        try:
            code = getattr(libcode, name)
            return code
        except AttributeError:
            try:
                code = getattr(libcode, alternative)
                return code
            except:
                return None

    def _get_handler_names(self, instance):
        try:
            return instance.get_keyword_names()
        except AttributeError:
            return instance.getKeywordNames()

    def _get_handler_method(self, instance, name_is_ignored):
        try:
            return instance.run_keyword
        except AttributeError:
            return instance.runKeyword

    def _create_handler(self, handler_name, handler_method):
        doc = self._get_kw_doc(handler_name)
        argspec = self._get_kw_argspec(handler_name)
        return DynamicHandler(self, handler_name, handler_method, doc, argspec)

    def _get_kw_doc(self, name):
        if self._get_keyword_documentation is not None:
            return self._get_keyword_documentation(self._libinst, name)
        return ''

    def _get_kw_argspec(self, name):
        if self._get_keyword_arguments is not None:
            return self._get_keyword_arguments(self._libinst, name)
        return None

class JavaLibrary(_BaseTestLibrary):

    def _create_handler(self, handler_name, handler_method):
        return JavaHandler(self, handler_name, handler_method)
    <commit_msg>Ignore all attrs starting with ROBOT_LIBRARY_. Prevents getting warnings to syslog about recently added _VERSION.<commit_after>#  Copyright 2008 Nokia Siemens Networks Oyj
#
#  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.

import os
import types

from robot import utils
from robot.errors import DataError
from robot.common import BaseLibrary
from robot.output import SystemLogger

from handlers import PythonHandler, JavaHandler, DynamicHandler

def TestLibrary(name, args=None, syslog=None):
    if syslog is None:
        syslog = SystemLogger()
    libcode, source = utils.import_(name)
    args = utils.to_list(args)
    if isinstance(libcode, types.ModuleType):
        if args:
            raise DataError('Libraries implemented as modules do not take '
                            'arguments, got: %s' % str(args))
        return ModuleLibrary(libcode, source, name, args, syslog)
    if _has_method(libcode, ['get_keyword_names', 'getKeywordNames']):
        if _has_method(libcode, ['run_keyword', 'runKeyword']):
            return DynamicLibrary(libcode, source, name, args, syslog)
        else:
            return HybridLibrary(libcode, source, name, args, syslog)
    # Using type check and not isinstance for ClassType, because it does not
    # match Java classes whose type is javaclass.
    if type(libcode) is types.ClassType or isinstance(libcode, types.TypeType):
        return PythonLibrary(libcode, source, name, args, syslog)
    return JavaLibrary(libcode, source, name, args, syslog)

def _has_method(code, names):
    for name in names:
        if hasattr(code, name) and _is_method(getattr(code, name)):
            return True
    return False

def _is_method(code):
    return isinstance(code, types.MethodType) or \
            (utils.is_jython and 'reflectedfunction' in str(type(code)))

class _BaseTestLibrary(BaseLibrary):

    def __init__(self, libcode, source, name, args, syslog):
        if os.path.exists(name):
            name = os.path.splitext(os.path.basename(name))[0]
        self.source = source
        self.version = self._get_version(libcode)
        self.name = name
        self.orig_name = name # Stores original name also after copying
        self.args = args
        self._instance_cache = []
        if libcode is not None:
            self.doc = utils.get_doc(libcode)
            self.scope = self._get_scope(libcode)
            self._libcode = libcode
            self._libinst = self.get_instance()
            self.handlers = self._create_handlers(syslog)
            self._init_scope_handling(self.scope)

    def _get_version(self, code):
        try:
            return str(code.ROBOT_LIBRARY_VERSION)
        except AttributeError:
            try:
                return str(code.__version__)
            except AttributeError:
                return '<unknown>'

    def _init_scope_handling(self, scope):
        if scope == 'GLOBAL':
            return
        self._libinst = None
        self.start_suite = self._caching_start
        self.end_suite = self._restoring_end
        if scope == 'TESTCASE':
            self.start_test = self._caching_start
            self.end_test = self._restoring_end

    def copy(self, name):
        lib = _BaseTestLibrary(None, self.version, name, self.args, None)
        lib.orig_name = self.name
        lib.doc = self.doc
        lib.scope = self.scope
        lib._libcode = self._libcode
        lib._libinst = self._libinst
        lib.handlers = utils.NormalizedDict(ignore=['_'])
        for name, handler in self.handlers.items():
            lib.handlers[name] = handler.copy(lib)
        return lib

    def start_suite(self):
        pass

    def end_suite(self):
        pass

    def start_test(self):
        pass

    def end_test(self):
        pass

    def _caching_start(self):
        self._instance_cache.append(self._libinst)
        self._libinst = None

    def _restoring_end(self):
        self._libinst = self._instance_cache.pop()

    def _get_scope(self, libcode):
        try:
            scope = libcode.ROBOT_LIBRARY_SCOPE
            scope = utils.normalize(scope, ignore=['_']).upper()
        except:
            scope = 'TESTCASE'
        return scope in ['GLOBAL','TESTSUITE'] and scope or 'TESTCASE'

    def get_instance(self):
        try:
            if self._libinst is None:
                self._libinst = self._get_instance()
        except AttributeError:
            self._libinst = self._get_instance()
        return self._libinst

    def _get_instance(self):
        try:
            return self._libcode(*self.args)
        except:
            self._raise_creating_instance_failed()

    def _create_handlers(self, syslog):
        success, failure, details = self._get_reporting_methods(syslog)
        handlers = utils.NormalizedDict(ignore=['_'])
        for name in self._get_handler_names(self._libinst):
            err_pre = "Adding keyword '%s' to library '%s' failed: " % (name, self.name)
            try:
                method = self._get_handler_method(self._libinst, name)
                success("Got handler method '%s'" % (name))
            except TypeError:
                failure(err_pre + 'Not a method or function')
                continue
            except:
                err_msg, err_details = utils.get_error_details()
                failure(err_pre + 'Getting handler method failed: ' + err_msg)
                details('Details:\n%s' % err_details)
                continue
            try:
                handlers[name] = self._create_handler(name, method)
                success("Created keyword '%s'" % handlers[name].name)
            except:
                err_msg, err_details = utils.get_error_details()
                failure(err_pre + 'Creating keyword failed: ' + err_msg)
                details('Details:\n%s' % err_details)
        return handlers

    def _get_reporting_methods(self, syslog):
        # success, failure, details
        return syslog.debug, syslog.info, syslog.debug

    def _get_handler_names(self, libcode):
        return [ name for name in dir(libcode)
                 if not (name.startswith('_') or
                         name.startswith('ROBOT_LIBRARY_')) ]

    def _get_handler_method(self, libcode, name):
        method = getattr(libcode, name)
        if not isinstance(method, (types.MethodType, types.FunctionType)):
            raise TypeError('Not a method or function')
        return method

    def _raise_creating_instance_failed(self):
        error_msg, error_details = utils.get_error_details()
        msg = "Creating an instance of the test library '%s' " % self.name
        if len(self.args) == 0:
            msg += "with no arguments "
        elif len(self.args) == 1:
            msg += "with argument '%s' " % self.args[0]
        else:
            msg += "with arguments %s " % utils.seq2str(self.args)
        msg += "failed: " + error_msg
        msg += "\n" + error_details
        raise DataError(msg)

class ModuleLibrary(_BaseTestLibrary):

    def _get_scope(self, libcode):
        return 'GLOBAL'

    def get_instance(self):
        return self._libcode

    def _create_handler(self, handler_name, handler_method):
        return PythonHandler(self, handler_name, handler_method)

class PythonLibrary(_BaseTestLibrary):

    def _create_handler(self, handler_name, handler_method):
        if self._is_java_method(handler_method):
            return JavaHandler(self, handler_name, handler_method)
        return PythonHandler(self, handler_name, handler_method)

    def _is_java_method(self, method):
        if not utils.is_jython:
            return False
        try:
            return 'reflectedfunction' in str(type(method.im_func))
        except AttributeError:
            return False

class HybridLibrary(PythonLibrary):

    def _get_handler_names(self, instance):
        try:
            return instance.get_keyword_names()
        except AttributeError:
            return instance.getKeywordNames()

    def _get_reporting_methods(self, syslog):
        # Use syslog.warn for reporting possible failures when creating kws
        # to make them visible. With hybrid libraries kw names are returned
        # explicitly so creating them should also pass.
        return syslog.debug, syslog.warn, syslog.info

class DynamicLibrary(_BaseTestLibrary):

    _get_kw_doc_names = ['get_keyword_documentation', 'getKeywordDocumentation']
    _get_kw_args_names = ['get_keyword_arguments', 'getKeywordArguments']

    def __init__(self, libcode, source, name, args, syslog):
        self._get_keyword_documentation = self._get_method(libcode, *self._get_kw_doc_names)
        self._get_keyword_arguments = self._get_method(libcode, *self._get_kw_args_names)
        _BaseTestLibrary.__init__(self, libcode, source, name, args, syslog)

    def _get_method(self, libcode, name, alternative):
        try:
            code = getattr(libcode, name)
            return code
        except AttributeError:
            try:
                code = getattr(libcode, alternative)
                return code
            except:
                return None

    def _get_handler_names(self, instance):
        try:
            return instance.get_keyword_names()
        except AttributeError:
            return instance.getKeywordNames()

    def _get_handler_method(self, instance, name_is_ignored):
        try:
            return instance.run_keyword
        except AttributeError:
            return instance.runKeyword

    def _create_handler(self, handler_name, handler_method):
        doc = self._get_kw_doc(handler_name)
        argspec = self._get_kw_argspec(handler_name)
        return DynamicHandler(self, handler_name, handler_method, doc, argspec)

    def _get_kw_doc(self, name):
        if self._get_keyword_documentation is not None:
            return self._get_keyword_documentation(self._libinst, name)
        return ''

    def _get_kw_argspec(self, name):
        if self._get_keyword_arguments is not None:
            return self._get_keyword_arguments(self._libinst, name)
        return None

class JavaLibrary(_BaseTestLibrary):

    def _create_handler(self, handler_name, handler_method):
        return JavaHandler(self, handler_name, handler_method)

<|endoftext|>


<commit_before># Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import telemetry.core.timeline.bounds as timeline_bounds
from telemetry.page.actions import page_action
from telemetry.page.actions import wait

class GestureAction(page_action.PageAction):
  def __init__(self, attributes=None):
    super(GestureAction, self).__init__(attributes)

    if hasattr(self, 'wait_after'):
      self.wait_action = wait.WaitAction(self.wait_after)
    else:
      self.wait_action = None

  def RunAction(self, page, tab, previous_action):
    tab.ExecuteJavaScript(
        'console.time("' + self._GetUniqueTimelineMarkerName() + '")')

    self.RunGesture(page, tab, previous_action)

    tab.ExecuteJavaScript(
        'console.timeEnd("' + self._GetUniqueTimelineMarkerName() + '")')

    if self.wait_action:
      self.wait_action.RunAction(page, tab, previous_action)

  def RunGesture(self, page, tab, previous_action):
    raise NotImplementedError()

  @staticmethod
  def GetGestureSourceTypeFromOptions(tab):
    gesture_source_type = tab.browser.synthetic_gesture_source_type
    return 'chrome.gpuBenchmarking.' + gesture_source_type.upper() + '_INPUT'

  def CustomizeBrowserOptionsForPageSet(self, options):
    options.AppendExtraBrowserArgs('--enable-gpu-benchmarking')

  def GetActiveRangeOnTimeline(self, timeline):
    action_range = super(GestureAction, self).GetActiveRangeOnTimeline(timeline)
    if action_range.is_empty:
      raise page_action.PageActionInvalidTimelineMarker(
          'Gesture action requires timeline marker to determine active range.')

    # The synthetic gesture controller inserts a trace marker to precisely
    # demarcate when the gesture was running. Find the trace marker that belongs
    # to this action.
    gesture_events = [
        ev for ev
        in timeline.GetAllEventsOfName('SyntheticGestureController::running',
                                       True)
        if ev.start >= action_range.min and
           ev.start + ev.duration <= action_range.max ]
    if len(gesture_events) == 0:
      raise page_action.PageActionInvalidTimelineMarker(
          'No valid synthetic gesture marker found in timeline.')
    if len(gesture_events) > 1:
      raise page_action.PageActionInvalidTimelineMarker(
          'More than one possible synthetic gesture marker found in timeline.')

    active_range = timeline_bounds.Bounds.CreateFromEvent(gesture_events[0])

    if self.wait_action:
      active_range.AddBounds(
          self.wait_action.GetActiveRangeOnTimeline(timeline))

    return active_range
<commit_msg>telemetry: Check for overlap, not inclusion, in GetActiveRangeOnTimeline.<commit_after># Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import telemetry.core.timeline.bounds as timeline_bounds
from telemetry.page.actions import page_action
from telemetry.page.actions import wait

class GestureAction(page_action.PageAction):
  def __init__(self, attributes=None):
    super(GestureAction, self).__init__(attributes)

    if hasattr(self, 'wait_after'):
      self.wait_action = wait.WaitAction(self.wait_after)
    else:
      self.wait_action = None

  def RunAction(self, page, tab, previous_action):
    tab.ExecuteJavaScript(
        'console.time("' + self._GetUniqueTimelineMarkerName() + '")')

    self.RunGesture(page, tab, previous_action)

    tab.ExecuteJavaScript(
        'console.timeEnd("' + self._GetUniqueTimelineMarkerName() + '")')

    if self.wait_action:
      self.wait_action.RunAction(page, tab, previous_action)

  def RunGesture(self, page, tab, previous_action):
    raise NotImplementedError()

  @staticmethod
  def GetGestureSourceTypeFromOptions(tab):
    gesture_source_type = tab.browser.synthetic_gesture_source_type
    return 'chrome.gpuBenchmarking.' + gesture_source_type.upper() + '_INPUT'

  def CustomizeBrowserOptionsForPageSet(self, options):
    options.AppendExtraBrowserArgs('--enable-gpu-benchmarking')

  def GetActiveRangeOnTimeline(self, timeline):
    action_range = super(GestureAction, self).GetActiveRangeOnTimeline(timeline)
    if action_range.is_empty:
      raise page_action.PageActionInvalidTimelineMarker(
          'Gesture action requires timeline marker to determine active range.')

    # The synthetic gesture controller inserts a trace marker to precisely
    # demarcate when the gesture was running. Find the trace marker that belongs
    # to this action. We check for overlap, not inclusion, because
    # gesture_actions can start/end slightly outside the action_range on
    # Windows. This problem is probably caused by a race condition between
    # the browser and renderer process submitting the trace events for the
    # markers.
    gesture_events = [
        ev for ev
        in timeline.GetAllEventsOfName('SyntheticGestureController::running',
                                       True)
        if ev.start <= action_range.max and
           ev.start + ev.duration >= action_range.min ]
    if len(gesture_events) == 0:
      raise page_action.PageActionInvalidTimelineMarker(
          'No valid synthetic gesture marker found in timeline.')
    if len(gesture_events) > 1:
      raise page_action.PageActionInvalidTimelineMarker(
          'More than one possible synthetic gesture marker found in timeline.')

    active_range = timeline_bounds.Bounds.CreateFromEvent(gesture_events[0])

    if self.wait_action:
      active_range.AddBounds(
          self.wait_action.GetActiveRangeOnTimeline(timeline))

    return active_range
<|endoftext|>


<commit_before>#  Copyright 2008-2012 Nokia Siemens Networks Oyj
#
#  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.

import sys
import os

UNIXY = os.sep == '/'
JYTHON = sys.platform.startswith('java')
if UNIXY:
    DEFAULT_SYSTEM_ENCODING = 'UTF-8'
    DEFAULT_OUTPUT_ENCODING = 'UTF-8'
else:
    DEFAULT_SYSTEM_ENCODING = 'cp1252'
    DEFAULT_OUTPUT_ENCODING = 'cp437'

def get_system_encoding():
    encoding = sys.getfilesystemencoding()
    if not encoding and JYTHON:
        from java.lang import System
        encoding = System.getProperty('file.encoding')
    return encoding or DEFAULT_SYSTEM_ENCODING

def get_output_encoding():
    if _on_buggy_jython():
        return DEFAULT_OUTPUT_ENCODING
    encoding = _get_encoding_from_standard_streams()
    if not encoding and UNIXY:
        encoding = _get_encoding_from_unix_environment_variables()
    return encoding or DEFAULT_OUTPUT_ENCODING

def _on_buggy_jython():
    # http://bugs.jython.org/issue1568
    if UNIXY or not JYTHON:
        return False
    return sys.platform.startswith('java1.5') or sys.version_info < (2, 5, 2)

def _get_encoding_from_standard_streams():
    # Stream may not have encoding attribute if it is intercepted outside RF
    # in Python. Encoding is None if process's outputs are redirected.
    for stream in sys.__stdout__, sys.__stderr__, sys.__stdin__:
        encoding = getattr(stream, 'encoding', None)
        if encoding and _is_valid_encoding(encoding):
            return encoding
    return None

def _get_encoding_from_unix_environment_variables():
    for name in 'LANG', 'LC_CTYPE', 'LANGUAGE', 'LC_ALL':
        if name in os.environ:
            # Encoding can be in format like `UTF-8` or `en_US.UTF-8`
            encoding = os.environ[name].split('.')[-1]
            if _is_valid_encoding(encoding):
                return encoding
    return None

def _is_valid_encoding(encoding):
    try:
        'test'.encode(encoding)
    except LookupError:
        return False
    else:
        return True
<commit_msg>System encoding: test sniffed encoding before using it<commit_after>#  Copyright 2008-2012 Nokia Siemens Networks Oyj
#
#  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.

import sys
import os

UNIXY = os.sep == '/'
JYTHON = sys.platform.startswith('java')
if UNIXY:
    DEFAULT_SYSTEM_ENCODING = 'UTF-8'
    DEFAULT_OUTPUT_ENCODING = 'UTF-8'
else:
    DEFAULT_SYSTEM_ENCODING = 'cp1252'
    DEFAULT_OUTPUT_ENCODING = 'cp437'

def get_system_encoding():
    encoding = _get_python_file_system_encoding()
    if not encoding and JYTHON:
        encoding = _get_java_file_system_encoding()
    return encoding or DEFAULT_SYSTEM_ENCODING

def _get_python_file_system_encoding():
    encoding = sys.getfilesystemencoding()
    return encoding if _is_valid(encoding) else None

def _get_java_file_system_encoding():
    from java.lang import System
    encoding = System.getProperty('file.encoding')
    return encoding if _is_valid(encoding) else None

def _is_valid(encoding):
    if not encoding:
        return False
    try:
        'test'.encode(encoding)
    except LookupError:
        return False
    else:
        return True

def get_output_encoding():
    if _on_buggy_jython():
        return DEFAULT_OUTPUT_ENCODING
    encoding = _get_encoding_from_standard_streams()
    if not encoding and UNIXY:
        encoding = _get_encoding_from_unix_environment_variables()
    return encoding or DEFAULT_OUTPUT_ENCODING

def _on_buggy_jython():
    # http://bugs.jython.org/issue1568
    if UNIXY or not JYTHON:
        return False
    return sys.platform.startswith('java1.5') or sys.version_info < (2, 5, 2)

def _get_encoding_from_standard_streams():
    # Stream may not have encoding attribute if it is intercepted outside RF
    # in Python. Encoding is None if process's outputs are redirected.
    for stream in sys.__stdout__, sys.__stderr__, sys.__stdin__:
        encoding = getattr(stream, 'encoding', None)
        if _is_valid(encoding):
            return encoding
    return None

def _get_encoding_from_unix_environment_variables():
    for name in 'LANG', 'LC_CTYPE', 'LANGUAGE', 'LC_ALL':
        if name in os.environ:
            # Encoding can be in format like `UTF-8` or `en_US.UTF-8`
            encoding = os.environ[name].split('.')[-1]
            if _is_valid(encoding):
                return encoding
    return None
<|endoftext|>


<commit_before># coding: utf-8
# pylint: disable=invalid-name, protected-access, fixme, too-many-arguments, W0221, W0201

"""NDArray interface of mxnet"""
from __future__ import absolute_import
from collections import namedtuple

import ctypes
import sys
import numpy as np
import logging
from .base import _LIB
from .base import c_array, c_str, mx_uint, py_str
from .base import DataIterHandle, NDArrayHandle
from .base import check_call, ctypes2docstring
from .ndarray import NDArray
from .ndarray import array

class DataIter(object):
    """DataIter object in mxnet. """

    def __init__(self):
        pass

    def __iter__(self):
        return self

    def reset(self):
        """Reset the iterator. """
        pass

    def next(self):
        """Get next data batch from iterator

        Returns
        -------
        data : NDArray
            The data of next batch.

        label : NDArray
            The label of next batch.
        """
        pass

    def __next__(self):
        return self.next()

    def iter_next(self):
        """Iterate to next batch.

        Returns
        -------
        has_next : boolean
            Whether the move is successful.
        """
        pass

    def getdata(self, index=0):
        """Get data of current batch.

        Parameters
        ----------
        index : int
            The index of data source to retrieve.

        Returns
        -------
        data : NDArray
            The data of current batch.
        """
        pass

    def getlabel(self):
        """Get label of current batch.

        Returns
        -------
        label : NDArray
            The label of current batch.
        """
        return self.getdata(-1)

    def getpad(self):
        """Get the number of padding examples in current batch.

        Returns
        -------
        pad : int
            Number of padding examples in current batch
        """
        pass

DataBatch = namedtuple('DataBatch', ['data', 'label', 'pad'])

def _init_data(data, allow_empty, default_name):
    """Convert data into canonical form."""
    assert (data is not None) or allow_empty
    if data is None:
        data = []

    if isinstance(data, (np.ndarray, NDArray)):
        data = [data]
    if isinstance(data, list):
        if not allow_empty:
            assert(len(data) > 0)
        if len(data) == 1:
            data = {default_name: data[0]}
        else:
            data = {'_%d_%s' % (i, default_name) : d for i, d in enumerate(data)}
    if not isinstance(data, dict):
        raise TypeError("Input must be NDArray, numpy.ndarray, " + \
                "a list of them or dict with them as values")
    for k, v in data.iteritems():
        if isinstance(v, NDArray):
            data[k] = v.asnumpy()
    for k, v in data.iteritems():
        if not isinstance(v, np.ndarray):
            raise TypeError(("Invalid type '%s' for %s, "  % (type(v), k)) + \
                    "should be NDArray or numpy.ndarray")
    return data.items()

class NDArrayIter(DataIter):
    """NDArrayIter object in mxnet. Taking NDArray or numpy array to get dataiter.

    Parameters
    ----------
    data_list or data, label: a list of, or two separate NDArray or numpy.ndarray
        list of NDArray for data. The last one is treated as label.

    batch_size: int
        Batch Size

    shuffle: bool
        Whether to shuffle the data

    data_pad_value: float, optional
        Padding value for data

    label_pad_value: float, optionl
        Padding value for label

    last_batch_handle: 'pad', 'discard' or 'roll_over'
        How to handle the last batch

    Note
    ----
    This iterator will pad, discard or roll over the last batch if
    the size of data does not match batch_size. Roll over is intended
    for training and can cause problems if used for prediction.
    """
    def __init__(self, data, label=None, batch_size=1, shuffle=False, last_batch_handle='pad'):
        # pylint: disable=W0201

        super(NDArrayIter, self).__init__()

        self.data = _init_data(data, allow_empty=False, default_name='data')
        self.label = _init_data(label, allow_empty=True, default_name='softmax_label')

        self.data_list = [x[1] for x in self.data] + [x[1] for x in self.label]
        self.num_source = len(self.data_list)

        # shuffle data
        if shuffle:
            idx = np.arange(self.data_list[0].shape[0])
            np.random.shuffle(idx)
            for i in range(self.num_source):
                assert self.data_list[i].shape[0] == len(idx)
                np.copyto(self.data_list[i], self.data_list[i][idx])

        # batching
        if last_batch_handle == 'discard':
            new_n = self.data_list[0].shape[0] - self.data_list[0].shape[0] % batch_size
            for k, _ in self.data.iteritems():
                self.data[k] = self.data[k][:new_n]
            for k, _ in self.label.iteritems():
                self.label[k] = self.label[k][:new_n]
        self.num_data = self.data_list[0].shape[0]
        assert self.num_data > batch_size, \
            "batch_size need to be smaller than data size when not padding."
        self.cursor = -batch_size
        self.batch_size = batch_size
        self.last_batch_handle = last_batch_handle

    @property
    def provide_data(self):
        """The name and shape of data provided by this iterator"""
        return [(k, tuple([self.batch_size] + list(v.shape[1:]))) for k, v in self.data]

    @property
    def provide_label(self):
        """The name and shape of label provided by this iterator"""
        return [(k, tuple([self.batch_size] + list(v.shape[1:]))) for k, v in self.label]

    def hard_reset(self):
        """Igore roll over data and set to start"""
        self.cursor = -self.batch_size

    def reset(self):
        if self.last_batch_handle == 'roll_over' and self.cursor > self.num_data:
            self.cursor = -self.batch_size + (self.cursor%self.num_data)%self.batch_size
        else:
            self.cursor = -self.batch_size

    def iter_next(self):
        self.cursor += self.batch_size
        if self.cursor < self.num_data:
            return True
        else:
            return False

    def next(self):
        if self.iter_next():
            return DataBatch(data=self.getdata(), label=self.getlabel(), pad=self.getpad())
        else:
            raise StopIteration

    def _getdata(self, data_source):
        """Load data from underlying arrays, internal use only"""
        assert(self.cursor < self.num_data), "DataIter needs reset."
        if self.cursor + self.batch_size <= self.num_data:
            return [array(x[1][self.cursor:self.cursor+self.batch_size]) for x in data_source]
        else:
            pad = self.batch_size - self.num_data + self.cursor
            return [array(np.concatenate((x[1][self.cursor:], x[1][:pad]),
                                         axis=0)) for x in data_source]

    def getdata(self):
        return self._getdata(self.data)

    def getlabel(self):
        return self._getdata(self.label)

    def getpad(self):
        if self.last_batch_handle == 'pad' and \
           self.cursor + self.batch_size > self.num_data:
            return self.cursor + self.batch_size - self.num_data
        else:
            return 0

class MXDataIter(DataIter):
    """DataIter built in MXNet. List all the needed functions here.

    Parameters
    ----------
    handle : DataIterHandle
        the handle to the underlying C++ Data Iterator
    """
    def __init__(self, handle, data_name='data', label_name='softmax_label', **_):
        super(MXDataIter, self).__init__()
        self.handle = handle
        # debug option, used to test the speed with io effect eliminated
        self._debug_skip_load = False

        # load the first batch to get shape information
        self.reset()
        batch = self.next()
        data = batch.data[0]
        label = batch.label[0]
        self.reset()

        # properties
        self.provide_data = [(data_name, data.shape)]
        self.provide_label = [(label_name, label.shape)]
        self.batch_size = data.shape[0]

    def __del__(self):
        check_call(_LIB.MXDataIterFree(self.handle))

    def debug_skip_load(self):
        """Set the iterator to simply return always first batch.

        Notes
        -----
        This can be used to test the speed of network without taking
        the loading delay into account.
        """
        self._debug_skip_load = True
        logging.info('Set debug_skip_load to be true, will simply return first batch')

    def reset(self):
        self._debug_at_begin = True
        check_call(_LIB.MXDataIterBeforeFirst(self.handle))

    def next(self):
        if self._debug_skip_load and not self._debug_at_begin:
            return  DataBatch(data=[self.getdata()], label=[self.getlabel()], pad=self.getpad())

        self._debug_at_begin = False
        next_res = ctypes.c_int(0)
        check_call(_LIB.MXDataIterNext(self.handle, ctypes.byref(next_res)))
        if next_res.value:
            return DataBatch(data=[self.getdata()], label=[self.getlabel()], pad=self.getpad())
        else:
            raise StopIteration

    def iter_next(self):
        next_res = ctypes.c_int(0)
        check_call(_LIB.MXDataIterNext(self.handle, ctypes.byref(next_res)))
        return next_res.value

    def getdata(self):
        hdl = NDArrayHandle()
        check_call(_LIB.MXDataIterGetData(self.handle, ctypes.byref(hdl)))
        return NDArray(hdl, False)

    def getlabel(self):
        hdl = NDArrayHandle()
        check_call(_LIB.MXDataIterGetLabel(self.handle, ctypes.byref(hdl)))
        return NDArray(hdl, False)

    def getpad(self):
        pad = ctypes.c_int(0)
        check_call(_LIB.MXDataIterGetPadNum(self.handle, ctypes.byref(pad)))
        return pad.value

def _make_io_iterator(handle):
    """Create an io iterator by handle."""
    name = ctypes.c_char_p()
    desc = ctypes.c_char_p()
    num_args = mx_uint()
    arg_names = ctypes.POINTER(ctypes.c_char_p)()
    arg_types = ctypes.POINTER(ctypes.c_char_p)()
    arg_descs = ctypes.POINTER(ctypes.c_char_p)()

    check_call(_LIB.MXDataIterGetIterInfo( \
            handle, ctypes.byref(name), ctypes.byref(desc), \
            ctypes.byref(num_args), \
            ctypes.byref(arg_names), \
            ctypes.byref(arg_types), \
            ctypes.byref(arg_descs)))
    iter_name = py_str(name.value)
    param_str = ctypes2docstring(num_args, arg_names, arg_types, arg_descs)

    doc_str = ('%s\n\n' +
               '%s\n' +
               'name : string, required.\n' +
               '    Name of the resulting data iterator.\n\n' +
               'Returns\n' +
               '-------\n' +
               'iterator: DataIter\n'+
               '    The result iterator.')
    doc_str = doc_str % (desc.value, param_str)

    def creator(*args, **kwargs):
        """Create an iterator.
        The parameters listed below can be passed in as keyword arguments.

        Parameters
        ----------
        name : string, required.
            Name of the resulting data iterator.

        Returns
        -------
        dataiter: Dataiter
            the resulting data iterator
        """
        param_keys = []
        param_vals = []

        for k, val in kwargs.items():
            param_keys.append(c_str(k))
            param_vals.append(c_str(str(val)))
        # create atomic symbol
        param_keys = c_array(ctypes.c_char_p, param_keys)
        param_vals = c_array(ctypes.c_char_p, param_vals)
        iter_handle = DataIterHandle()
        check_call(_LIB.MXDataIterCreateIter(
            handle,
            mx_uint(len(param_keys)),
            param_keys, param_vals,
            ctypes.byref(iter_handle)))

        if len(args):
            raise TypeError('%s can only accept keyword arguments' % iter_name)

        return MXDataIter(iter_handle, **kwargs)

    creator.__name__ = iter_name
    creator.__doc__ = doc_str
    return creator

def _init_io_module():
    """List and add all the data iterators to current module."""
    plist = ctypes.POINTER(ctypes.c_void_p)()
    size = ctypes.c_uint()
    check_call(_LIB.MXListDataIters(ctypes.byref(size), ctypes.byref(plist)))
    module_obj = sys.modules[__name__]
    for i in range(size.value):
        hdl = ctypes.c_void_p(plist[i])
        dataiter = _make_io_iterator(hdl)
        setattr(module_obj, dataiter.__name__, dataiter)

# Initialize the io in startups
_init_io_module()
<commit_msg>fix lint<commit_after># coding: utf-8
# pylint: disable=invalid-name, protected-access, fixme, too-many-arguments, W0221, W0201

"""NDArray interface of mxnet"""
from __future__ import absolute_import
from collections import namedtuple

import ctypes
import sys
import numpy as np
import logging
from .base import _LIB
from .base import c_array, c_str, mx_uint, py_str
from .base import DataIterHandle, NDArrayHandle
from .base import check_call, ctypes2docstring
from .ndarray import NDArray
from .ndarray import array

class DataIter(object):
    """DataIter object in mxnet. """

    def __init__(self):
        pass

    def __iter__(self):
        return self

    def reset(self):
        """Reset the iterator. """
        pass

    def next(self):
        """Get next data batch from iterator

        Returns
        -------
        data : NDArray
            The data of next batch.

        label : NDArray
            The label of next batch.
        """
        pass

    def __next__(self):
        return self.next()

    def iter_next(self):
        """Iterate to next batch.

        Returns
        -------
        has_next : boolean
            Whether the move is successful.
        """
        pass

    def getdata(self, index=0):
        """Get data of current batch.

        Parameters
        ----------
        index : int
            The index of data source to retrieve.

        Returns
        -------
        data : NDArray
            The data of current batch.
        """
        pass

    def getlabel(self):
        """Get label of current batch.

        Returns
        -------
        label : NDArray
            The label of current batch.
        """
        return self.getdata(-1)

    def getpad(self):
        """Get the number of padding examples in current batch.

        Returns
        -------
        pad : int
            Number of padding examples in current batch
        """
        pass

DataBatch = namedtuple('DataBatch', ['data', 'label', 'pad'])

def _init_data(data, allow_empty, default_name):
    """Convert data into canonical form."""
    assert (data is not None) or allow_empty
    if data is None:
        data = []

    if isinstance(data, (np.ndarray, NDArray)):
        data = [data]
    if isinstance(data, list):
        if not allow_empty:
            assert(len(data) > 0)
        if len(data) == 1:
            data = {default_name: data[0]}
        else:
            data = {'_%d_%s' % (i, default_name) : d for i, d in enumerate(data)}
    if not isinstance(data, dict):
        raise TypeError("Input must be NDArray, numpy.ndarray, " + \
                "a list of them or dict with them as values")
    for k, v in data.iteritems():
        if isinstance(v, NDArray):
            data[k] = v.asnumpy()
    for k, v in data.iteritems():
        if not isinstance(v, np.ndarray):
            raise TypeError(("Invalid type '%s' for %s, "  % (type(v), k)) + \
                    "should be NDArray or numpy.ndarray")
    return data.items()

class NDArrayIter(DataIter):
    """NDArrayIter object in mxnet. Taking NDArray or numpy array to get dataiter.

    Parameters
    ----------
    data_list or data, label: a list of, or two separate NDArray or numpy.ndarray
        list of NDArray for data. The last one is treated as label.

    batch_size: int
        Batch Size

    shuffle: bool
        Whether to shuffle the data

    data_pad_value: float, optional
        Padding value for data

    label_pad_value: float, optionl
        Padding value for label

    last_batch_handle: 'pad', 'discard' or 'roll_over'
        How to handle the last batch

    Note
    ----
    This iterator will pad, discard or roll over the last batch if
    the size of data does not match batch_size. Roll over is intended
    for training and can cause problems if used for prediction.
    """
    def __init__(self, data, label=None, batch_size=1, shuffle=False, last_batch_handle='pad'):
        # pylint: disable=W0201, E1101
        # disable E1101 b/c lint says numpy has no copyto function, maybe due to different version?

        super(NDArrayIter, self).__init__()

        self.data = _init_data(data, allow_empty=False, default_name='data')
        self.label = _init_data(label, allow_empty=True, default_name='softmax_label')

        self.data_list = [x[1] for x in self.data] + [x[1] for x in self.label]
        self.num_source = len(self.data_list)

        # shuffle data
        if shuffle:
            idx = np.arange(self.data_list[0].shape[0])
            np.random.shuffle(idx)
            for i in range(self.num_source):
                assert self.data_list[i].shape[0] == len(idx)
                np.copyto(self.data_list[i], self.data_list[i][idx])

        # batching
        if last_batch_handle == 'discard':
            new_n = self.data_list[0].shape[0] - self.data_list[0].shape[0] % batch_size
            for k, _ in self.data.iteritems():
                self.data[k] = self.data[k][:new_n]
            for k, _ in self.label.iteritems():
                self.label[k] = self.label[k][:new_n]
        self.num_data = self.data_list[0].shape[0]
        assert self.num_data > batch_size, \
            "batch_size need to be smaller than data size when not padding."
        self.cursor = -batch_size
        self.batch_size = batch_size
        self.last_batch_handle = last_batch_handle

    @property
    def provide_data(self):
        """The name and shape of data provided by this iterator"""
        return [(k, tuple([self.batch_size] + list(v.shape[1:]))) for k, v in self.data]

    @property
    def provide_label(self):
        """The name and shape of label provided by this iterator"""
        return [(k, tuple([self.batch_size] + list(v.shape[1:]))) for k, v in self.label]

    def hard_reset(self):
        """Igore roll over data and set to start"""
        self.cursor = -self.batch_size

    def reset(self):
        if self.last_batch_handle == 'roll_over' and self.cursor > self.num_data:
            self.cursor = -self.batch_size + (self.cursor%self.num_data)%self.batch_size
        else:
            self.cursor = -self.batch_size

    def iter_next(self):
        self.cursor += self.batch_size
        if self.cursor < self.num_data:
            return True
        else:
            return False

    def next(self):
        if self.iter_next():
            return DataBatch(data=self.getdata(), label=self.getlabel(), pad=self.getpad())
        else:
            raise StopIteration

    def _getdata(self, data_source):
        """Load data from underlying arrays, internal use only"""
        assert(self.cursor < self.num_data), "DataIter needs reset."
        if self.cursor + self.batch_size <= self.num_data:
            return [array(x[1][self.cursor:self.cursor+self.batch_size]) for x in data_source]
        else:
            pad = self.batch_size - self.num_data + self.cursor
            return [array(np.concatenate((x[1][self.cursor:], x[1][:pad]),
                                         axis=0)) for x in data_source]

    def getdata(self):
        return self._getdata(self.data)

    def getlabel(self):
        return self._getdata(self.label)

    def getpad(self):
        if self.last_batch_handle == 'pad' and \
           self.cursor + self.batch_size > self.num_data:
            return self.cursor + self.batch_size - self.num_data
        else:
            return 0

class MXDataIter(DataIter):
    """DataIter built in MXNet. List all the needed functions here.

    Parameters
    ----------
    handle : DataIterHandle
        the handle to the underlying C++ Data Iterator
    """
    def __init__(self, handle, data_name='data', label_name='softmax_label', **_):
        super(MXDataIter, self).__init__()
        self.handle = handle
        # debug option, used to test the speed with io effect eliminated
        self._debug_skip_load = False

        # load the first batch to get shape information
        self.reset()
        batch = self.next()
        data = batch.data[0]
        label = batch.label[0]
        self.reset()

        # properties
        self.provide_data = [(data_name, data.shape)]
        self.provide_label = [(label_name, label.shape)]
        self.batch_size = data.shape[0]

    def __del__(self):
        check_call(_LIB.MXDataIterFree(self.handle))

    def debug_skip_load(self):
        """Set the iterator to simply return always first batch.

        Notes
        -----
        This can be used to test the speed of network without taking
        the loading delay into account.
        """
        self._debug_skip_load = True
        logging.info('Set debug_skip_load to be true, will simply return first batch')

    def reset(self):
        self._debug_at_begin = True
        check_call(_LIB.MXDataIterBeforeFirst(self.handle))

    def next(self):
        if self._debug_skip_load and not self._debug_at_begin:
            return  DataBatch(data=[self.getdata()], label=[self.getlabel()], pad=self.getpad())

        self._debug_at_begin = False
        next_res = ctypes.c_int(0)
        check_call(_LIB.MXDataIterNext(self.handle, ctypes.byref(next_res)))
        if next_res.value:
            return DataBatch(data=[self.getdata()], label=[self.getlabel()], pad=self.getpad())
        else:
            raise StopIteration

    def iter_next(self):
        next_res = ctypes.c_int(0)
        check_call(_LIB.MXDataIterNext(self.handle, ctypes.byref(next_res)))
        return next_res.value

    def getdata(self):
        hdl = NDArrayHandle()
        check_call(_LIB.MXDataIterGetData(self.handle, ctypes.byref(hdl)))
        return NDArray(hdl, False)

    def getlabel(self):
        hdl = NDArrayHandle()
        check_call(_LIB.MXDataIterGetLabel(self.handle, ctypes.byref(hdl)))
        return NDArray(hdl, False)

    def getpad(self):
        pad = ctypes.c_int(0)
        check_call(_LIB.MXDataIterGetPadNum(self.handle, ctypes.byref(pad)))
        return pad.value

def _make_io_iterator(handle):
    """Create an io iterator by handle."""
    name = ctypes.c_char_p()
    desc = ctypes.c_char_p()
    num_args = mx_uint()
    arg_names = ctypes.POINTER(ctypes.c_char_p)()
    arg_types = ctypes.POINTER(ctypes.c_char_p)()
    arg_descs = ctypes.POINTER(ctypes.c_char_p)()

    check_call(_LIB.MXDataIterGetIterInfo( \
            handle, ctypes.byref(name), ctypes.byref(desc), \
            ctypes.byref(num_args), \
            ctypes.byref(arg_names), \
            ctypes.byref(arg_types), \
            ctypes.byref(arg_descs)))
    iter_name = py_str(name.value)
    param_str = ctypes2docstring(num_args, arg_names, arg_types, arg_descs)

    doc_str = ('%s\n\n' +
               '%s\n' +
               'name : string, required.\n' +
               '    Name of the resulting data iterator.\n\n' +
               'Returns\n' +
               '-------\n' +
               'iterator: DataIter\n'+
               '    The result iterator.')
    doc_str = doc_str % (desc.value, param_str)

    def creator(*args, **kwargs):
        """Create an iterator.
        The parameters listed below can be passed in as keyword arguments.

        Parameters
        ----------
        name : string, required.
            Name of the resulting data iterator.

        Returns
        -------
        dataiter: Dataiter
            the resulting data iterator
        """
        param_keys = []
        param_vals = []

        for k, val in kwargs.items():
            param_keys.append(c_str(k))
            param_vals.append(c_str(str(val)))
        # create atomic symbol
        param_keys = c_array(ctypes.c_char_p, param_keys)
        param_vals = c_array(ctypes.c_char_p, param_vals)
        iter_handle = DataIterHandle()
        check_call(_LIB.MXDataIterCreateIter(
            handle,
            mx_uint(len(param_keys)),
            param_keys, param_vals,
            ctypes.byref(iter_handle)))

        if len(args):
            raise TypeError('%s can only accept keyword arguments' % iter_name)

        return MXDataIter(iter_handle, **kwargs)

    creator.__name__ = iter_name
    creator.__doc__ = doc_str
    return creator

def _init_io_module():
    """List and add all the data iterators to current module."""
    plist = ctypes.POINTER(ctypes.c_void_p)()
    size = ctypes.c_uint()
    check_call(_LIB.MXListDataIters(ctypes.byref(size), ctypes.byref(plist)))
    module_obj = sys.modules[__name__]
    for i in range(size.value):
        hdl = ctypes.c_void_p(plist[i])
        dataiter = _make_io_iterator(hdl)
        setattr(module_obj, dataiter.__name__, dataiter)

# Initialize the io in startups
_init_io_module()
<|endoftext|>