<commit_before># Copyright 2015.  All rights reserved.
#
#    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 uuid

from cfgm_common import exceptions as vnc_exc
from neutron.common import constants as n_constants
from vnc_api import vnc_api

import contrail_res_handler as res_handler
import router_res_handler as router_handler
import vmi_res_handler as vmi_handler

class FloatingIpMixin(object):

    def _neutron_dict_to_fip_obj(self, fip_q, is_admin=False,
                                 tenant_id=None, fip_obj=None):
        if not fip_obj:
            fip_obj = self._resource_get(id=fip_q['id'])

        vmi_get_handler = vmi_handler.VMInterfaceGetHandler(
            self._vnc_lib)
        port_id = fip_q.get('port_id')
        if port_id:
            vmi_obj = vmi_get_handler.get_vmi_obj(port_id)

            if not is_admin:
                vmi_tenant_id = vmi_get_handler.get_vmi_tenant_id(vmi_obj)
                if vmi_tenant_id != tenant_id:
                    self._raise_contrail_exception('PortNotFound',
                                                   resource='floatingip',
                                                   port_id=port_id)
            fip_obj.set_virtual_machine_interface(vmi_obj)
        else:
            fip_obj.set_virtual_machine_interface_list([])

        if fip_q.get('fixed_ip_address'):
            fip_obj.set_floating_ip_fixed_ip_address(fip_q['fixed_ip_address'])
        else:
            # fixed_ip_address not specified, pick from vmi_obj in create,
            # reset in case of disassociate
            vmi_refs = fip_obj.get_virtual_machine_interface_refs()
            if not vmi_refs:
                fip_obj.set_floating_ip_fixed_ip_address(None)
            else:
                vmi_obj = vmi_get_handler.get_vmi_obj(
                    vmi_refs[0]['uuid'], fields=['instance_ip_back_refs'])

                iip_refs = vmi_obj.get_instance_ip_back_refs()
                if iip_refs:
                    iip_obj = self._vnc_lib.instance_ip_read(
                        id=iip_refs[0]['uuid'])
                    fip_obj.set_floating_ip_fixed_ip_address(
                        iip_obj.get_instance_ip_address())

        return fip_obj

    def _fip_obj_to_neutron_dict(self, fip_obj, fields=None):
        fip_q_dict = {}
        vmi_get_handler = vmi_handler.VMInterfaceGetHandler(
            self._vnc_lib)

        floating_net_id = self._vnc_lib.fq_name_to_id(
            'virtual-network', fip_obj.get_fq_name()[:-2])
        tenant_id = self._project_id_vnc_to_neutron(
            fip_obj.get_project_refs()[0]['uuid'])

        port_id = None
        router_id = None
        vmi_obj = None
        vmi_refs = fip_obj.get_virtual_machine_interface_refs()
        for vmi_ref in vmi_refs or []:
            try:
                vmi_obj = vmi_get_handler.get_vmi_obj(vmi_ref['uuid'])
                port_id = vmi_ref['uuid']
                break
            except vnc_exc.NoIdError:
                pass

        if vmi_obj:
            router_get_handler = router_handler.LogicalRouterGetHandler(
                self._vnc_lib)
            router_id = router_get_handler.get_vmi_obj_router_id(vmi_obj)

        fip_q_dict['id'] = fip_obj.uuid
        fip_q_dict['tenant_id'] = tenant_id
        fip_q_dict['floating_ip_address'] = fip_obj.get_floating_ip_address()
        fip_q_dict['floating_network_id'] = floating_net_id
        fip_q_dict['router_id'] = router_id
        fip_q_dict['port_id'] = port_id
        fip_q_dict['fixed_ip_address'] = (
            fip_obj.get_floating_ip_fixed_ip_address())
        fip_q_dict['status'] = n_constants.PORT_STATUS_ACTIVE

        if fields:
            fip_q_dict = self._filter_res_dict(fip_q_dict, fields)
        return fip_q_dict

class FloatingIpCreateHandler(res_handler.ResourceCreateHandler,
                              FloatingIpMixin):
    resource_create_method = 'floating_ip_create'

    def _create_fip_obj(self, fip_q):
        # TODO() for now create from default pool, later
        # use first available pool on net
        net_id = fip_q['floating_network_id']
        try:
            fq_name = self._vnc_lib.floating_ip_pools_list(
                parent_id=net_id)['floating-ip-pools'][0]['fq_name']
        except IndexError:
            # IndexError could happens when an attempt to
            # retrieve a floating ip pool from a private network.
            msg = "Network %s doesn't provide a floatingip pool" % net_id
            self._raise_contrail_exception('BadRequest',
                                           resource="floatingip", msg=msg)

        fip_pool_obj = self._vnc_lib.floating_ip_pool_read(fq_name=fq_name)
        fip_name = str(uuid.uuid4())
        fip_obj = vnc_api.FloatingIp(fip_name, fip_pool_obj)
        fip_obj.uuid = fip_name

        proj_id = self._project_id_neutron_to_vnc(fip_q['tenant_id'])
        proj_obj = self._project_read(proj_id=proj_id)
        fip_obj.set_project(proj_obj)
        return fip_obj

    def resource_create(self, context, fip_q):
        try:
            fip_obj = self._create_fip_obj(fip_q)
            fip_obj = self._neutron_dict_to_fip_obj(fip_q, context['is_admin'],
                                                    context['tenant'],
                                                    fip_obj=fip_obj)
        except Exception:
            msg = ('Internal error when trying to create floating ip. '
                   'Please be sure the network %s is an external '
                   'network.') % (fip_q['floating_network_id'])
            self._raise_contrail_exception('BadRequest',
                                           resource='floatingip', msg=msg)
        try:
            fip_uuid = self._vnc_lib.floating_ip_create(fip_obj)
        except Exception:
            self._raise_contrail_exception('IpAddressGenerationFailure',
                                           net_id=fip_q['floating_network_id'])
        fip_obj = self._vnc_lib.floating_ip_read(id=fip_uuid)

        return self._fip_obj_to_neutron_dict(fip_obj)

class FloatingIpDeleteHandler(res_handler.ResourceDeleteHandler):
    resource_delete_method = 'floating_ip_delete'

    def resource_delete(self, context, fip_id):
        self._resource_delete(id=fip_id)

class FloatingIpUpdateHandler(res_handler.ResourceUpdateHandler,
                              FloatingIpMixin):
    resource_update_method = 'floating_ip_update'

    def resource_update(self, context, fip_id, fip_q):
        fip_q['id'] = fip_id
        fip_obj = self._neutron_dict_to_fip_obj(fip_q, context['is_admin'],
                                                context['tenant'])
        self._resource_update(fip_obj)
        return self._fip_obj_to_neutron_dict(fip_obj)

class FloatingIpGetHandler(res_handler.ResourceGetHandler, FloatingIpMixin):
    resource_list_method = 'floating_ips_list'
    resource_get_method = 'floating_ip_read'

    def resource_get(self, context, fip_uuid, fields=None):
        try:
            fip_obj = self._resource_get(id=fip_uuid)
        except vnc_exc.NoIdError:
            self._raise_contrail_exception('FloatingIPNotFound',
                                           floatingip_id=fip_uuid)

        return self._fip_obj_to_neutron_dict(fip_obj, fields=fields)

    def resource_list(self, context, filters=None, fields=None):
        # Read in floating ips with either
        # - port(s) as anchor
        # - project(s) as anchor
        # - none as anchor (floating-ip collection)
        ret_list = []

        proj_ids = None
        port_ids = None
        if filters:
            if 'tenant_id' in filters:
                proj_ids = self._validate_project_ids(
                    context, filters['tenant_id'])
            elif 'port_id' in filters:
                port_ids = filters['port_id']
        else:  # no filters
            if not context['is_admin']:
                proj_ids = [self._project_id_neutron_to_vnc(context['tenant'])]

        if port_ids:
            fip_objs = self._resource_list(back_ref_id=port_ids)
        elif proj_ids:
            fip_objs = self._resource_list(back_ref_id=proj_ids)
        else:
            fip_objs = self._resource_list()

        for fip_obj in fip_objs:
            if 'floating_ip_address' in filters:
                if (fip_obj.get_floating_ip_address() not in
                        filters['floating_ip_address']):
                    continue
            ret_list.append(self._fip_obj_to_neutron_dict(fip_obj,
                                                          fields=fields))

        return ret_list

    def resource_count(self, context, filters):
        count = self._resource_count_optimized(filters)
        if count is not None:
            return count

        floatingip_info = self.resource_list(context=context, filters=filters)
        return len(floatingip_info)

class FloatingIpHandler(FloatingIpGetHandler,
                        FloatingIpCreateHandler,
                        FloatingIpDeleteHandler,
                        FloatingIpUpdateHandler):

    pass
<commit_msg>Return 404 errors when FIPs are not found<commit_after># Copyright 2015.  All rights reserved.
#
#    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 uuid

from cfgm_common import exceptions as vnc_exc
from neutron.common import constants as n_constants
from vnc_api import vnc_api

import contrail_res_handler as res_handler
import router_res_handler as router_handler
import vmi_res_handler as vmi_handler

class FloatingIpMixin(object):

    def _neutron_dict_to_fip_obj(self, fip_q, is_admin=False,
                                 tenant_id=None, fip_obj=None):
        if not fip_obj:
            try:
                fip_obj = self._resource_get(id=fip_q['id'])
            except vnc_exc.NoIdError:
                self._raise_contrail_exception('FloatingIPNotFound',
                                               floatingip_id=fip_q['id'])

        vmi_get_handler = vmi_handler.VMInterfaceGetHandler(
            self._vnc_lib)
        port_id = fip_q.get('port_id')
        if port_id:
            vmi_obj = vmi_get_handler.get_vmi_obj(port_id)

            if not is_admin:
                vmi_tenant_id = vmi_get_handler.get_vmi_tenant_id(vmi_obj)
                if vmi_tenant_id != tenant_id:
                    self._raise_contrail_exception('PortNotFound',
                                                   resource='floatingip',
                                                   port_id=port_id)
            fip_obj.set_virtual_machine_interface(vmi_obj)
        else:
            fip_obj.set_virtual_machine_interface_list([])

        if fip_q.get('fixed_ip_address'):
            fip_obj.set_floating_ip_fixed_ip_address(fip_q['fixed_ip_address'])
        else:
            # fixed_ip_address not specified, pick from vmi_obj in create,
            # reset in case of disassociate
            vmi_refs = fip_obj.get_virtual_machine_interface_refs()
            if not vmi_refs:
                fip_obj.set_floating_ip_fixed_ip_address(None)
            else:
                vmi_obj = vmi_get_handler.get_vmi_obj(
                    vmi_refs[0]['uuid'], fields=['instance_ip_back_refs'])

                iip_refs = vmi_obj.get_instance_ip_back_refs()
                if iip_refs:
                    iip_obj = self._vnc_lib.instance_ip_read(
                        id=iip_refs[0]['uuid'])
                    fip_obj.set_floating_ip_fixed_ip_address(
                        iip_obj.get_instance_ip_address())

        return fip_obj

    def _fip_obj_to_neutron_dict(self, fip_obj, fields=None):
        fip_q_dict = {}
        vmi_get_handler = vmi_handler.VMInterfaceGetHandler(
            self._vnc_lib)

        floating_net_id = self._vnc_lib.fq_name_to_id(
            'virtual-network', fip_obj.get_fq_name()[:-2])
        tenant_id = self._project_id_vnc_to_neutron(
            fip_obj.get_project_refs()[0]['uuid'])

        port_id = None
        router_id = None
        vmi_obj = None
        vmi_refs = fip_obj.get_virtual_machine_interface_refs()
        for vmi_ref in vmi_refs or []:
            try:
                vmi_obj = vmi_get_handler.get_vmi_obj(vmi_ref['uuid'])
                port_id = vmi_ref['uuid']
                break
            except vnc_exc.NoIdError:
                pass

        if vmi_obj:
            router_get_handler = router_handler.LogicalRouterGetHandler(
                self._vnc_lib)
            router_id = router_get_handler.get_vmi_obj_router_id(vmi_obj)

        fip_q_dict['id'] = fip_obj.uuid
        fip_q_dict['tenant_id'] = tenant_id
        fip_q_dict['floating_ip_address'] = fip_obj.get_floating_ip_address()
        fip_q_dict['floating_network_id'] = floating_net_id
        fip_q_dict['router_id'] = router_id
        fip_q_dict['port_id'] = port_id
        fip_q_dict['fixed_ip_address'] = (
            fip_obj.get_floating_ip_fixed_ip_address())
        fip_q_dict['status'] = n_constants.PORT_STATUS_ACTIVE

        if fields:
            fip_q_dict = self._filter_res_dict(fip_q_dict, fields)
        return fip_q_dict

class FloatingIpCreateHandler(res_handler.ResourceCreateHandler,
                              FloatingIpMixin):
    resource_create_method = 'floating_ip_create'

    def _create_fip_obj(self, fip_q):
        # TODO() for now create from default pool, later
        # use first available pool on net
        net_id = fip_q['floating_network_id']
        try:
            fq_name = self._vnc_lib.floating_ip_pools_list(
                parent_id=net_id)['floating-ip-pools'][0]['fq_name']
        except IndexError:
            # IndexError could happens when an attempt to
            # retrieve a floating ip pool from a private network.
            msg = "Network %s doesn't provide a floatingip pool" % net_id
            self._raise_contrail_exception('BadRequest',
                                           resource="floatingip", msg=msg)

        fip_pool_obj = self._vnc_lib.floating_ip_pool_read(fq_name=fq_name)
        fip_name = str(uuid.uuid4())
        fip_obj = vnc_api.FloatingIp(fip_name, fip_pool_obj)
        fip_obj.uuid = fip_name

        proj_id = self._project_id_neutron_to_vnc(fip_q['tenant_id'])
        proj_obj = self._project_read(proj_id=proj_id)
        fip_obj.set_project(proj_obj)
        return fip_obj

    def resource_create(self, context, fip_q):
        try:
            fip_obj = self._create_fip_obj(fip_q)
            fip_obj = self._neutron_dict_to_fip_obj(fip_q, context['is_admin'],
                                                    context['tenant'],
                                                    fip_obj=fip_obj)
        except Exception:
            msg = ('Internal error when trying to create floating ip. '
                   'Please be sure the network %s is an external '
                   'network.') % (fip_q['floating_network_id'])
            self._raise_contrail_exception('BadRequest',
                                           resource='floatingip', msg=msg)
        try:
            fip_uuid = self._vnc_lib.floating_ip_create(fip_obj)
        except Exception:
            self._raise_contrail_exception('IpAddressGenerationFailure',
                                           net_id=fip_q['floating_network_id'])
        fip_obj = self._vnc_lib.floating_ip_read(id=fip_uuid)

        return self._fip_obj_to_neutron_dict(fip_obj)

class FloatingIpDeleteHandler(res_handler.ResourceDeleteHandler):
    resource_delete_method = 'floating_ip_delete'

    def resource_delete(self, context, fip_id):
        try:
            self._resource_delete(id=fip_id)
        except vnc_exc.NoIdError:
            self._raise_contrail_exception('FloatingIPNotFound',
                                           floatingip_id=fip_id)

class FloatingIpUpdateHandler(res_handler.ResourceUpdateHandler,
                              FloatingIpMixin):
    resource_update_method = 'floating_ip_update'

    def resource_update(self, context, fip_id, fip_q):
        fip_q['id'] = fip_id
        fip_obj = self._neutron_dict_to_fip_obj(fip_q, context['is_admin'],
                                                context['tenant'])
        self._resource_update(fip_obj)
        return self._fip_obj_to_neutron_dict(fip_obj)

class FloatingIpGetHandler(res_handler.ResourceGetHandler, FloatingIpMixin):
    resource_list_method = 'floating_ips_list'
    resource_get_method = 'floating_ip_read'

    def resource_get(self, context, fip_uuid, fields=None):
        try:
            fip_obj = self._resource_get(id=fip_uuid)
        except vnc_exc.NoIdError:
            self._raise_contrail_exception('FloatingIPNotFound',
                                           floatingip_id=fip_uuid)

        return self._fip_obj_to_neutron_dict(fip_obj, fields=fields)

    def resource_list(self, context, filters=None, fields=None):
        # Read in floating ips with either
        # - port(s) as anchor
        # - project(s) as anchor
        # - none as anchor (floating-ip collection)
        ret_list = []

        proj_ids = None
        port_ids = None
        if filters:
            if 'tenant_id' in filters:
                proj_ids = self._validate_project_ids(
                    context, filters['tenant_id'])
            elif 'port_id' in filters:
                port_ids = filters['port_id']
        else:  # no filters
            if not context['is_admin']:
                proj_ids = [self._project_id_neutron_to_vnc(context['tenant'])]

        if port_ids:
            fip_objs = self._resource_list(back_ref_id=port_ids)
        elif proj_ids:
            fip_objs = self._resource_list(back_ref_id=proj_ids)
        else:
            fip_objs = self._resource_list()

        for fip_obj in fip_objs:
            if 'floating_ip_address' in filters:
                if (fip_obj.get_floating_ip_address() not in
                        filters['floating_ip_address']):
                    continue
            ret_list.append(self._fip_obj_to_neutron_dict(fip_obj,
                                                          fields=fields))

        return ret_list

    def resource_count(self, context, filters):
        count = self._resource_count_optimized(filters)
        if count is not None:
            return count

        floatingip_info = self.resource_list(context=context, filters=filters)
        return len(floatingip_info)

class FloatingIpHandler(FloatingIpGetHandler,
                        FloatingIpCreateHandler,
                        FloatingIpDeleteHandler,
                        FloatingIpUpdateHandler):

    pass
<|endoftext|>


<commit_before>VERSION = (0, 3, 0)

__all__ = ('register',)

class AlreadyRegistered(Exception):
    """
    An attempt was made to register a model for MPTT more than once.
    """
    pass

registry = []

def register(model, parent_attr='parent', left_attr='lft', right_attr='rght',
             tree_id_attr='tree_id', level_attr='level',
             tree_manager_attr='tree', order_insertion_by=None):
    """
    Sets the given model class up for Modified Preorder Tree Traversal.
    """
    try:
        from functools import wraps
    except ImportError:
        from django.utils.functional import wraps # Python 2.3, 2.4 fallback

    from django.db.models import signals as model_signals
    from django.db.models import FieldDoesNotExist, PositiveIntegerField
    from django.utils.translation import ugettext as _

    from mptt import models
    from mptt.signals import pre_save
    from mptt.managers import TreeManager

    if model in registry:
        raise AlreadyRegistered(
            _('The model %s has already been registered.') % model.__name__)
    registry.append(model)

    # Add tree options to the model's Options
    opts = model._meta
    opts.parent_attr = parent_attr
    opts.right_attr = right_attr
    opts.left_attr = left_attr
    opts.tree_id_attr = tree_id_attr
    opts.level_attr = level_attr
    opts.tree_manager_attr = tree_manager_attr
    opts.order_insertion_by = order_insertion_by

    # Add tree fields if they do not exist
    for attr in [left_attr, right_attr, tree_id_attr, level_attr]:
        try:
            opts.get_field(attr)
        except FieldDoesNotExist:
            PositiveIntegerField(
                db_index=True, editable=False).contribute_to_class(model, attr)

    # Add tree methods for model instances
    setattr(model, 'get_ancestors', models.get_ancestors)
    setattr(model, 'get_children', models.get_children)
    setattr(model, 'get_descendants', models.get_descendants)
    setattr(model, 'get_descendant_count', models.get_descendant_count)
    setattr(model, 'get_leafnodes', models.get_leafnodes)
    setattr(model, 'get_next_sibling', models.get_next_sibling)
    setattr(model, 'get_previous_sibling', models.get_previous_sibling)
    setattr(model, 'get_root', models.get_root)
    setattr(model, 'get_siblings', models.get_siblings)
    setattr(model, 'insert_at', models.insert_at)
    setattr(model, 'is_child_node', models.is_child_node)
    setattr(model, 'is_leaf_node', models.is_leaf_node)
    setattr(model, 'is_root_node', models.is_root_node)
    setattr(model, 'move_to', models.move_to)

    # Add a custom tree manager
    TreeManager(parent_attr, left_attr, right_attr, tree_id_attr,
                level_attr).contribute_to_class(model, tree_manager_attr)
    setattr(model, '_tree_manager', getattr(model, tree_manager_attr))

    # Set up signal receiver to manage the tree when instances of the
    # model are about to be saved.
    model_signals.pre_save.connect(pre_save, sender=model)

    # Wrap the model's delete method to manage the tree structure before
    # deletion. This is icky, but the pre_delete signal doesn't currently
    # provide a way to identify which model delete was called on and we
    # only want to manage the tree based on the topmost node which is
    # being deleted.
    def wrap_delete(delete):
        def _wrapped_delete(self):
            opts = self._meta
            tree_width = (getattr(self, opts.right_attr) -
                          getattr(self, opts.left_attr) + 1)
            target_right = getattr(self, opts.right_attr)
            tree_id = getattr(self, opts.tree_id_attr)
            self._tree_manager._close_gap(tree_width, target_right, tree_id)
            delete(self)
        return wraps(delete)(_wrapped_delete)
    model.delete = wrap_delete(model.delete)
<commit_msg>Allow order_insertion_by to be either a string, a list/tuple or None<commit_after>VERSION = (0, 3, 0)

__all__ = ('register',)

class AlreadyRegistered(Exception):
    """
    An attempt was made to register a model for MPTT more than once.
    """
    pass

registry = []

def register(model, parent_attr='parent', left_attr='lft', right_attr='rght',
             tree_id_attr='tree_id', level_attr='level',
             tree_manager_attr='tree', order_insertion_by=None):
    """
    Sets the given model class up for Modified Preorder Tree Traversal.
    """
    try:
        from functools import wraps
    except ImportError:
        from django.utils.functional import wraps # Python 2.3, 2.4 fallback

    from django.db.models import signals as model_signals
    from django.db.models import FieldDoesNotExist, PositiveIntegerField
    from django.utils.translation import ugettext as _

    from mptt import models
    from mptt.signals import pre_save
    from mptt.managers import TreeManager

    if model in registry:
        raise AlreadyRegistered(
            _('The model %s has already been registered.') % model.__name__)
    registry.append(model)

    # Allow order_insertion_by to be either a string, a list/tuple or None
    if order_insertion_by is not None:
        if isinstance(order_insertion_by, basestring):
            order_insertion_by = [order_insertion_by]
        else:
            order_insertion_by = list(order_insertion_by)

    # Add tree options to the model's Options
    opts = model._meta
    opts.parent_attr = parent_attr
    opts.right_attr = right_attr
    opts.left_attr = left_attr
    opts.tree_id_attr = tree_id_attr
    opts.level_attr = level_attr
    opts.tree_manager_attr = tree_manager_attr
    opts.order_insertion_by = order_insertion_by

    # Add tree fields if they do not exist
    for attr in [left_attr, right_attr, tree_id_attr, level_attr]:
        try:
            opts.get_field(attr)
        except FieldDoesNotExist:
            PositiveIntegerField(
                db_index=True, editable=False).contribute_to_class(model, attr)

    # Add tree methods for model instances
    setattr(model, 'get_ancestors', models.get_ancestors)
    setattr(model, 'get_children', models.get_children)
    setattr(model, 'get_descendants', models.get_descendants)
    setattr(model, 'get_descendant_count', models.get_descendant_count)
    setattr(model, 'get_leafnodes', models.get_leafnodes)
    setattr(model, 'get_next_sibling', models.get_next_sibling)
    setattr(model, 'get_previous_sibling', models.get_previous_sibling)
    setattr(model, 'get_root', models.get_root)
    setattr(model, 'get_siblings', models.get_siblings)
    setattr(model, 'insert_at', models.insert_at)
    setattr(model, 'is_child_node', models.is_child_node)
    setattr(model, 'is_leaf_node', models.is_leaf_node)
    setattr(model, 'is_root_node', models.is_root_node)
    setattr(model, 'move_to', models.move_to)

    # Add a custom tree manager
    TreeManager(parent_attr, left_attr, right_attr, tree_id_attr,
                level_attr).contribute_to_class(model, tree_manager_attr)
    setattr(model, '_tree_manager', getattr(model, tree_manager_attr))

    # Set up signal receiver to manage the tree when instances of the
    # model are about to be saved.
    model_signals.pre_save.connect(pre_save, sender=model)

    # Wrap the model's delete method to manage the tree structure before
    # deletion. This is icky, but the pre_delete signal doesn't currently
    # provide a way to identify which model delete was called on and we
    # only want to manage the tree based on the topmost node which is
    # being deleted.
    def wrap_delete(delete):
        def _wrapped_delete(self):
            opts = self._meta
            tree_width = (getattr(self, opts.right_attr) -
                          getattr(self, opts.left_attr) + 1)
            target_right = getattr(self, opts.right_attr)
            tree_id = getattr(self, opts.tree_id_attr)
            self._tree_manager._close_gap(tree_width, target_right, tree_id)
            delete(self)
        return wraps(delete)(_wrapped_delete)
    model.delete = wrap_delete(model.delete)
<|endoftext|>


<commit_before>from idiokit.xmlcore import Element
from abusehelper.core import serialize

import operator
import functools
import re

class RuleError(Exception):
    pass

class _Rule(object):
    @classmethod
    def serialize_register(cls):
        name = "rule-" + cls.__name__.lower()
        serialize.register(cls.dump_rule, cls.load_rule, cls, name)

    def __init__(self, *args, **keys):
        self.arguments = tuple(args), frozenset(keys.items())

    def __eq__(self, other):
        if other.__class__ is not self.__class__:
            return NotImplemented
        return self.arguments == other.arguments

    def __ne__(self, other):
        result = self.__eq__(other)
        if result is NotImplemented:
            return result
        return not result

    def __repr__(self):
        args, keys = self.arguments

        reprs = list()
        reprs.extend(map(repr, args))
        reprs.extend(key + "=" + repr(value) for (key, value) in keys)
        return self.__class__.__name__ + "(" + ", ".join(reprs) + ")"

    def __hash__(self):
        return hash(self.__class__) ^ hash(self.arguments)

class MATCH(_Rule):
    _flags = tuple(sorted("UIXLMS"))

    @classmethod
    def dump_rule(cls, dump, name, rule):
        flags = rule.flags
        if flags is not None:
            flags = "".join(x for x in cls._flags if flags & getattr(re, x, 0))
        return serialize.dump_list(dump, name, [rule.key, rule.pattern, flags])

    @classmethod
    def load_rule(cls, load, element):
        key, pattern, flags = serialize.load_list(load, element)
        if flags is not None:
            flags = reduce(lambda x, y: getattr(re, y, 0) | x, flags, 0)
            pattern = re.compile(pattern, flags)
        return cls(key, pattern)

    def __init__(self, key=None, value=None):
        if value is None:
            filter = None
            pattern = None
            flags = None
        elif isinstance(value, basestring):
            filter = functools.partial(operator.eq, value)
            pattern = value
            flags = None
        else:
            filter = value.search
            pattern = value.pattern
            flags = value.flags

        _Rule.__init__(self, key, value)
        self.key = key
        self.filter = filter
        self.pattern = pattern
        self.flags = flags

    def __call__(self, event):
        if self.key is None:
            return event.contains(filter=self.filter)
        return event.contains(self.key, filter=self.filter)
MATCH.serialize_register()

class NOT(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(dump(rule.child))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 1:
            raise RuleError(element)
        return cls(load(children[0]))

    def __init__(self, child):
        _Rule.__init__(self, child)
        self.child = child

    def __call__(self, *args, **keys):
        return not self.child(*args, **keys)
NOT.serialize_register()

class OR(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if child(*args, **keys):
                return True
        return False
OR.serialize_register()

class AND(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if not child(*args, **keys):
                return False
        return True
AND.serialize_register()

class CONTAINS(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(serialize.dump_list(dump, "keys", rule.keys))
        element.add(serialize.dump_dict(dump, "key-values", rule.key_values))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 2:
            raise RuleError(element)
        keys = set(serialize.load_list(load, children[0]))
        key_values = serialize.load_dict(load, children[1])
        return cls(*keys, **key_values)

    def __init__(self, *keys, **key_values):
        _Rule.__init__(self, *keys, **key_values)
        self.keys = keys
        self.key_values = key_values

    def __call__(self, event):
        for key in self.keys:
            if not event.contains(key):
                return False
        for key, value in self.key_values.items():
            if not event.contains(key, value):
                return False
        return True
CONTAINS.serialize_register()

import struct
from socket import inet_pton, AF_INET, AF_INET6, error

class NETBLOCKError(Exception):
    pass

class NETBLOCK(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name, ip=rule.ip, bits=rule.bits)
        if rule.keys is not None:
            element.add(serialize.dump_list(dump, "keys", rule.keys))
        return element

    @classmethod
    def load_rule(cls, load, element):
        ip = element.get_attr("ip", None)
        bits = element.get_attr("bits", None)
        if None in (ip, bits):
            raise RuleError(element)
        try:
            bits = int(bits)
        except ValueError:
            raise RuleError(element)

        keys = None
        for child in element.children():
            keys = serialize.load_list(load, child)
        return cls(ip, bits, keys)

    _unpack_ipv4 = struct.Struct("!I").unpack
    _unpack_ipv6 = struct.Struct("!QQ").unpack
    def parse_ip(self, ip):
        try:
            packed = inet_pton(AF_INET, ip)
            return 4, self._unpack_ipv4(packed)[0]
        except error:
            try:
                packed = inet_pton(AF_INET6, ip)
                hi, lo = self._unpack_ipv6(packed)
                return 6, ((hi << 64) | lo)
            except error:
                return None

    def __init__(self, ip, bits, keys=None):
        if keys is not None:
            keys = frozenset(keys)

        _Rule.__init__(self, ip, bits, keys)

        self.ip = ip
        self.bits = bits
        self.keys = keys

        parsed = self.parse_ip(ip)
        if parsed is None:
            raise NETBLOCKError("could not parse IP %r" % ip)
        self.version, self.ip_num = parsed
        assert self.version in (4, 6)

        if self.version == 4:
            self.mask = ((1<<32)-1) ^ ((1<<(32-bits))-1)
        elif self.version == 6:
            self.mask = ((1<<128)-1) ^ ((1<<(128-bits))-1)
        self.ip_num &= self.mask

    def __call__(self, event):
        if self.keys is None:
            keys = event.keys()
        else:
            keys = self.keys

        for key in keys:
            for value in event.values(key):
                parsed = self.parse_ip(value)
                if parsed is None:
                    continue
                version, ip_num = parsed
                if version != self.version:
                    continue
                if ip_num & self.mask == self.ip_num:
                    return True
        return False
NETBLOCK.serialize_register()

if __name__ == "__main__":
    import unittest
    from abusehelper.core import events

    class MockEvent(events.Event):
        def __init__(self, **keys):
            events.Event.__init__(self)
            for key, values in keys.iteritems():
                self.update(key, values)

    class MatchTests(unittest.TestCase):
        def test_match(self):
            rule = MATCH()
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["b"]))
            assert rule(MockEvent(x=["y"]))

            rule = MATCH("a")
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["b"]))
            assert not rule(MockEvent(x=["y"]))

            rule = MATCH("a", "b")
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["b"]))
            assert not rule(MockEvent(a=["a"]))

            rule = MATCH("a", re.compile("b"))
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["abba"]))
            assert not rule(MockEvent(a=["aaaa"]))

        def test_eq(self):
            assert MATCH() == MATCH()
            assert MATCH() != MATCH("a")
            assert MATCH("a") == MATCH("a")
            assert MATCH("a") != MATCH("x")
            assert MATCH("a") != MATCH("a", "b")
            assert MATCH("a", "b") == MATCH("a", "b")
            assert MATCH("a", "b") != MATCH("x", "y")
            assert MATCH("a", re.compile("b")) == MATCH("a", re.compile("b"))
            assert MATCH("a", re.compile("b")) != MATCH("x", re.compile("y"))
            assert MATCH("a", re.compile("b")) != MATCH("a", "b")
            assert MATCH("a", re.compile("b", re.I)) != MATCH("a", re.compile("a"))
            assert MATCH("a", re.compile("b", re.I)) != MATCH("a", re.compile("a", re.I))

        def test_serialize(self):
            rule = MATCH()
            assert serialize.load(serialize.dump(rule)) == rule

            rule = MATCH("a", "b")
            assert serialize.load(serialize.dump(rule)) == rule

            rule = MATCH("a", re.compile("b", re.I))
            assert serialize.load(serialize.dump(rule)) == rule

    class NetblockTests(unittest.TestCase):
        def test_match_ipv6(self):
            rule = NETBLOCK("2001:0db8:ac10:fe01::", 32)
            assert rule(MockEvent(ip=["2001:0db8:aaaa:bbbb:cccc::"]))

        def test_non_match_ipv6(self):
            rule = NETBLOCK("2001:0db8:ac10:fe01::", 32)
            assert not rule(MockEvent(ip=["::1"]))

        def test_match_arbitrary_key(self):
            rule = NETBLOCK("1.2.3.4", 24)
            assert rule(MockEvent(somekey=["1.2.3.255"]))

        def test_non_match_arbitrary_key(self):
            rule = NETBLOCK("1.2.3.4", 24)
            assert not rule(MockEvent(somekey=["1.2.4.255"]))

        def test_match_ip_key(self):
            rule = NETBLOCK("1.2.3.4", 24, keys=["ip"])
            assert rule(MockEvent(somekey=["4.5.6.255"], ip=["1.2.3.255"]))

        def test_nonmatch_ip_key(self):
            rule = NETBLOCK("1.2.3.4", 24, keys=["ip"])
            assert not rule(MockEvent(somekey=["1.2.3.4"], ip=["1.2.4.255"]))

        def test_non_match_non_ip_data(self):
            rule = NETBLOCK("1.2.3.4", 24)
            assert not rule(MockEvent(ip=["this is just some data"]))

    unittest.main()
<commit_msg>abusehelper.core.rules: Added some basic tests for rules.CONTAINS.<commit_after>from idiokit.xmlcore import Element
from abusehelper.core import serialize

import operator
import functools
import re

class RuleError(Exception):
    pass

class _Rule(object):
    @classmethod
    def serialize_register(cls):
        name = "rule-" + cls.__name__.lower()
        serialize.register(cls.dump_rule, cls.load_rule, cls, name)

    def __init__(self, *args, **keys):
        self.arguments = tuple(args), frozenset(keys.items())

    def __eq__(self, other):
        if other.__class__ is not self.__class__:
            return NotImplemented
        return self.arguments == other.arguments

    def __ne__(self, other):
        result = self.__eq__(other)
        if result is NotImplemented:
            return result
        return not result

    def __repr__(self):
        args, keys = self.arguments

        reprs = list()
        reprs.extend(map(repr, args))
        reprs.extend(key + "=" + repr(value) for (key, value) in keys)
        return self.__class__.__name__ + "(" + ", ".join(reprs) + ")"

    def __hash__(self):
        return hash(self.__class__) ^ hash(self.arguments)

class MATCH(_Rule):
    _flags = tuple(sorted("UIXLMS"))

    @classmethod
    def dump_rule(cls, dump, name, rule):
        flags = rule.flags
        if flags is not None:
            flags = "".join(x for x in cls._flags if flags & getattr(re, x, 0))
        return serialize.dump_list(dump, name, [rule.key, rule.pattern, flags])

    @classmethod
    def load_rule(cls, load, element):
        key, pattern, flags = serialize.load_list(load, element)
        if flags is not None:
            flags = reduce(lambda x, y: getattr(re, y, 0) | x, flags, 0)
            pattern = re.compile(pattern, flags)
        return cls(key, pattern)

    def __init__(self, key=None, value=None):
        if value is None:
            filter = None
            pattern = None
            flags = None
        elif isinstance(value, basestring):
            filter = functools.partial(operator.eq, value)
            pattern = value
            flags = None
        else:
            filter = value.search
            pattern = value.pattern
            flags = value.flags

        _Rule.__init__(self, key, value)
        self.key = key
        self.filter = filter
        self.pattern = pattern
        self.flags = flags

    def __call__(self, event):
        if self.key is None:
            return event.contains(filter=self.filter)
        return event.contains(self.key, filter=self.filter)
MATCH.serialize_register()

class NOT(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(dump(rule.child))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 1:
            raise RuleError(element)
        return cls(load(children[0]))

    def __init__(self, child):
        _Rule.__init__(self, child)
        self.child = child

    def __call__(self, *args, **keys):
        return not self.child(*args, **keys)
NOT.serialize_register()

class OR(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if child(*args, **keys):
                return True
        return False
OR.serialize_register()

class AND(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        return serialize.dump_list(dump, name, rule.children)

    @classmethod
    def load_rule(cls, load, element):
        children = serialize.load_list(load, element)
        if len(children) < 1:
            raise RuleError(element)
        return cls(*children)

    def __init__(self, first, *rest):
        _Rule.__init__(self, first, *rest)
        self.children = (first,) + rest

    def __call__(self, *args, **keys):
        for child in self.children:
            if not child(*args, **keys):
                return False
        return True
AND.serialize_register()

class CONTAINS(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name)
        element.add(serialize.dump_list(dump, "keys", rule.keys))
        element.add(serialize.dump_dict(dump, "key-values", rule.key_values))
        return element

    @classmethod
    def load_rule(cls, load, element):
        children = list(element.children())
        if len(children) != 2:
            raise RuleError(element)
        keys = set(serialize.load_list(load, children[0]))
        key_values = serialize.load_dict(load, children[1])
        return cls(*keys, **key_values)

    def __init__(self, *keys, **key_values):
        _Rule.__init__(self, *keys, **key_values)
        self.keys = keys
        self.key_values = key_values

    def __call__(self, event):
        for key in self.keys:
            if not event.contains(key):
                return False
        for key, value in self.key_values.items():
            if not event.contains(key, value):
                return False
        return True
CONTAINS.serialize_register()

import struct
from socket import inet_pton, AF_INET, AF_INET6, error

class NETBLOCKError(Exception):
    pass

class NETBLOCK(_Rule):
    @classmethod
    def dump_rule(cls, dump, name, rule):
        element = Element(name, ip=rule.ip, bits=rule.bits)
        if rule.keys is not None:
            element.add(serialize.dump_list(dump, "keys", rule.keys))
        return element

    @classmethod
    def load_rule(cls, load, element):
        ip = element.get_attr("ip", None)
        bits = element.get_attr("bits", None)
        if None in (ip, bits):
            raise RuleError(element)
        try:
            bits = int(bits)
        except ValueError:
            raise RuleError(element)

        keys = None
        for child in element.children():
            keys = serialize.load_list(load, child)
        return cls(ip, bits, keys)

    _unpack_ipv4 = struct.Struct("!I").unpack
    _unpack_ipv6 = struct.Struct("!QQ").unpack
    def parse_ip(self, ip):
        try:
            packed = inet_pton(AF_INET, ip)
            return 4, self._unpack_ipv4(packed)[0]
        except error:
            try:
                packed = inet_pton(AF_INET6, ip)
                hi, lo = self._unpack_ipv6(packed)
                return 6, ((hi << 64) | lo)
            except error:
                return None

    def __init__(self, ip, bits, keys=None):
        if keys is not None:
            keys = frozenset(keys)

        _Rule.__init__(self, ip, bits, keys)

        self.ip = ip
        self.bits = bits
        self.keys = keys

        parsed = self.parse_ip(ip)
        if parsed is None:
            raise NETBLOCKError("could not parse IP %r" % ip)
        self.version, self.ip_num = parsed
        assert self.version in (4, 6)

        if self.version == 4:
            self.mask = ((1<<32)-1) ^ ((1<<(32-bits))-1)
        elif self.version == 6:
            self.mask = ((1<<128)-1) ^ ((1<<(128-bits))-1)
        self.ip_num &= self.mask

    def __call__(self, event):
        if self.keys is None:
            keys = event.keys()
        else:
            keys = self.keys

        for key in keys:
            for value in event.values(key):
                parsed = self.parse_ip(value)
                if parsed is None:
                    continue
                version, ip_num = parsed
                if version != self.version:
                    continue
                if ip_num & self.mask == self.ip_num:
                    return True
        return False
NETBLOCK.serialize_register()

if __name__ == "__main__":
    import unittest
    from abusehelper.core import events

    class MockEvent(events.Event):
        def __init__(self, **keys):
            events.Event.__init__(self)
            for key, values in keys.iteritems():
                self.update(key, values)

    class ContainsTests(unittest.TestCase):
        def test_match(self):
            rule = CONTAINS()
            assert rule(MockEvent())

            rule = CONTAINS("a")
            assert not rule(MockEvent())
            assert not rule(MockEvent(x=["y"]))
            assert rule(MockEvent(a=["b"]))

            rule = CONTAINS(a="b")
            assert not rule(MockEvent(a=["a"]))
            assert rule(MockEvent(a=["b"]))

            rule = CONTAINS(a="b", x="y")
            assert not rule(MockEvent(x=["y"]))
            assert not rule(MockEvent(a=["b"]))
            assert rule(MockEvent(a=["b"], x=["y"]))

        def test_serialize(self):
            rule = CONTAINS()
            assert serialize.load(serialize.dump(rule)) == rule

            rule = CONTAINS("a")
            assert serialize.load(serialize.dump(rule)) == rule

            rule = CONTAINS(a="a")
            assert serialize.load(serialize.dump(rule)) == rule

            rule = CONTAINS("key", a="a", x="y")
            assert serialize.load(serialize.dump(rule)) == rule

    class MatchTests(unittest.TestCase):
        def test_match(self):
            rule = MATCH()
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["b"]))
            assert rule(MockEvent(x=["y"]))

            rule = MATCH("a")
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["b"]))
            assert not rule(MockEvent(x=["y"]))

            rule = MATCH("a", "b")
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["b"]))
            assert not rule(MockEvent(a=["a"]))

            rule = MATCH("a", re.compile("b"))
            assert not rule(MockEvent())
            assert rule(MockEvent(a=["abba"]))
            assert not rule(MockEvent(a=["aaaa"]))

        def test_eq(self):
            assert MATCH() == MATCH()
            assert MATCH() != MATCH("a")
            assert MATCH("a") == MATCH("a")
            assert MATCH("a") != MATCH("x")
            assert MATCH("a") != MATCH("a", "b")
            assert MATCH("a", "b") == MATCH("a", "b")
            assert MATCH("a", "b") != MATCH("x", "y")
            assert MATCH("a", re.compile("b")) == MATCH("a", re.compile("b"))
            assert MATCH("a", re.compile("b")) != MATCH("x", re.compile("y"))
            assert MATCH("a", re.compile("b")) != MATCH("a", "b")
            assert MATCH("a", re.compile("b", re.I)) != MATCH("a", re.compile("a"))
            assert MATCH("a", re.compile("b", re.I)) != MATCH("a", re.compile("a", re.I))

        def test_serialize(self):
            rule = MATCH()
            assert serialize.load(serialize.dump(rule)) == rule

            rule = MATCH("a", "b")
            assert serialize.load(serialize.dump(rule)) == rule

            rule = MATCH("a", re.compile("b", re.I))
            assert serialize.load(serialize.dump(rule)) == rule

    class NetblockTests(unittest.TestCase):
        def test_match_ipv6(self):
            rule = NETBLOCK("2001:0db8:ac10:fe01::", 32)
            assert rule(MockEvent(ip=["2001:0db8:aaaa:bbbb:cccc::"]))

        def test_non_match_ipv6(self):
            rule = NETBLOCK("2001:0db8:ac10:fe01::", 32)
            assert not rule(MockEvent(ip=["::1"]))

        def test_match_arbitrary_key(self):
            rule = NETBLOCK("1.2.3.4", 24)
            assert rule(MockEvent(somekey=["1.2.3.255"]))

        def test_non_match_arbitrary_key(self):
            rule = NETBLOCK("1.2.3.4", 24)
            assert not rule(MockEvent(somekey=["1.2.4.255"]))

        def test_match_ip_key(self):
            rule = NETBLOCK("1.2.3.4", 24, keys=["ip"])
            assert rule(MockEvent(somekey=["4.5.6.255"], ip=["1.2.3.255"]))

        def test_nonmatch_ip_key(self):
            rule = NETBLOCK("1.2.3.4", 24, keys=["ip"])
            assert not rule(MockEvent(somekey=["1.2.3.4"], ip=["1.2.4.255"]))

        def test_non_match_non_ip_data(self):
            rule = NETBLOCK("1.2.3.4", 24)
            assert not rule(MockEvent(ip=["this is just some data"]))

    unittest.main()
<|endoftext|>


<commit_before>"""
Dataset class that wraps the dataset class found in
image_directory_dataset to support models with branched
input layers; allowing different versions of images as
input to those layers.

Based on dev work in the Interactive Pylearn2 notebook.
"""
__authors__ = "Gavin Gray"
__copyright__ = "Copyright 2015 - University of Edinburgh"
__credits__ = ["Gavin Gray"]
__license__ = "MIT"
__maintainer__ = "Gavin Gray"
__email__ = "gavingray1729@gmail.com"

import numpy as np

import neukrill_net.image_directory_dataset

class ParallelIterator(neukrill_net.image_directory_dataset.FlyIterator):
    """
    A simple version of FlyIterator that is able to deal with multiple
    images being returned by the processing function.
    """
    def next(self):
        # return one batch
        if len(self.indices) >= self.batch_size:
            batch_indices = [self.indices.pop(0) for i in range(self.batch_size)]
            # preallocate array
            if len(self.final_shape) == 2:
                Xbatch1 = np.zeros([self.batch_size]+list(self.final_shape)+[1])
                Xbatch2 = np.zeros([self.batch_size]+list(self.final_shape)+[1])
            elif len(self.final_shape) == 3:
                Xbatch1 = np.zeros([self.batch_size]+list(self.final_shape))
                Xbatch2 = np.zeros([self.batch_size]+list(self.final_shape))
            # iterate over indices, applying the dataset's processing function
            for i,j in enumerate(batch_indices):
                import pdb
                pdb.set_trace()
                Xbatch1[i],Xbatch2[i] = [image.reshape(Xbatch1.shape[1:]) for
                        image in self.dataset.fn(self.dataset.X[j])]
            # get the batch for y as well
            ybatch = self.dataset.y[batch_indices,:].astype(np.float32)
            Xbatch1 = Xbatch1.astype(np.float32)
            Xbatch2 = Xbatch2.astype(np.float32)
            return Xbatch1,Xbatch2,ybatch
        else:
            raise StopIteration

class ParallelDataset(neukrill_net.image_directory_dataset.ListDataset):
    """
    Only difference here is that it will use the above Parallel
    iterator instead of the existing iterator.
    """
    def iterator(self, mode=None, batch_size=None, num_batches=None, rng=None,
                        data_specs=None, return_tuple=False):
        """
        Returns iterator object with standard Pythonic interface; iterates
        over the dataset over batches, popping off batches from a shuffled
        list of indices.
        Inputs:
            - mode: 'sequential' or 'shuffled_sequential'.
            - batch_size: required, size of the minibatches produced.
            - num_batches: supply if you want, the dataset will make as many
        as it can if you don't.
            - rng: not used, as above.
            - data_specs: not used, as above
            - return_tuple: not used, as above
        Outputs:
            - instance of FlyIterator, see above.
        """
        if not num_batches:
            # guess that we want to use all of them
            num_batches = int(len(self.X)/batch_size)
        iterator = ParallelIterator(dataset=self, batch_size=batch_size,
                                num_batches=num_batches,
                                final_shape=self.run_settings["final_shape"],
                                rng=self.rng)
        return iterator
<commit_msg>fixed modes<commit_after>"""
Dataset class that wraps the dataset class found in
image_directory_dataset to support models with branched
input layers; allowing different versions of images as
input to those layers.

Based on dev work in the Interactive Pylearn2 notebook.
"""
__authors__ = "Gavin Gray"
__copyright__ = "Copyright 2015 - University of Edinburgh"
__credits__ = ["Gavin Gray"]
__license__ = "MIT"
__maintainer__ = "Gavin Gray"
__email__ = "gavingray1729@gmail.com"

import numpy as np

import neukrill_net.image_directory_dataset

class ParallelIterator(neukrill_net.image_directory_dataset.FlyIterator):
    """
    A simple version of FlyIterator that is able to deal with multiple
    images being returned by the processing function.
    """
    def next(self):
        # return one batch
        if len(self.indices) >= self.batch_size:
            batch_indices = [self.indices.pop(0) for i in range(self.batch_size)]
            # preallocate array
            if len(self.final_shape) == 2:
                Xbatch1 = np.zeros([self.batch_size]+list(self.final_shape)+[1])
                Xbatch2 = np.zeros([self.batch_size]+list(self.final_shape)+[1])
            elif len(self.final_shape) == 3:
                Xbatch1 = np.zeros([self.batch_size]+list(self.final_shape))
                Xbatch2 = np.zeros([self.batch_size]+list(self.final_shape))
            # iterate over indices, applying the dataset's processing function
            for i,j in enumerate(batch_indices):
                import pdb
                pdb.set_trace()
                Xbatch1[i],Xbatch2[i] = [image.reshape(Xbatch1.shape[1:]) for
                        image in self.dataset.fn(self.dataset.X[j])]
            # get the batch for y as well
            ybatch = self.dataset.y[batch_indices,:].astype(np.float32)
            Xbatch1 = Xbatch1.astype(np.float32)
            Xbatch2 = Xbatch2.astype(np.float32)
            return Xbatch1,Xbatch2,ybatch
        else:
            raise StopIteration

class ParallelDataset(neukrill_net.image_directory_dataset.ListDataset):
    """
    Only difference here is that it will use the above Parallel
    iterator instead of the existing iterator.
    """
    def iterator(self, mode=None, batch_size=None, num_batches=None, rng=None,
                        data_specs=None, return_tuple=False):
        """
        Returns iterator object with standard Pythonic interface; iterates
        over the dataset over batches, popping off batches from a shuffled
        list of indices.
        Inputs:
            - mode: 'sequential' or 'shuffled_sequential'.
            - batch_size: required, size of the minibatches produced.
            - num_batches: supply if you want, the dataset will make as many
        as it can if you don't.
            - rng: not used, as above.
            - data_specs: not used, as above
            - return_tuple: not used, as above
        Outputs:
            - instance of FlyIterator, see above.
        """
        if not num_batches:
            # guess that we want to use all of them
            num_batches = int(len(self.X)/batch_size)
        iterator = ParallelIterator(dataset=self, batch_size=batch_size,
                                num_batches=num_batches,
                                final_shape=self.run_settings["final_shape"],
                                rng=self.rng, mode=mode)
        return iterator
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
__author__ = 'study_sun'

import sys
from spider_base.convenient import *
reload(sys)
sys.setdefaultencoding('utf-8')

from enum import Enum, unique
# 数据抹平策略:
# 0.常规:大于100的舍弃,小于0的舍弃
# 1.小于0的算0
# 2.过高的数据(暂定pe>100)抛弃
# 3.小于0的抛弃
# 默认是后两种的组合,注意1和3是不相容的,我也懒得交验
@unique
class StockDataFlatPolicy(Enum):
    RAW = 0
    ZERO_NEGATIVE = 1 << 0
    DISCARD_NEGATIVE = 1 << 1
    DISCARD_LARGE = 1 << 2
    NORMAL = (DISCARD_NEGATIVE | DISCARD_LARGE)

# 数据输出策略,不过有点尴尬的是,如果是中位数,其实就不需要抹平数据了
# 1.算数平均
# 2.中位数
@unique
class StockDataAveragePolicy(Enum):
    MEAN = 0
    MEDIAN = 1

class IndexQuotation:
    pass

# 这个输出当然是输出指数啦,放进来的都是原始数据,看输出要怎么处理
# 一般输入的形式都是
class IndexOutputer(object):

    def __init__(self):
        self.flat_policy = 0

    # 暂时不知道什么策略,先留白吧,pe和pb可能策略也不一样,所以写了两个函数
    def _flat_pe(self, pes, flat_policy = StockDataFlatPolicy.RAW):
        pes = to_container(pes)
        if flat_policy == StockDataFlatPolicy.RAW:
            return pes
        for (index, value) in enumerate(pes):
            to_remove = []
            if value < 0:
                if flat_policy & StockDataFlatPolicy.DISCARD_NEGATIVE > 0:
                    to_remove.append(value)
                elif flat_policy & StockDataFlatPolicy.ZERO_NEGATIVE > 0:
                    pes[index] = 0
            elif value > 100:
                if flat_policy & StockDataFlatPolicy.DISCARD_LARGE > 0:
                    to_remove.append(value)
        pes.remove()

    def _flat_pb(selfself, pbs, flat_policy = StockDataFlatPolicy.RAW):
        pass

    # 纯粹按数据份数来百分比值,可能以后会有按年份分的
    # 数值请尽量多一些,不然我不好10%的取
    def _segment(self, data):
        sumed = []
        for one in data:
            sumed.append(sum(one)/len(one))
        sumed.sort()
        length = len(sumed)
        r = range(0, 10)
        s = [sumed[int(float(length)/10*i)] for i in r]
        s.append(sumed[LAST_ELEMENT_INDEX])
        return s

    # 打印输入时间段内的估值百分比,入参还是一如analysis的结果
    # 结果大致是:
    # 数据时间段xxx天
    # 0%pe为9.2, 10%pe为10.2, 20%pe威武12, ... 100%pe为30
    # 0#pb为xx...
    def print_index_quotations(self, index_quotations, flat_policy=StockDataFlatPolicy.NORMAL, average_policy=StockDataAveragePolicy.MEAN,
                               today_pes=[], today_pbs=[]):
        index_quotations = to_container(index_quotations)
        for index_quotation in index_quotations:
            index_info = index_quotation[0]
            print index_info
            pes = []
            pbs = []
            for index_day_quotation in index_quotation[1]:
                pes.append(index_day_quotation[1])
                pbs.append(index_day_quotation[2])

            pes = sorted(pes)
            pbs = sorted(pbs)
            # 中位数无需进行数据抹平,排个序就好了
            mid_pe = median(pes)
            mid_pb = median(pbs)
            # 此时根据输出策略进行数据抹平,才好算平均数
            self._flat_pe(pes, flat_policy)
            self._flat_pb(pbs, flat_policy)
            segmented_pes = self._segment(pes)
            segmented_pbs = self._segment(pbs)

            print 'pe 百分点位为 ' + ','.join([str(i) for i in segmented_pes])
            print 'pb 百分点位为 ' + ','.join([str(i) for i in segmented_pbs])

            # 最后应该再来个当天处于什么百分位比较好,就以最后一个行情为当天吧

    # 绘图形式输出,先占个位吧
    def draw_index_quotations(self, index_quotations):
        pass

    def xixi(self, s, b, c='hehe', d='zz'):
        print s, b, c, d

if __name__ == '__main__':
    a = IndexOutputer()
    # a.xixi(b=2,s=1,c='zz',d='hehe')
    l = [1,2,3,-1,-3,1,2,-9]
    z = StockDataFlatPolicy.NORMAL.value

    if z & StockDataFlatPolicy.ZERO_NEGATIVE.value > 0:
        x = [max(0, i) for i in l]
    if z & StockDataFlatPolicy.DISCARD_NEGATIVE.value > 0:
        x = [i for i in l if i > 0]
    print x<commit_msg>implement print version index value percent<commit_after># -*- coding: utf-8 -*-
__author__ = 'study_sun'

import sys
from spider_base.convenient import *
reload(sys)
sys.setdefaultencoding('utf-8')

from enum import Enum, unique
# 数据抹平策略:
# 0.常规:大于100的舍弃,小于0的舍弃
# 1.小于0的算0
# 2.过高的数据(暂定pe>100)抛弃
# 3.小于0的抛弃
# 默认是后两种的组合,注意1和3是不相容的,我也懒得交验
@unique
class StockDataFlatPolicy(Enum):
    RAW = 0
    ZERO_NEGATIVE = 1 << 0
    DISCARD_NEGATIVE = 1 << 1
    DISCARD_LARGE = 1 << 2
    NORMAL = (DISCARD_NEGATIVE | DISCARD_LARGE)

# 数据输出策略,不过有点尴尬的是,如果是中位数,其实就不需要抹平数据了
# 1.算数平均
# 2.中位数
@unique
class StockDataAveragePolicy(Enum):
    MEAN = 0
    MEDIAN = 1

class IndexQuotation:
    pass

# 这个输出当然是输出指数啦,放进来的都是原始数据,看输出要怎么处理
# 一般输入的形式都是
class IndexOutputer(object):

    def __init__(self):
        self.flat_policy = 0

    #对数组的数组取中位数,自然不需要什么抹平
    def _median(self, data):
        mids = []
        for one_list in data:
            one_list = sorted(one_list)
            mids.append(median(one_list))
        return mids

    # 暂时不知道什么策略,先留白吧,pe和pb可能策略也不一样,所以写了两个函数
    def _flat(self, pes, flat_policy = StockDataFlatPolicy.RAW.value):
        pes = to_container(pes)
        if flat_policy != StockDataFlatPolicy.RAW.value:
            for (index, value) in enumerate(pes):
                if flat_policy & StockDataFlatPolicy.DISCARD_NEGATIVE.value:
                    value = [i for i in value if i >= 0]
                if flat_policy & StockDataFlatPolicy.DISCARD_LARGE.value:
                    value = [i for i in value if i <= 100]
                if flat_policy & StockDataFlatPolicy.ZERO_NEGATIVE.value:
                    value = [max(i, 0) for i in value]
                pes[index] = value
        flat_pes = []
        for value in pes:
            flat_pes.append(sum(value)/len(value))
        return flat_pes

    def _flat_pb(selfself, pbs, flat_policy = StockDataFlatPolicy.RAW.value):
        pass

    # 纯粹按数据份数来百分比值,可能以后会有按年份分的
    # 数值请尽量多一些,不然我不好10%的取
    def _segment(self, data):
        sumed = sorted(data)
        length = len(sumed)
        r = range(0, 10)
        s = [sumed[int(float(length)/10*i)] for i in r]
        s.append(sumed[LAST_ELEMENT_INDEX])
        return s

    # 打印输入时间段内的估值百分比,入参还是一如analysis的结果
    # 结果大致是:
    # 数据时间段xxx天
    # 0%pe为9.2, 10%pe为10.2, 20%pe威武12, ... 100%pe为30
    # 0#pb为xx...
    def print_index_quotations(self, index_quotations, flat_policy=StockDataFlatPolicy.NORMAL.value, average_policy=StockDataAveragePolicy.MEAN.value,
                               today_pes=[], today_pbs=[]):
        index_quotations = to_container(index_quotations)
        for index_quotation in index_quotations:
            index_info = index_quotation[0]
            print index_info
            pes = []
            pbs = []
            for index_day_quotation in index_quotation[1]:
                pes.append(index_day_quotation[1])
                pbs.append(index_day_quotation[2])
            # 此时的pes和pbs是数组的数组,且未经处理
            # 请注意中位数不是指把他们加起来之后的中位数,而是每天的中位数哦
            mid_pes = self._median(pes)
            mid_pbs = self._median(pbs)
            # 再根据输出策略进行数据抹平,才好算平均数,pb就不抹平了吧
            flat_pes = self._flat(pes, flat_policy)
            flat_pbs = self._flat(pbs, StockDataFlatPolicy.RAW.value)
            mid_segmented_pes = self._segment(mid_pes)
            mid_segmented_pbs = self._segment(mid_pbs)
            flat_segmented_pes = self._segment(flat_pes)
            flat_segmented_pbs = self._segment(flat_pbs)

            print 'pe 平均数百分点位为 ' + ','.join([str(i) for i in flat_segmented_pes])
            print 'pe 中位数百分点位为 ' + ','.join([str(i) for i in mid_segmented_pes])
            print 'pb 平均数百分点位为 ' + ','.join([str(i) for i in flat_segmented_pbs])
            print 'pb 中位数百分点位为 ' + ','.join([str(i) for i in mid_segmented_pbs])
            # 最后应该再来个当天处于什么百分位比较好,就以最后一个行情为当天吧

    # 绘图形式输出,先占个位吧
    def draw_index_quotations(self, index_quotations):
        pass

    def xixi(self, s, b, c='hehe', d='zz'):
        print s, b, c, d

if __name__ == '__main__':
    a = IndexOutputer()<|endoftext|>


<commit_before>"""
ViewOperations manipulate dataviews, typically for the purposes of
visualization. Such operations often apply to SheetViews or
SheetStacks and compose the data together in ways that can be viewed
conveniently, often by creating or manipulating color channels.
"""

import colorsys
import numpy as np
import matplotlib

import param
from param import ParamOverrides

from views import Overlay
from sheetviews import SheetView, SheetStack, SheetLayer, DataGrid
from dataviews import DataLayer, DataStack, Stack, TableView, TableStack
from sheetviews import GridLayout, CoordinateGrid

from options import options, GrayNearest

rgb_to_hsv = np.vectorize(colorsys.rgb_to_hsv)
hsv_to_rgb = np.vectorize(colorsys.hsv_to_rgb)

stack_mapping = {SheetLayer:SheetStack,
                 DataLayer:DataStack,
                 TableView:TableStack}

class ViewOperation(param.ParameterizedFunction):
    """
    A ViewOperation takes one or more views as inputs and processes
    them, returning arbitrary new view objects as output. Individual
    dataviews may be passed in directly while multiple dataviews must
    be passed in as a Stack of the appropriate type. A ViewOperation
    may be used to implement simple dataview manipulations or perform
    complex analysis.

    Internally, ViewOperations operate on the level of individual
    dataviews, processing each layer on an input Stack independently.
    """

    def _process(self, view):
        """
        Process a single input view and output a list of views. When
        multiple views are returned as a list, they will be returned
        to the user as a GridLayout. If a Stack is passed into a
        ViewOperation, the individual layers are processed
        sequentially.
        """
        raise NotImplementedError

    def get_views(self, view, pattern, view_type=SheetView):
        """
        Helper method that return a list of views with labels ending
        with the given pattern and which have the specified type. This
        may be useful to check is a single view satisfies some
        condition or to extract the appropriate views from an Overlay.
        """
        if isinstance(view, Overlay):
            matches = [v for v in view.data if v.label.endswith(pattern)]
        elif isinstance(view, SheetView):
            matches = [view] if view.label.endswith(pattern) else []

        return [match for match in matches if isinstance(match, view_type)]

    def _get_signature(self, view_lists):
        """
        Given a sequence of view lists generated by the _process
        method, find the fundamental view types required to select the
        appropriate Stack type. All lists are examined to ensure that
        the fundamental types are consistent across all items.
        """
        signature = None
        for views in view_lists:
            view_signature = []
            for v in views:
                if type(v) in [CoordinateGrid, DataGrid]:
                    raise NotImplementedError("CoordinateGrid and DataGrid not supported yet")
                matches = [k for k in stack_mapping if issubclass(type(v),k)]
                assert len(matches) == 1
                view_signature.append(matches[0])

            if signature is None:
                signature = view_signature

            if signature != view_signature:
                raise Exception('The ViewOperation is returning inconsistent types.')

        return signature

    def __call__(self, view, **params):
        self.p = ParamOverrides(self, params)

        if not isinstance(view, Stack):
            views = self._process(view)
            if len(views) > 1:
                return GridLayout(views)
            else:
                return views[0]
        else:
            mapped_items = [(k, self._process(el)) for k, el in view.items()]
            signature = self._get_signature(el[1] for el in mapped_items)

            stack_types = [stack_mapping[tp] for tp in signature]
            stacks = [stack_tp(dimensions=view.dimensions,
                               metadata=view.metadata) for stack_tp in stack_types]

            for k, views in mapped_items:
                for ind, v in enumerate(views):
                    stacks[ind][k] = v

            if len(stacks) == 1:
                return stacks[0]
            else:
                return GridLayout(stacks)

class RGBA(ViewOperation):
    """
    Accepts an overlay containing either 3 or 4 layers. The first
    three layers are the R,G, B channels and the last layer (if given)
    is the alpha channel.
    """

    def _process(self, overlay):
        if len(overlay) not in [3, 4]:
            raise Exception("Requires 3 or 4 layers to convert to RGB(A)")
        if not all(isinstance(el, SheetView) for el in overlay.data):
            raise Exception("All layers must be SheetViews to convert"
                            " to RGB(A) format")
        if not all(el.depth == 1 for el in overlay.data):
            raise Exception("All SheetViews must have a depth of one for"
                            " conversion to RGB(A) format")

        return [SheetView(np.dstack([el.data for el in overlay.data]), overlay.bounds,
                          label='RGBA',
                          roi_bounds=overlay.roi_bounds)]

class AlphaOverlay(ViewOperation):
    """
    Accepts an overlay of a SheetView defined with a cmap and converts
    it to an RGBA SheetView. The alpha channel of the result is
    defined by the second layer of the overlay.
    """

    def _process(self, overlay):
        R,G,B,_ = split(cmap2rgb(overlay[0]))
        return [SheetView(RGBA(R*G*B*overlay[1]).data,
                          overlay.bounds,
                          label='AlphaOverlay')]

class HCS(ViewOperation):
    """
    Hue-Confidence-Strength plot.

    Accepts an overlay containing either 2 or 3 layers. The first two
    layers are hue and confidence and the third layer (if available)
    is the strength channel.
    """

    S_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the strength value.""")

    C_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the confidence value.""")

    flipSC = param.Boolean(default=False, doc="""
        Whether to flip the strength and confidence channels""")

    def _process(self, overlay):
        hue = overlay[0]
        confidence = overlay[1]

        strength_data = overlay[2].data if (len(overlay) == 3) else np.ones(hue.shape)

        if hue.shape != confidence.shape:
            raise Exception("Cannot combine plots with different shapes")

        (h,s,v)= (hue.N.data.clip(0.0, 1.0),
                  (confidence.data * self.p.C_multiplier).clip(0.0, 1.0),
                  (strength_data * self.p.S_multiplier).clip(0.0, 1.0))

        if self.p.flipSC:
            (h,s,v) = (h,v,s.clip(0,1.0))

        r, g, b = hsv_to_rgb(h, s, v)
        rgb = np.dstack([r,g,b])
        return [SheetView(rgb, hue.bounds, roi_bounds=overlay.roi_bounds,
                          label=hue.label+' HCS')]

class Colorize(ViewOperation):
    """
    Given a SheetOverlay consisting of a grayscale colormap and a
    second Sheetview with some specified colour map, use the second
    layer to colorize the data of the first layer.

    Currently, colorize only support the 'hsv' color map and is just a
    shortcut to the HCS operation using a constant confidence
    value. Arbitrary colorization will be supported in future.
    """

    def _process(self, overlay):

         if len(overlay) != 2 and overlay[0].mode != 'cmap':
             raise Exception("Can only colorize grayscale overlayed with colour map.")
         if [overlay[0].depth, overlay[1].depth ] != [1,1]:
             raise Exception("Depth one layers required.")
         if overlay[0].shape != overlay[1].shape:
             raise Exception("Shapes don't match.")

         # Needs a general approach which works with any color map
         C = SheetView(np.ones(overlay[1].data.shape),
                       bounds=overlay.bounds)
         hcs = HCS(overlay[1] * C * overlay[0].N)

         return [SheetView(hcs.data, hcs.bounds,
                           roi_bounds=hcs.roi_bounds,
                           label= overlay[0].label+' Colorize')]

class cmap2rgb(ViewOperation):
    """
    Convert SheetViews using colormaps to RGBA mode.  The colormap of
    the style is used, if available. Otherwise, the colormap may be
    forced as a parameter.
    """

    cmap = param.String(default=None, allow_None=True, doc="""
          Force the use of a specific color map. Otherwise, the cmap
          property of the applicable style is used.""")

    def _process(self, sheetview):
        if sheetview.depth != 1:
            raise Exception("Can only apply colour maps to SheetViews with depth of 1.")

        style_cmap = options.style[sheetview][0].get('cmap', None)
        if not any([self.p.cmap, style_cmap]):
            raise Exception("No color map supplied and no cmap in the active style.")

        cmap = matplotlib.cm.get_cmap(style_cmap if self.p.cmap is None else self.p.cmap)
        return [SheetView(cmap(sheetview.data),
                         bounds=sheetview.bounds,
                         cyclic_range=sheetview.cyclic_range,
                         style=sheetview.style,
                         metadata=sheetview.metadata,
                         label = sheetview.label+' RGB')]

class split(ViewOperation):
    """
    Given SheetViews in RGBA mode, return the R,G,B and A channels as
    a GridLayout.
    """
    def _process(self, sheetview):
        if sheetview.mode not in ['rgb','rgba']:
            raise Exception("Can only split SheetViews with a depth of 3 or 4")
        return [SheetView(sheetview.data[:,:,i],
                          bounds=sheetview.bounds,
                          label='RGBA'[i] + ' Channel')
                for i in range(sheetview.depth)]

options.R_Channel = GrayNearest
options.G_Channel = GrayNearest
options.B_Channel = GrayNearest
options.A_Channel = GrayNearest
<commit_msg>Simplified stack type matching in ViewOperation<commit_after>"""
ViewOperations manipulate dataviews, typically for the purposes of
visualization. Such operations often apply to SheetViews or
SheetStacks and compose the data together in ways that can be viewed
conveniently, often by creating or manipulating color channels.
"""

import colorsys
import numpy as np
import matplotlib

import param
from param import ParamOverrides

from views import Overlay
from sheetviews import SheetView, SheetStack, SheetLayer, DataGrid
from dataviews import DataLayer, DataStack, Stack, TableView, TableStack
from sheetviews import GridLayout, CoordinateGrid

from options import options, GrayNearest

rgb_to_hsv = np.vectorize(colorsys.rgb_to_hsv)
hsv_to_rgb = np.vectorize(colorsys.hsv_to_rgb)

stack_mapping = {SheetLayer:SheetStack,
                 DataLayer:DataStack,
                 TableView:TableStack}

class ViewOperation(param.ParameterizedFunction):
    """
    A ViewOperation takes one or more views as inputs and processes
    them, returning arbitrary new view objects as output. Individual
    dataviews may be passed in directly while multiple dataviews must
    be passed in as a Stack of the appropriate type. A ViewOperation
    may be used to implement simple dataview manipulations or perform
    complex analysis.

    Internally, ViewOperations operate on the level of individual
    dataviews, processing each layer on an input Stack independently.
    """

    def _process(self, view):
        """
        Process a single input view and output a list of views. When
        multiple views are returned as a list, they will be returned
        to the user as a GridLayout. If a Stack is passed into a
        ViewOperation, the individual layers are processed
        sequentially.
        """
        raise NotImplementedError

    def get_views(self, view, pattern, view_type=SheetView):
        """
        Helper method that return a list of views with labels ending
        with the given pattern and which have the specified type. This
        may be useful to check is a single view satisfies some
        condition or to extract the appropriate views from an Overlay.
        """
        if isinstance(view, Overlay):
            matches = [v for v in view.data if v.label.endswith(pattern)]
        elif isinstance(view, SheetView):
            matches = [view] if view.label.endswith(pattern) else []

        return [match for match in matches if isinstance(match, view_type)]

    def _get_signature(self, view_lists):
        """
        Given a sequence of view lists generated by the _process
        method, find the fundamental view types required to select the
        appropriate Stack type. All lists are examined to ensure that
        the fundamental types are consistent across all items.
        """
        signature = None
        for views in view_lists:
            view_signature = []
            for v in views:
                if type(v) in [CoordinateGrid, DataGrid]:
                    raise NotImplementedError("CoordinateGrid and DataGrid not supported yet")
                stack_type = [k for k in stack_mapping if issubclass(type(v), k)][0]
                view_signature.append(stack_type)

            if signature is None:
                signature = view_signature

            if signature != view_signature:
                raise Exception('The ViewOperation is returning inconsistent types.')

        return signature

    def __call__(self, view, **params):
        self.p = ParamOverrides(self, params)

        if not isinstance(view, Stack):
            views = self._process(view)
            if len(views) > 1:
                return GridLayout(views)
            else:
                return views[0]
        else:
            mapped_items = [(k, self._process(el)) for k, el in view.items()]
            signature = self._get_signature(el[1] for el in mapped_items)

            stack_types = [stack_mapping[tp] for tp in signature]
            stacks = [stack_tp(dimensions=view.dimensions,
                               metadata=view.metadata) for stack_tp in stack_types]

            for k, views in mapped_items:
                for ind, v in enumerate(views):
                    stacks[ind][k] = v

            if len(stacks) == 1:
                return stacks[0]
            else:
                return GridLayout(stacks)

class RGBA(ViewOperation):
    """
    Accepts an overlay containing either 3 or 4 layers. The first
    three layers are the R,G, B channels and the last layer (if given)
    is the alpha channel.
    """

    def _process(self, overlay):
        if len(overlay) not in [3, 4]:
            raise Exception("Requires 3 or 4 layers to convert to RGB(A)")
        if not all(isinstance(el, SheetView) for el in overlay.data):
            raise Exception("All layers must be SheetViews to convert"
                            " to RGB(A) format")
        if not all(el.depth == 1 for el in overlay.data):
            raise Exception("All SheetViews must have a depth of one for"
                            " conversion to RGB(A) format")

        return [SheetView(np.dstack([el.data for el in overlay.data]), overlay.bounds,
                          label='RGBA',
                          roi_bounds=overlay.roi_bounds)]

class AlphaOverlay(ViewOperation):
    """
    Accepts an overlay of a SheetView defined with a cmap and converts
    it to an RGBA SheetView. The alpha channel of the result is
    defined by the second layer of the overlay.
    """

    def _process(self, overlay):
        R,G,B,_ = split(cmap2rgb(overlay[0]))
        return [SheetView(RGBA(R*G*B*overlay[1]).data,
                          overlay.bounds,
                          label='AlphaOverlay')]

class HCS(ViewOperation):
    """
    Hue-Confidence-Strength plot.

    Accepts an overlay containing either 2 or 3 layers. The first two
    layers are hue and confidence and the third layer (if available)
    is the strength channel.
    """

    S_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the strength value.""")

    C_multiplier = param.Number(default=1.0, bounds=(0.0,None), doc="""
        Multiplier for the confidence value.""")

    flipSC = param.Boolean(default=False, doc="""
        Whether to flip the strength and confidence channels""")

    def _process(self, overlay):
        hue = overlay[0]
        confidence = overlay[1]

        strength_data = overlay[2].data if (len(overlay) == 3) else np.ones(hue.shape)

        if hue.shape != confidence.shape:
            raise Exception("Cannot combine plots with different shapes")

        (h,s,v)= (hue.N.data.clip(0.0, 1.0),
                  (confidence.data * self.p.C_multiplier).clip(0.0, 1.0),
                  (strength_data * self.p.S_multiplier).clip(0.0, 1.0))

        if self.p.flipSC:
            (h,s,v) = (h,v,s.clip(0,1.0))

        r, g, b = hsv_to_rgb(h, s, v)
        rgb = np.dstack([r,g,b])
        return [SheetView(rgb, hue.bounds, roi_bounds=overlay.roi_bounds,
                          label=hue.label+' HCS')]

class Colorize(ViewOperation):
    """
    Given a SheetOverlay consisting of a grayscale colormap and a
    second Sheetview with some specified colour map, use the second
    layer to colorize the data of the first layer.

    Currently, colorize only support the 'hsv' color map and is just a
    shortcut to the HCS operation using a constant confidence
    value. Arbitrary colorization will be supported in future.
    """

    def _process(self, overlay):

         if len(overlay) != 2 and overlay[0].mode != 'cmap':
             raise Exception("Can only colorize grayscale overlayed with colour map.")
         if [overlay[0].depth, overlay[1].depth ] != [1,1]:
             raise Exception("Depth one layers required.")
         if overlay[0].shape != overlay[1].shape:
             raise Exception("Shapes don't match.")

         # Needs a general approach which works with any color map
         C = SheetView(np.ones(overlay[1].data.shape),
                       bounds=overlay.bounds)
         hcs = HCS(overlay[1] * C * overlay[0].N)

         return [SheetView(hcs.data, hcs.bounds,
                           roi_bounds=hcs.roi_bounds,
                           label= overlay[0].label+' Colorize')]

class cmap2rgb(ViewOperation):
    """
    Convert SheetViews using colormaps to RGBA mode.  The colormap of
    the style is used, if available. Otherwise, the colormap may be
    forced as a parameter.
    """

    cmap = param.String(default=None, allow_None=True, doc="""
          Force the use of a specific color map. Otherwise, the cmap
          property of the applicable style is used.""")

    def _process(self, sheetview):
        if sheetview.depth != 1:
            raise Exception("Can only apply colour maps to SheetViews with depth of 1.")

        style_cmap = options.style[sheetview][0].get('cmap', None)
        if not any([self.p.cmap, style_cmap]):
            raise Exception("No color map supplied and no cmap in the active style.")

        cmap = matplotlib.cm.get_cmap(style_cmap if self.p.cmap is None else self.p.cmap)
        return [SheetView(cmap(sheetview.data),
                         bounds=sheetview.bounds,
                         cyclic_range=sheetview.cyclic_range,
                         style=sheetview.style,
                         metadata=sheetview.metadata,
                         label = sheetview.label+' RGB')]

class split(ViewOperation):
    """
    Given SheetViews in RGBA mode, return the R,G,B and A channels as
    a GridLayout.
    """
    def _process(self, sheetview):
        if sheetview.mode not in ['rgb','rgba']:
            raise Exception("Can only split SheetViews with a depth of 3 or 4")
        return [SheetView(sheetview.data[:,:,i],
                          bounds=sheetview.bounds,
                          label='RGBA'[i] + ' Channel')
                for i in range(sheetview.depth)]

options.R_Channel = GrayNearest
options.G_Channel = GrayNearest
options.B_Channel = GrayNearest
options.A_Channel = GrayNearest
<|endoftext|>


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

import rospy
import subprocess
import signal
import os

import actionlib
from strands_executive_msgs.abstract_task_server import AbstractTaskServer
import aaf_logging.msg

class LoggingScriptServer(AbstractTaskServer):

    def __init__(self, name):
        rospy.loginfo("Starting node: %s" % name)
        #self._action_name = name
        #self._as = actionlib.SimpleActionServer(self._action_name, aaf_logging.msg.RunNodeAction, execute_cb = self.execute_cb)
        #self._as.start()
        super(LoggingScriptServer, self).__init__(
            name=name,
            action_type=aaf_logging.msg.RunNodeAction,
            interruptible=True
        )
        rospy.loginfo('Server is up')

    def execute(self, goal):
        # decide whether recording should be started or stopped
        if goal.command == "start":
            rospy.loginfo('now the logging should start')
            command = "rosrun aaf_logging run_aaf_logger.bash"
            self.p = subprocess.Popen(command, stdin=subprocess.PIPE, preexec_fn=os.setsid, shell=True)
            rospy.loginfo(self.p.pid)

            # check if the goal is preempted
            rate = rospy.Rate(1.0)
            while not rospy.is_shutdown() and not self.server.is_preempt_requested(): # and self.p.poll() is None:
                rate.sleep()

            rospy.loginfo('Logging is preempted')
            os.killpg(os.getpgid(self.p.pid), signal.SIGINT)
            #if rospy.is_shutdown():
            #    self.server.set_preempted()
            #    return

        elif goal.command == "stop":
            rospy.loginfo('now the logging should stop')
            rospy.loginfo(self.p.pid)
            os.killpg(os.getpgid(self.p.pid), signal.SIGINT)
            rospy.loginfo("I'm done")

        else:
            rospy.loginfo('goal.command is not valid')

        if not self.server.is_preempt_requested():
            self.server.set_succeeded()
        else:
            self.server.set_preempted()

    def create(self, req):
        task = super(LoggingScriptServer, self).create(req)
        if task.start_node_id == "":
            task.start_node_id = "ChargingPoint"
        if task.end_node_id == "":
            task.end_node_id = task.start_node_id
        if task.max_duration.secs == 0:
            task.max_duration.secs = 60 # Default execution time: 1min
        if task.priority == 0:
            task.priority = 5 # Always start and stop loggin with high priority.
        return task

if __name__ == '__main__':
    rospy.init_node('logging_script_server')
    LoggingScriptServer(rospy.get_name())
    rospy.spin()
<commit_msg>Added some checks so that it doesn't crash if you try to call invalid commands<commit_after>#! /usr/bin/env python

import rospy
import subprocess
import signal
import os

import actionlib
from strands_executive_msgs.abstract_task_server import AbstractTaskServer
import aaf_logging.msg

class LoggingScriptServer(AbstractTaskServer):

    def __init__(self, name):
        rospy.loginfo("Starting node: %s" % name)
        #self._action_name = name
        #self._as = actionlib.SimpleActionServer(self._action_name, aaf_logging.msg.RunNodeAction, execute_cb = self.execute_cb)
        #self._as.start()
        super(LoggingScriptServer, self).__init__(
            name=name,
            action_type=aaf_logging.msg.RunNodeAction,
            interruptible=True
        )
        rospy.loginfo('Server is up')
        self.running = False

    def execute(self, goal):
        # decide whether recording should be started or stopped
        if goal.command == "start" and not self.running:
            rospy.loginfo('now the logging should start')
            command = "rosrun aaf_logging run_aaf_logger.bash"
            self.p = subprocess.Popen(command, stdin=subprocess.PIPE, preexec_fn=os.setsid, shell=True)
            self.running = True
            rospy.loginfo(self.p.pid)

            # check if the goal is preempted
            rate = rospy.Rate(1.0)
            while not rospy.is_shutdown() and not self.server.is_preempt_requested(): # and self.p.poll() is None:
                rate.sleep()

            rospy.loginfo('Logging is preempted')
            os.killpg(os.getpgid(self.p.pid), signal.SIGINT)
            #if rospy.is_shutdown():
            #    self.server.set_preempted()
            #    return
            self.running = False

        elif goal.command == "stop" and self.running:
            rospy.loginfo('now the logging should stop')
            rospy.loginfo(self.p.pid)
            os.killpg(os.getpgid(self.p.pid), signal.SIGINT)
            rospy.loginfo("I'm done")
            self.running = False

        else:
            rospy.loginfo('goal.command is not valid')

        if not self.server.is_preempt_requested():
            self.server.set_succeeded()
        else:
            self.server.set_preempted()

    def create(self, req):
        task = super(LoggingScriptServer, self).create(req)
        if task.start_node_id == "":
            task.start_node_id = "ChargingPoint"
        if task.end_node_id == "":
            task.end_node_id = task.start_node_id
        if task.max_duration.secs == 0:
            task.max_duration.secs = 60 # Default execution time: 1min
        if task.priority == 0:
            task.priority = 5 # Always start and stop loggin with high priority.
        return task

if __name__ == '__main__':
    rospy.init_node('logging_script_server')
    LoggingScriptServer(rospy.get_name())
    rospy.spin()
<|endoftext|>


<commit_before>from data_structures.stack import Stack

class Node(object):
    def __init__(self, key=None, val=None):
        self.count = 0
        if key is not None:
            self.count = 1
        self.elems = [(key, val), (None, None), (None, None), ]

    def __repr__(self):
        """For printing out the heap and its nodes
        It's here to save me typing during debugging"""
        result = "["
        for i in range(self.count):
            result += '%s, ' % str(self.elems[i])
        return '%s]' % result

    def add_to_node(self, key, val):
        for i in range(3):
            if self.elems[i][0] is None or self.elems[i][0] > key:
                self.elems.pop()
                self.elems.insert(i, (key, val))
                self.count += 1
                break

    def del_from_node(self, idx):
        pass

    def has(self, key):
        pass

    def split_node(self):
        pass

class BTree(object):
    def __init__(self):
        self.root = Node()

    def search(self, key):
        pass

    def insert(self, key, val):
        pass

    def delete(self, key):
        pass
<commit_msg>Add del_from_node<commit_after>from data_structures.stack import Stack

class Node(object):
    def __init__(self, key=None, val=None):
        self.count = 0
        if key is not None:
            self.count = 1
        self.elems = [(key, val), (None, None), (None, None), ]

    def __repr__(self):
        """For printing out the heap and its nodes
        It's here to save me typing during debugging"""
        result = "["
        for i in range(self.count):
            result += '%s, ' % str(self.elems[i])
        return '%s]' % result

    def add_to_node(self, key, val):
        for i in range(3):
            if self.elems[i][0] is None or self.elems[i][0] > key:
                self.elems.pop()
                self.elems.insert(i, (key, val))
                self.count += 1
                break

    def del_from_node(self, idx):
        self.elems.pop(idx)
        self.count -= 1
        self.elems.append((None, None))

    def has(self, key):
        pass

    def split_node(self):
        pass

class BTree(object):
    def __init__(self):
        self.root = Node()

    def search(self, key):
        pass

    def insert(self, key, val):
        pass

    def delete(self, key):
        pass
<|endoftext|>