<commit_before>import functools
import json
import uuid

import jsonschema
import panglery

from cardboard import core
from cardboard.util import ANY

class NotAuthorized(Exception):
    """
    The user was not authorized to call an API method.

    """

    def __str__(self):
        return "Authorization failed or was not provided."

class User(object):
    """
    A user is an enraged animal sitting across the network.

    If you truly wish to interact with one, be my guest.

    """

    def event_triggered(self, **event):
        pass

    def select(self, choices, how_many=1, duplicates=False):
        return self.protocol.callRemote(
            Select, choices=choices, how_many=how_many, duplicates=duplicates
        )

    def select_cards(
        self, zone=None, match=ANY, how_many=1, duplicates=False, bad=True
    ):
        cards = [card for card in zone if match(card)]
        return self.protocol.callRemote(
            SelectCards, cards=cards, how_many=how_many,
            duplicates=duplicates, bad=bad,
        )

    def select_players(
        self, match=ANY, how_many=1, duplicates=False, bad=True
    ):
        players = [player for player in self.game.players if match(player)]
        return self.protocol.callRemote(
            SelectPlayers, players=players, how_many=how_many,
            duplicates=duplicates, bad=bad,
        )

    def select_combined(
        self,
        zone=None,
        match_cards=ANY,
        how_many_cards=1,
        duplicate_cards=False,
        match_players=ANY,
        how_many_players=1,
        duplicate_players=False,
        bad=True,
    ):
        cards = [card for card in zone if match_cards(card)]
        players = [play for play in self.game.players if match_players(play)]
        return self.protocol.callRemote(
            SelectCombined, cards=cards, players=players,
            how_many_cards=how_many_cards, duplicate_cards=duplicate_cards,
            how_many_players=how_many_players,
            duplicate_players=duplicate_players, bad=bad,
        )

    def select_range(self, start, stop, how_many=1, duplicates=False):
        return self.protocol.callRemote(
            SelectRange, start=start, stop=stop,
            how_many=how_many, duplicates=duplicates,
        )

def document_schema(schema, type, indent=0):
    """
    Document a schema for a JSON RPC object of the given type.

    :argument schema: a JSON Schema
    :argument type: the type of object that the schema validates
    :type type: "request", "response" or "notification"
    :argument indent: an indent level to move the schema string over
    :type indent: integer
    :rtype: a str in the Sphinx markup language

    """

    if type not in {"request", "response", "notification"}:
        raise ValueError(type)

    ind = (indent + 4) * " "
    schema = ind + json.dumps(schema, indent=4).replace("\n", "\n" + ind)
    return indent * " " + ".. {}::\n\n{}".format(type, schema)

def exposed(request_schema, response_schema, validate=jsonschema.validate):
    """
    Document and validate an exposed API method.

    :argument request_schema: a schema to use to validate the incoming request
    :argument response_schema: a schema used only to document the response
                               object that will be sent back

    """

    def _expose(fn):
        if not fn.__doc__:
            raise ValueError(
                "Get off your lazy ass and write some documentation!"
            )
        else:
            for indent, c in enumerate(fn.__doc__.lstrip("\n")):
                if c != " ":
                    break

            req_doc = document_schema(request_schema, "request", indent)
            res_doc = document_schema(response_schema, "response", indent)
            fn.__doc__  = "\n\n".join([fn.__doc__.rstrip(), req_doc, res_doc])

        # TODO: Document schema["properties"] too

        @functools.wraps(fn)
        def exposed_fn(self, **request):
            validate(request, request_schema)
            return fn(self, **request)

        exposed_fn.request_schema = request_schema
        exposed_fn.response_schema = response_schema
        return exposed_fn
    return _expose

class APIController(object):
    def __init__(self):
        self.drafts = []
        self.games = []
        self.players = []

    def lookup_method(self, name):
        """
        Lookup the appropriate API method with the given name.

        """

        return getattr(self, "api_" + name.replace(".", "_"))

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "auth" : {"type" : "string", "required" : True},
             "gameID" : {"type" : "integer", "required" : True},
             "playerID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
    )
    def api_concede(self, auth, gameID, playerID):
        """
        Concede from the current game.

        """

        expected_auth, player = self.players[gameID][playerID]
        if auth != expected_auth:
            raise NotAuthorized()
        player.concede()
        return {}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "verbose" : {"type" : "boolean", "default" : False},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "started" : {"type" : "boolean", "required" : True},
             "teams" : {"required" : True},
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_info(self, gameID, verbose=False):
        """
        Retrieve game info about a specific game.

        """

        # XXX: verbose
        game = self.games[gameID]
        return {
            "gameID" : gameID, "started" : game.started, "teams" : game.teams,
        }

    @exposed(
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
        {
         "type" : "array",
        },
    )
    def api_Game_list(self, verbose=False):
        """
        List the currently open games.

        """

        info = self.lookup_method("Game.info")
        return [
            info(gameID=i, verbose=verbose) for i in xrange(len(self.games))
        ]

    @exposed(
        {
         "type" : "object",
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {"gameID" : {"type" : "integer", "required" : True}}
        },
    )
    def api_Game_create(self):
        """
        Create a new game.

        """

        self.games.append(core.Game(panglery.Pangler()))
        self.players.append([])
        return {"gameID" : len(self.games) - 1}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_start(self, gameID):
        """
        Start a game.

        """

        # XXX: Unknown game, not authorized
        self.games[gameID].start()
        return {}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "name" : {"type" : "string", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
             "auth" : {"type" : "string", "required" : True},
             "playerID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_join(self, gameID, name):
        """
        Join a currently open game.

        """

        # XXX: Can't join a started game, can't join twice, library
        game, players = self.games[gameID], self.players[gameID]
        player = game.add_player(library=[], user=User(), name=name)
        auth = uuid.uuid4().bytes
        players.append((auth, player))
        return {"playerID" : len(players) - 1, "auth" : auth}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_end(self, gameID):
        """
        End the current game.

        """

        self.games[gameID].end()
        return {}

    @exposed({}, {})
    def api_Draft_info(self):
        """
        Retrieve game info about a specific draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_list(self):
        """
        List the currently open drafts.

        """

        pass

    @exposed({}, {})
    def api_Draft_create(self):
        """
        Create a new draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_start(self):
        """
        Start a draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_join(self):
        """
        Join a currently open draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_end(self):
        """
        End the current draft.

        """

        pass

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "playerID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
             "name" : {"type" : "string", "required" : True},
             "dead" : {"type" : "boolean", "required" : True},
             "handSize" : {"type" : "integer", "required" : True},
             "life" : {"type" : "integer", "required" : True},
             "poison" : {
                 "type" : "integer", "required" : True,
                 "minimum" : 0, "maximum" : 10,
             },
         },
         "additionalProperties" : False,
        },
    )
    def api_Player_info(self, gameID, playerID):
        """
        Retrieve info about a given player.

        """

        # XXX: No such player
        player = self.players[gameID][playerID][1]
        return {
            "name" : player.name, "handSize" : player.hand_size,
            "life" : player.life, "poison" : player.poison,
            "dead" : player.dead,
        }
<commit_msg>Minor: Sort schema keys.<commit_after>import functools
import json
import uuid

import jsonschema
import panglery

from cardboard import core
from cardboard.util import ANY

class NotAuthorized(Exception):
    """
    The user was not authorized to call an API method.

    """

    def __str__(self):
        return "Authorization failed or was not provided."

class User(object):
    """
    A user is an enraged animal sitting across the network.

    If you truly wish to interact with one, be my guest.

    """

    def event_triggered(self, **event):
        pass

    def select(self, choices, how_many=1, duplicates=False):
        return self.protocol.callRemote(
            Select, choices=choices, how_many=how_many, duplicates=duplicates
        )

    def select_cards(
        self, zone=None, match=ANY, how_many=1, duplicates=False, bad=True
    ):
        cards = [card for card in zone if match(card)]
        return self.protocol.callRemote(
            SelectCards, cards=cards, how_many=how_many,
            duplicates=duplicates, bad=bad,
        )

    def select_players(
        self, match=ANY, how_many=1, duplicates=False, bad=True
    ):
        players = [player for player in self.game.players if match(player)]
        return self.protocol.callRemote(
            SelectPlayers, players=players, how_many=how_many,
            duplicates=duplicates, bad=bad,
        )

    def select_combined(
        self,
        zone=None,
        match_cards=ANY,
        how_many_cards=1,
        duplicate_cards=False,
        match_players=ANY,
        how_many_players=1,
        duplicate_players=False,
        bad=True,
    ):
        cards = [card for card in zone if match_cards(card)]
        players = [play for play in self.game.players if match_players(play)]
        return self.protocol.callRemote(
            SelectCombined, cards=cards, players=players,
            how_many_cards=how_many_cards, duplicate_cards=duplicate_cards,
            how_many_players=how_many_players,
            duplicate_players=duplicate_players, bad=bad,
        )

    def select_range(self, start, stop, how_many=1, duplicates=False):
        return self.protocol.callRemote(
            SelectRange, start=start, stop=stop,
            how_many=how_many, duplicates=duplicates,
        )

def document_schema(schema, type, indent=0):
    """
    Document a schema for a JSON RPC object of the given type.

    :argument schema: a JSON Schema
    :argument type: the type of object that the schema validates
    :type type: "request", "response" or "notification"
    :argument indent: an indent level to move the schema string over
    :type indent: integer
    :rtype: a str in the Sphinx markup language

    """

    if type not in {"request", "response", "notification"}:
        raise ValueError(type)

    ind = (indent + 4) * " "
    schema = ind + json.dumps(schema, sort_keys=True, indent=4)
    schema = schema.replace("\n", "\n" + ind)
    return indent * " " + ".. {}::\n\n{}".format(type, schema)

def exposed(request_schema, response_schema, validate=jsonschema.validate):
    """
    Document and validate an exposed API method.

    :argument request_schema: a schema to use to validate the incoming request
    :argument response_schema: a schema used only to document the response
                               object that will be sent back

    """

    def _expose(fn):
        if not fn.__doc__:
            raise ValueError(
                "Get off your lazy ass and write some documentation!"
            )
        else:
            for indent, c in enumerate(fn.__doc__.lstrip("\n")):
                if c != " ":
                    break

            req_doc = document_schema(request_schema, "request", indent)
            res_doc = document_schema(response_schema, "response", indent)
            fn.__doc__  = "\n\n".join([fn.__doc__.rstrip(), req_doc, res_doc])

        # TODO: Document schema["properties"] too

        @functools.wraps(fn)
        def exposed_fn(self, **request):
            validate(request, request_schema)
            return fn(self, **request)

        exposed_fn.request_schema = request_schema
        exposed_fn.response_schema = response_schema
        return exposed_fn
    return _expose

class APIController(object):
    def __init__(self):
        self.drafts = []
        self.games = []
        self.players = []

    def lookup_method(self, name):
        """
        Lookup the appropriate API method with the given name.

        """

        return getattr(self, "api_" + name.replace(".", "_"))

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "auth" : {"type" : "string", "required" : True},
             "gameID" : {"type" : "integer", "required" : True},
             "playerID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
    )
    def api_concede(self, auth, gameID, playerID):
        """
        Concede from the current game.

        """

        expected_auth, player = self.players[gameID][playerID]
        if auth != expected_auth:
            raise NotAuthorized()
        player.concede()
        return {}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "verbose" : {"type" : "boolean", "default" : False},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "started" : {"type" : "boolean", "required" : True},
             "teams" : {"required" : True},
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_info(self, gameID, verbose=False):
        """
        Retrieve game info about a specific game.

        """

        # XXX: verbose
        game = self.games[gameID]
        return {
            "gameID" : gameID, "started" : game.started, "teams" : game.teams,
        }

    @exposed(
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
        {
         "type" : "array",
        },
    )
    def api_Game_list(self, verbose=False):
        """
        List the currently open games.

        """

        info = self.lookup_method("Game.info")
        return [
            info(gameID=i, verbose=verbose) for i in xrange(len(self.games))
        ]

    @exposed(
        {
         "type" : "object",
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {"gameID" : {"type" : "integer", "required" : True}}
        },
    )
    def api_Game_create(self):
        """
        Create a new game.

        """

        self.games.append(core.Game(panglery.Pangler()))
        self.players.append([])
        return {"gameID" : len(self.games) - 1}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_start(self, gameID):
        """
        Start a game.

        """

        # XXX: Unknown game, not authorized
        self.games[gameID].start()
        return {}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "name" : {"type" : "string", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
             "auth" : {"type" : "string", "required" : True},
             "playerID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_join(self, gameID, name):
        """
        Join a currently open game.

        """

        # XXX: Can't join a started game, can't join twice, library
        game, players = self.games[gameID], self.players[gameID]
        player = game.add_player(library=[], user=User(), name=name)
        auth = uuid.uuid4().bytes
        players.append((auth, player))
        return {"playerID" : len(players) - 1, "auth" : auth}

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
         },
         "additionalProperties" : False,
        },
    )
    def api_Game_end(self, gameID):
        """
        End the current game.

        """

        self.games[gameID].end()
        return {}

    @exposed({}, {})
    def api_Draft_info(self):
        """
        Retrieve game info about a specific draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_list(self):
        """
        List the currently open drafts.

        """

        pass

    @exposed({}, {})
    def api_Draft_create(self):
        """
        Create a new draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_start(self):
        """
        Start a draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_join(self):
        """
        Join a currently open draft.

        """

        pass

    @exposed({}, {})
    def api_Draft_end(self):
        """
        End the current draft.

        """

        pass

    @exposed(
        {
         "type" : "object",
         "properties" : {
             "gameID" : {"type" : "integer", "required" : True},
             "playerID" : {"type" : "integer", "required" : True},
         },
         "additionalProperties" : False,
        },
        {
         "type" : "object",
         "properties" : {
             "name" : {"type" : "string", "required" : True},
             "dead" : {"type" : "boolean", "required" : True},
             "handSize" : {"type" : "integer", "required" : True},
             "life" : {"type" : "integer", "required" : True},
             "poison" : {
                 "type" : "integer", "required" : True,
                 "minimum" : 0, "maximum" : 10,
             },
         },
         "additionalProperties" : False,
        },
    )
    def api_Player_info(self, gameID, playerID):
        """
        Retrieve info about a given player.

        """

        # XXX: No such player
        player = self.players[gameID][playerID][1]
        return {
            "name" : player.name, "handSize" : player.hand_size,
            "life" : player.life, "poison" : player.poison,
            "dead" : player.dead,
        }
<|endoftext|>


<commit_before>#
# -*- coding: utf-8 -*-
#
# This file is part of reclass (http://github.com/madduck/reclass)
#
# Copyright © 2007–13 martin f. krafft <madduck@madduck.net>
# Released under the terms of the Artistic Licence 2.0
#

from output import OutputLoader
from storage import StorageBackendLoader

def get_data(storage_type, nodes_uri, classes_uri, applications_postfix, node):
    storage_class = StorageBackendLoader(storage_type).load()
    storage = storage_class(nodes_uri, classes_uri, applications_postfix)
    if node is False:
        ret = storage.inventory()
    else:
        ret = storage.nodeinfo(node)

    return ret

def output(data, fmt, pretty_print=False):
    output_class = OutputLoader(fmt).load()
    outputter = output_class()
    return outputter.dump(data, pretty_print=pretty_print)
<commit_msg>Allow node to be None to trigger inventory<commit_after>#
# -*- coding: utf-8 -*-
#
# This file is part of reclass (http://github.com/madduck/reclass)
#
# Copyright © 2007–13 martin f. krafft <madduck@madduck.net>
# Released under the terms of the Artistic Licence 2.0
#

from output import OutputLoader
from storage import StorageBackendLoader

def get_data(storage_type, nodes_uri, classes_uri, applications_postfix, node):
    storage_class = StorageBackendLoader(storage_type).load()
    storage = storage_class(nodes_uri, classes_uri, applications_postfix)
    if not node:
        ret = storage.inventory()
    else:
        ret = storage.nodeinfo(node)

    return ret

def output(data, fmt, pretty_print=False):
    output_class = OutputLoader(fmt).load()
    outputter = output_class()
    return outputter.dump(data, pretty_print=pretty_print)
<|endoftext|>


<commit_before># Copyright (c) 2016 Cisco Systems, Inc.
# 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.

# This is a simple Flask application that provides REST APIs by which
# compute and network services can communicate, plus a REST API for
# debugging using a CLI client.

# Note that it does *NOT* at this point have a persistent database, so
# restarting this process will make Gluon forget about every port it's
# learned, which will not do your system much good (the data is in the
# global 'backends' and 'ports' objects).  This is for simplicity of
# demonstration; we have a second codebase already defined that is
# written to OpenStack endpoint principles and includes its ORM, so
# that work was not repeated here where the aim was to get the APIs
# worked out.  The two codebases will merge in the future.

from flask import Flask
from flask_restful import Api
from flask_restful import reqparse
from flask_restful import Resource
import os
import sys
import vpp

from neutron.agent.linux import bridge_lib
from neutron.agent.linux import ip_lib
from neutron.common import constants as n_const
from networking_vpp import config_opts
from oslo_config import cfg
from oslo_log import log as logging

######################################################################

# This mirrors functionality in Neutron so that we're creating a name
# that Neutron can find for its agents.

DEV_NAME_PREFIX = n_const.TAP_DEVICE_PREFIX

def get_tap_name(uuid):
    return n_const.TAP_DEVICE_PREFIX + uuid[0:11]

# This mirrors functionality in Nova so that we're creating a vhostuser
# name that it will be able to locate

VHOSTUSER_DIR = '/tmp'

def get_vhostuser_name(uuid):
    return os.path.join(VHOSTUSER_DIR, uuid)

######################################################################

class VPPForwarder(object):

    def __init__(self,
                 flat_network_if=None,
                 vlan_trunk_if=None,
                 vxlan_src_addr=None,
                 vxlan_bcast_addr=None,
                 vxlan_vrf=None):
        self.vpp = vpp.VPPInterface()
        # This is the flat network interface for providing FLAT networking
        self.flat_if = flat_network_if
        # This is the trunk interface for VLAN networking
        self.trunk_if = vlan_trunk_if

        # This is the address we'll use if we plan on broadcasting
        # vxlan packets
        self.vxlan_bcast_addr = vxlan_bcast_addr
        self.vxlan_src_addr = vxlan_src_addr
        self.vxlan_vrf = vxlan_vrf

        # Used as a unique number for bridge IDs
        self.next_bridge_id = 5678

        self.networks = {}      # vlan: bridge index
        self.interfaces = {}    # uuid: if idx

        for (ifname, f) in self.vpp.get_interfaces():
            # Clean up interfaces from previous runs
            # TODO(ijw) can't easily SPOT VLAN subifs to delete
            if ifname.startswith('tap-'):
                # all VPP tap interfaces are of this form
                self.vpp.delete_tap(f.sw_if_index)
            elif ifname.startswith('VirtualEthernet'):
                # all VPP vhostuser interfaces are of this form
                self.vpp.delete_vhostuser(f.sw_if_index)

        trunk_ifstruct = self.vpp.get_interface(self.trunk_if) if self.trunk_if else None
        flat_ifstruct = self.vpp.get_interface(self.flat_if) if self.flat_if else None
	    if trunk_ifstruct is None and flat_ifstruct is None:
		  raise Exception("Could not find a VPP uplink interface:%s-%s" % (self.trunk_if, self.flat_if))
        if trunk_ifstruct is not None:
            self.trunk_ifidx = trunk_ifstruct.sw_if_index
            # This interface is not automatically up just because
            # we've started and we need to ensure it is before
            # proceeding.

            # TODO(ijw): when we start up in a recovery mode we may
            # want to check the local VPP config and bring it up when
            # confirmed.
            self.vpp.ifup(self.trunk_ifidx)
        if flat_ifstruct is not None:
            self.flat_ifidx = flat_ifstruct.sw_if_index
            self.vpp.ifup(self.flat_ifidx)

    # This, here, is us creating a VLAN backed network
    def network_on_host(self, type, seg_id):
        if (type, seg_id) not in self.networks:
            # TODO(ijw): bridge domains have no distinguishing marks.
            # VPP needs to allow us to name or label them so that we
            # can find them when we restart
            if type == 'flat':
                if_upstream = self.flat_ifidx
                app.logger.debug('Adding upstream interface:%s to bridge for flat networking' % self.flat_if)
            elif type == 'vlan':
                # TODO(ijw): this VLAN subinterface may already exist, and
                # may even be in another bridge domain already (see
                # above).
                if_upstream = self.vpp.create_vlan_subif(self.trunk_ifidx,
                                                         seg_id)
                app.logger.debug('Adding upstream trunk interface:%s.%s \
                to bridge for vlan networking' % (self.trunk_if,seg_id))
            elif type == 'vxlan':
                if_upstream = \
                    self.vpp.create_srcrep_vxlan_subif(self, self.vxlan_vrf,
                                                       self.vxlan_src_addr,
                                                       self.vxlan_bcast_addr,
                                                       seg_id)
            else:
                raise Exception('type %s not supported', type)

            self.vpp.ifup(if_upstream)

            # May not remain this way but we use the VLAN ID as the
            # bridge ID; TODO(ijw): bridge ID can already exist, we
            # should check till we find a free one
            id = self.next_bridge_id
            self.next_bridge_id += 1

            self.vpp.create_bridge_domain(id)

            self.vpp.add_to_bridge(id, if_upstream)

            self.networks[(type, seg_id)] = id

        return self.networks[(type, seg_id)]

    ########################################
    # stolen from LB driver
    def _bridge_exists_and_ensure_up(self, bridge_name):
        """Check if the bridge exists and make sure it is up."""
        br = ip_lib.IPDevice(bridge_name)
        br.set_log_fail_as_error(False)
        try:
            # If the device doesn't exist this will throw a RuntimeError
            br.link.set_up()
        except RuntimeError:
            return False
        return True

    def ensure_bridge(self, bridge_name):
        """Create a bridge unless it already exists."""
        # _bridge_exists_and_ensure_up instead of device_exists is used here
        # because there are cases where the bridge exists but it's not UP,
        # for example:
        # 1) A greenthread was executing this function and had not yet executed
        # "ip link set bridge_name up" before eventlet switched to this
        # thread running the same function
        # 2) The Nova VIF driver was running concurrently and had just created
        #    the bridge, but had not yet put it UP
        if not self._bridge_exists_and_ensure_up(bridge_name):
            bridge_device = bridge_lib.BridgeDevice.addbr(bridge_name)
            if bridge_device.setfd(0):
                return
            if bridge_device.disable_stp():
                return
            if bridge_device.disable_ipv6():
                return
            if bridge_device.link.set_up():
                return
        else:
            bridge_device = bridge_lib.BridgeDevice(bridge_name)

        # TODO(ijw): should be checking this all succeeded

    # end theft
    ########################################

    def create_interface_on_host(self, if_type, uuid, mac):
        if uuid in self.interfaces:
            app.logger.debug('port %s repeat binding request - ignored' % uuid)
        else:
            app.logger.debug('binding port %s as type %s' % (uuid, if_type))

            # TODO(ijw): naming not obviously consistent with
            # Neutron's naming
            name = uuid[0:11]
            bridge_name = 'br-' + name
            tap_name = 'tap' + name

            if if_type == 'maketap' or if_type == 'plugtap':
                if if_type == 'maketap':
                    iface = self.vpp.create_tap(tap_name, mac)
                    props = {'bind_type': 'maketap', 'name': tap_name}
                else:
                    int_tap_name = 'vpp' + name

                    props = {'bind_type': 'plugtap',
                             'bridge_name': bridge_name,
                             'ext_tap_name': tap_name,
                             'int_tap_name': int_tap_name
                    }

                    iface = self.vpp.create_tap(int_tap_name, mac)

                    # TODO(ijw): someone somewhere ought to be sorting
                    # the MTUs out

                    br = self.ensure_bridge(bridge_name)
                    # This is the external TAP device that has already been precreated
                    br.addif(tap_name)
                    # This is the device that we just created with VPP
                    br.addif(int_tap_name)

            elif if_type == 'vhostuser':
                path = get_vhostuser_name(uuid)
                iface = self.vpp.create_vhostuser(path, mac)
                props = {'bind_type': 'vhostuser', 'path': uuid}
            else:
                raise Exception('unsupported interface type')

            self.interfaces[uuid] = (iface, props)
        return self.interfaces[uuid]

    def bind_interface_on_host(self, if_type, uuid, mac, net_type, seg_id):
        net_br_idx = self.network_on_host(net_type, seg_id)
        (iface, props) = self.create_interface_on_host(if_type, uuid, mac)

        self.vpp.ifup(iface)
        self.vpp.add_to_bridge(net_br_idx, iface)

        return props

    def unbind_interface_on_host(self, uuid):
        if uuid not in self.interfaces:
            app.logger.debug("unknown port %s unbinding request - ignored" % uuid)
        else:

            iface_idx, props = self.interfaces[uuid]

            app.logger.debug("unbinding port %s, recorded as type %s" % (uuid, props['bind_type']))

            # We no longer need this interface.  Specifically if it's
            # a vhostuser interface it's annoying to have it around
            # because the VM's memory (hugepages) will not be
            # released.  So, here, we destroy it.

            # TODO(ijw): I'm assuming deleting an interface in a
            # bridge domain works...
            if props['bind_type'] == 'vhostuser':
                self.vpp.delete_vhostuser(iface_idx)
            elif props['bind_type'] in ['maketap', 'plugtap']:
                self.vpp.delete_tap(iface_idx)
                if props['bind_type'] == 'plugtap':
                    bridge = bridge_lib.BridgeDevice(bridge_name)
                    if bridge_device.exists():
                        # These may fail, don't care much
                        try:
                            bridge.delif(props['int_tap_name'])
                            bridge.delif(props['ext_tap_name'])
                            bridge.delbr(props['bridge_name'])
                        except:
                            pass
            else:
                app.logger.error('Unknown port type %s during unbind' % props['bind_type'])

######################################################################

class PortBind(Resource):
    bind_args = reqparse.RequestParser()
    bind_args.add_argument('mac_address', type=str, required=True)
    bind_args.add_argument('mtu', type=str, required=True)
    bind_args.add_argument('segmentation_id', type=int, required=True)
    bind_args.add_argument('network_type', type=str, required=True)
    bind_args.add_argument('host', type=str, required=True)

    def put(self, id):
        global vppf

        args = self.bind_args.parse_args()
        app.logger.debug('on host %s, binding %s %d to mac %s id %s'
                         % (args['host'],
                            args['network_type'],
                            args['segmentation_id'],
                            args['mac_address'], id))
        vppf.bind_interface_on_host('vhostuser',
                                    id,
                                    args['mac_address'],
                                    args['network_type'],
                                    args['segmentation_id'])

class PortUnbind(Resource):

    def __init(self, *args, **kwargs):
        super('PortBind', self).__init__(*args, **kwargs)

    def put(self, id):
        global vppf

        vppf.unbind_interface_on_host(id)

LOG = logging.getLogger('vpp-agent')

# Basic Flask RESTful app setup
app = Flask('vpp-agent')

def main():
    app.debug = True

#    logger = logging.getLogger('werkzeug')
#    logger.setLevel(logging.INFO)

    # Basic log config
    app.logger.debug('Debug logging enabled')
    # TODO(ijw) port etc. should probably be configurable.

    cfg.CONF(sys.argv[1:])
    global vppf
    vppf = VPPForwarder(flat_network_if=cfg.CONF.ml2_vpp.flat_network_if,
                        vlan_trunk_if=cfg.CONF.ml2_vpp.vlan_trunk_if,
                        vxlan_src_addr=cfg.CONF.ml2_vpp.vxlan_src_addr,
                        vxlan_bcast_addr=cfg.CONF.ml2_vpp.vxlan_bcast_addr,
                        vxlan_vrf=cfg.CONF.ml2_vpp.vxlan_vrf)

    api = Api(app)

    api.add_resource(PortBind, '/ports/<id>/bind')
    api.add_resource(PortUnbind, '/ports/<id>/unbind')

    app.run(host='0.0.0.0',port=2704)

if __name__ == '__main__':

    main()
<commit_msg>add flat networking<commit_after># Copyright (c) 2016 Cisco Systems, Inc.
# 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.

# This is a simple Flask application that provides REST APIs by which
# compute and network services can communicate, plus a REST API for
# debugging using a CLI client.

# Note that it does *NOT* at this point have a persistent database, so
# restarting this process will make Gluon forget about every port it's
# learned, which will not do your system much good (the data is in the
# global 'backends' and 'ports' objects).  This is for simplicity of
# demonstration; we have a second codebase already defined that is
# written to OpenStack endpoint principles and includes its ORM, so
# that work was not repeated here where the aim was to get the APIs
# worked out.  The two codebases will merge in the future.

from flask import Flask
from flask_restful import Api
from flask_restful import reqparse
from flask_restful import Resource
import os
import sys
import vpp

from neutron.agent.linux import bridge_lib
from neutron.agent.linux import ip_lib
from neutron.common import constants as n_const
from networking_vpp import config_opts
from oslo_config import cfg
from oslo_log import log as logging

######################################################################

# This mirrors functionality in Neutron so that we're creating a name
# that Neutron can find for its agents.

DEV_NAME_PREFIX = n_const.TAP_DEVICE_PREFIX

def get_tap_name(uuid):
    return n_const.TAP_DEVICE_PREFIX + uuid[0:11]

# This mirrors functionality in Nova so that we're creating a vhostuser
# name that it will be able to locate

VHOSTUSER_DIR = '/tmp'

def get_vhostuser_name(uuid):
    return os.path.join(VHOSTUSER_DIR, uuid)

######################################################################

class VPPForwarder(object):

    def __init__(self,
                 flat_network_if=None,
                 vlan_trunk_if=None,
                 vxlan_src_addr=None,
                 vxlan_bcast_addr=None,
                 vxlan_vrf=None):
        self.vpp = vpp.VPPInterface()
        # This is the flat network interface for providing FLAT networking
        self.flat_if = flat_network_if
        # This is the trunk interface for VLAN networking
        self.trunk_if = vlan_trunk_if

        # This is the address we'll use if we plan on broadcasting
        # vxlan packets
        self.vxlan_bcast_addr = vxlan_bcast_addr
        self.vxlan_src_addr = vxlan_src_addr
        self.vxlan_vrf = vxlan_vrf

        # Used as a unique number for bridge IDs
        self.next_bridge_id = 5678

        self.networks = {}      # vlan: bridge index
        self.interfaces = {}    # uuid: if idx

        for (ifname, f) in self.vpp.get_interfaces():
            # Clean up interfaces from previous runs
            # TODO(ijw) can't easily SPOT VLAN subifs to delete
            if ifname.startswith('tap-'):
                # all VPP tap interfaces are of this form
                self.vpp.delete_tap(f.sw_if_index)
            elif ifname.startswith('VirtualEthernet'):
                # all VPP vhostuser interfaces are of this form
                self.vpp.delete_vhostuser(f.sw_if_index)

        trunk_ifstruct = self.vpp.get_interface(self.trunk_if) if self.trunk_if else None
        flat_ifstruct = self.vpp.get_interface(self.flat_if) if self.flat_if else None
	    if trunk_ifstruct is None and flat_ifstruct is None:
		    raise Exception("Could not find a VPP uplink interface:%s-%s" % (self.trunk_if, self.flat_if))
        if trunk_ifstruct is not None:
            self.trunk_ifidx = trunk_ifstruct.sw_if_index
            # This interface is not automatically up just because
            # we've started and we need to ensure it is before
            # proceeding.

            # TODO(ijw): when we start up in a recovery mode we may
            # want to check the local VPP config and bring it up when
            # confirmed.
            self.vpp.ifup(self.trunk_ifidx)
        if flat_ifstruct is not None:
            self.flat_ifidx = flat_ifstruct.sw_if_index
            self.vpp.ifup(self.flat_ifidx)

    # This, here, is us creating a VLAN backed network
    def network_on_host(self, type, seg_id):
        if (type, seg_id) not in self.networks:
            # TODO(ijw): bridge domains have no distinguishing marks.
            # VPP needs to allow us to name or label them so that we
            # can find them when we restart
            if type == 'flat':
                if_upstream = self.flat_ifidx
                app.logger.debug('Adding upstream interface:%s to bridge for flat networking' % self.flat_if)
            elif type == 'vlan':
                # TODO(ijw): this VLAN subinterface may already exist, and
                # may even be in another bridge domain already (see
                # above).
                if_upstream = self.vpp.create_vlan_subif(self.trunk_ifidx,
                                                         seg_id)
                app.logger.debug('Adding upstream trunk interface:%s.%s \
                to bridge for vlan networking' % (self.trunk_if,seg_id))
            elif type == 'vxlan':
                if_upstream = \
                    self.vpp.create_srcrep_vxlan_subif(self, self.vxlan_vrf,
                                                       self.vxlan_src_addr,
                                                       self.vxlan_bcast_addr,
                                                       seg_id)
            else:
                raise Exception('type %s not supported', type)

            self.vpp.ifup(if_upstream)

            # May not remain this way but we use the VLAN ID as the
            # bridge ID; TODO(ijw): bridge ID can already exist, we
            # should check till we find a free one
            id = self.next_bridge_id
            self.next_bridge_id += 1

            self.vpp.create_bridge_domain(id)

            self.vpp.add_to_bridge(id, if_upstream)

            self.networks[(type, seg_id)] = id

        return self.networks[(type, seg_id)]

    ########################################
    # stolen from LB driver
    def _bridge_exists_and_ensure_up(self, bridge_name):
        """Check if the bridge exists and make sure it is up."""
        br = ip_lib.IPDevice(bridge_name)
        br.set_log_fail_as_error(False)
        try:
            # If the device doesn't exist this will throw a RuntimeError
            br.link.set_up()
        except RuntimeError:
            return False
        return True

    def ensure_bridge(self, bridge_name):
        """Create a bridge unless it already exists."""
        # _bridge_exists_and_ensure_up instead of device_exists is used here
        # because there are cases where the bridge exists but it's not UP,
        # for example:
        # 1) A greenthread was executing this function and had not yet executed
        # "ip link set bridge_name up" before eventlet switched to this
        # thread running the same function
        # 2) The Nova VIF driver was running concurrently and had just created
        #    the bridge, but had not yet put it UP
        if not self._bridge_exists_and_ensure_up(bridge_name):
            bridge_device = bridge_lib.BridgeDevice.addbr(bridge_name)
            if bridge_device.setfd(0):
                return
            if bridge_device.disable_stp():
                return
            if bridge_device.disable_ipv6():
                return
            if bridge_device.link.set_up():
                return
        else:
            bridge_device = bridge_lib.BridgeDevice(bridge_name)

        # TODO(ijw): should be checking this all succeeded

    # end theft
    ########################################

    def create_interface_on_host(self, if_type, uuid, mac):
        if uuid in self.interfaces:
            app.logger.debug('port %s repeat binding request - ignored' % uuid)
        else:
            app.logger.debug('binding port %s as type %s' % (uuid, if_type))

            # TODO(ijw): naming not obviously consistent with
            # Neutron's naming
            name = uuid[0:11]
            bridge_name = 'br-' + name
            tap_name = 'tap' + name

            if if_type == 'maketap' or if_type == 'plugtap':
                if if_type == 'maketap':
                    iface = self.vpp.create_tap(tap_name, mac)
                    props = {'bind_type': 'maketap', 'name': tap_name}
                else:
                    int_tap_name = 'vpp' + name

                    props = {'bind_type': 'plugtap',
                             'bridge_name': bridge_name,
                             'ext_tap_name': tap_name,
                             'int_tap_name': int_tap_name
                    }

                    iface = self.vpp.create_tap(int_tap_name, mac)

                    # TODO(ijw): someone somewhere ought to be sorting
                    # the MTUs out

                    br = self.ensure_bridge(bridge_name)
                    # This is the external TAP device that has already been precreated
                    br.addif(tap_name)
                    # This is the device that we just created with VPP
                    br.addif(int_tap_name)

            elif if_type == 'vhostuser':
                path = get_vhostuser_name(uuid)
                iface = self.vpp.create_vhostuser(path, mac)
                props = {'bind_type': 'vhostuser', 'path': uuid}
            else:
                raise Exception('unsupported interface type')

            self.interfaces[uuid] = (iface, props)
        return self.interfaces[uuid]

    def bind_interface_on_host(self, if_type, uuid, mac, net_type, seg_id):
        net_br_idx = self.network_on_host(net_type, seg_id)
        (iface, props) = self.create_interface_on_host(if_type, uuid, mac)

        self.vpp.ifup(iface)
        self.vpp.add_to_bridge(net_br_idx, iface)

        return props

    def unbind_interface_on_host(self, uuid):
        if uuid not in self.interfaces:
            app.logger.debug("unknown port %s unbinding request - ignored" % uuid)
        else:

            iface_idx, props = self.interfaces[uuid]

            app.logger.debug("unbinding port %s, recorded as type %s" % (uuid, props['bind_type']))

            # We no longer need this interface.  Specifically if it's
            # a vhostuser interface it's annoying to have it around
            # because the VM's memory (hugepages) will not be
            # released.  So, here, we destroy it.

            # TODO(ijw): I'm assuming deleting an interface in a
            # bridge domain works...
            if props['bind_type'] == 'vhostuser':
                self.vpp.delete_vhostuser(iface_idx)
            elif props['bind_type'] in ['maketap', 'plugtap']:
                self.vpp.delete_tap(iface_idx)
                if props['bind_type'] == 'plugtap':
                    bridge = bridge_lib.BridgeDevice(bridge_name)
                    if bridge_device.exists():
                        # These may fail, don't care much
                        try:
                            bridge.delif(props['int_tap_name'])
                            bridge.delif(props['ext_tap_name'])
                            bridge.delbr(props['bridge_name'])
                        except:
                            pass
            else:
                app.logger.error('Unknown port type %s during unbind' % props['bind_type'])

######################################################################

class PortBind(Resource):
    bind_args = reqparse.RequestParser()
    bind_args.add_argument('mac_address', type=str, required=True)
    bind_args.add_argument('mtu', type=str, required=True)
    bind_args.add_argument('segmentation_id', type=int, required=True)
    bind_args.add_argument('network_type', type=str, required=True)
    bind_args.add_argument('host', type=str, required=True)

    def put(self, id):
        global vppf

        args = self.bind_args.parse_args()
        app.logger.debug('on host %s, binding %s %d to mac %s id %s'
                         % (args['host'],
                            args['network_type'],
                            args['segmentation_id'],
                            args['mac_address'], id))
        vppf.bind_interface_on_host('vhostuser',
                                    id,
                                    args['mac_address'],
                                    args['network_type'],
                                    args['segmentation_id'])

class PortUnbind(Resource):

    def __init(self, *args, **kwargs):
        super('PortBind', self).__init__(*args, **kwargs)

    def put(self, id):
        global vppf

        vppf.unbind_interface_on_host(id)

LOG = logging.getLogger('vpp-agent')

# Basic Flask RESTful app setup
app = Flask('vpp-agent')

def main():
    app.debug = True

#    logger = logging.getLogger('werkzeug')
#    logger.setLevel(logging.INFO)

    # Basic log config
    app.logger.debug('Debug logging enabled')
    # TODO(ijw) port etc. should probably be configurable.

    cfg.CONF(sys.argv[1:])
    global vppf
    vppf = VPPForwarder(flat_network_if=cfg.CONF.ml2_vpp.flat_network_if,
                        vlan_trunk_if=cfg.CONF.ml2_vpp.vlan_trunk_if,
                        vxlan_src_addr=cfg.CONF.ml2_vpp.vxlan_src_addr,
                        vxlan_bcast_addr=cfg.CONF.ml2_vpp.vxlan_bcast_addr,
                        vxlan_vrf=cfg.CONF.ml2_vpp.vxlan_vrf)

    api = Api(app)

    api.add_resource(PortBind, '/ports/<id>/bind')
    api.add_resource(PortUnbind, '/ports/<id>/unbind')

    app.run(host='0.0.0.0',port=2704)

if __name__ == '__main__':

    main()
<|endoftext|>


<commit_before>from copy import deepcopy
import numpy as np
from six import string_types
from sklearn import decomposition as sk_decomp
from sklearn import cluster as sk_cluster
from sklearn.metrics import pairwise_distances
from os.path import exists, join
from os import makedirs
from nibabel import nifti1
from neurosynth.analysis import meta

class Clusterable(object):

    '''
    Args:
        dataset: The Dataset instance to extract data from.
        mask: A mask defining the voxels to cluster. Can be a filename,
            nibabel image, or numpy array (see Masker.mask() for details).
        features (str or list): Optional string or list of strings specifying
            any feature names to use for study selection. E.g., passing
            ['emotion', 'reward'] would retain for analysis only those studies
            associated with the features emotion or reward at a frequency
            greater than feature_threshold.
        feature_threshold (float): The threshold to use when selecting studies
            on the basis of features.
        min_voxels (int): Minimum number of active voxels a study
            must report in order to be retained in the dataset. By default,
            all studies are used.
        min_studies (int): Minimum number of studies a voxel must be
            active in in order to be retained in analysis. By default, all
            voxels are used.
    '''

    def __init__(self, dataset, mask=None, features=None,
                 feature_threshold=None, min_voxels=None, min_studies=None):

        self.dataset = dataset
        self.masker = deepcopy(dataset.masker)

        # Condition study inclusion on specific features
        if features is not None:
            ids = dataset.get_studies(features=features,
                                      frequency_threshold=feature_threshold)
            data = dataset.get_image_data(ids, dense=False)
        else:
            data = dataset.image_table.data

        # Trim data based on minimum number of voxels or studies
        if min_studies is not None:
            av = self.masker.unmask(
                data.sum(1) >= min_studies, output='vector')
            self.masker.add({'voxels': av})

        if min_voxels is not None:
            data = data[:, np.array(data.sum(0) >= min_voxels).squeeze()]

        if mask is not None:
            self.masker.add({'roi': mask})

        self.data = data[self.masker.get_mask(['voxels', 'roi']), :].toarray()

    def transform(self, transformer, transpose=False):
        ''' Apply a transformation to the Clusterable instance. Accepts any
        scikit-learn-style class that implements a fit_transform() method. '''
        data = self.data.T if transpose else self.data
        data = transformer.fit_transform(data)
        self.data = data.T if transpose else data
        return self

def magic(dataset, method='coactivation', roi_mask=None,
          coactivation_mask=None, features=None, feature_threshold=0.05,
          min_voxels_per_study=None, min_studies_per_voxel=None,
          reduce_reference='pca', n_components=100,
          distance_metric='correlation', clustering_algorithm='kmeans',
          n_clusters=5, clustering_kwargs={}, output_dir=None, filename=None,
          coactivation_images=False, coactivation_threshold=0.1):
    ''' Execute a full clustering analysis pipeline.
    Args:
        dataset: a Dataset instance to extract all data from.
        method (str): the overall clustering approach to use. Valid options:
            'coactivation' (default): Clusters voxel within the ROI mask based
                on shared pattern of coactivation with the rest of the brain.
            'studies': Treat each study as a feature in an n-dimensional space.
                I.e., voxels will be assigned to the same cluster if they tend
                to be co-reported in similar studies.
        roi_mask: A string, nibabel image, or numpy array providing an
            inclusion mask of voxels to cluster. If None, the default mask
            in the Dataset instance is used (typically, all in-brain voxels).
        coactivation_mask: If method='coactivation', this mask defines the
            voxels to use when generating the pairwise distance matrix. For
            example, if a PFC mask is passed, all voxels in the roi_mask will
            be clustered based on how similar their patterns of coactivation
            with PFC voxels are. Can be a str, nibabel image, or numpy array.
        features (str or list): Optional string or list of strings specifying
            any feature names to use for study selection. E.g., passing
            ['emotion', 'reward'] would retain for analysis only those studies
            associated with the features emotion or reward at a frequency
            greater than feature_threshold.
        feature_threshold (float): The threshold to use when selecting studies
            on the basis of features.
        min_voxels_per_study (int): Minimum number of active voxels a study
            must report in order to be retained in the dataset. By default,
            all studies are used.
        min_studies_per_voxel (int): Minimum number of studies a voxel must be
            active in in order to be retained in analysis. By default, all
            voxels are used.
        reduce_reference (str, scikit-learn object or None): The dimensionality
            reduction algorithm to apply to the feature space prior to the
            computation of pairwise distances. If a string is passed (either
            'pca' or 'ica'), n_components must be specified. If None, no
            dimensionality reduction will be applied. Otherwise, must be a
            scikit-learn-style object that exposes a transform() method.
        n_components (int): Number of components to extract during the
            dimensionality reduction step. Only used if reduce_reference is
            a string.
        distance_metric (str): The distance metric to use when computing
            pairwise distances on the to-be-clustered voxels. Can be any of the
            metrics supported by sklearn.metrics.pairwise_distances.
        clustering_algorithm (str or scikit-learn object): the clustering
            algorithm to use. If a string, must be one of 'kmeans' or 'minik'.
            Otherwise, any sklearn class that exposes a fit_predict() method.
        n_clusters (int): If clustering_algorithm is a string, the number of
            clusters to extract.
        clustering_kwargs (dict): Additional keywords to pass to the clustering
            object.
        output_dir (str): The directory to write results to. If None (default),
            returns the cluster label image rather than saving to disk.
        filename (str): Name of cluster label image file. Defaults to
            cluster_labels_k{k}.nii.gz, where k is the number of clusters.
        coactivation_images (bool): If True, saves a meta-analytic coactivation
            map for every ROI in the resulting cluster map.
        coactivation_threshold (float or int): If coactivation_images is True,
            this is the threshold used to define whether or not a study is
            considered to activation within a cluster ROI. Integer values are
            interpreted as minimum number of voxels within the ROI; floats
            are interpreted as the proportion of voxels. Defaults to 0.1 (i.e.,
            10% of all voxels within ROI must be active).
    '''

    roi = Clusterable(dataset, roi_mask, min_voxels=min_voxels_per_study,
                      min_studies=min_studies_per_voxel, features=features,
                      feature_threshold=feature_threshold)

    if method == 'coactivation':
        reference = Clusterable(dataset, coactivation_mask,
                                min_voxels=min_voxels_per_study,
                                min_studies=min_studies_per_voxel,
                                features=features,
                                feature_threshold=feature_threshold)
    elif method == 'features':
        reference = deepcopy(roi)
        feature_data = dataset.feature_table.data
        n_studies = len(feature_data)
        reference.data = reference.data.dot(feature_data.values) / n_studies
    elif method == 'studies':
        reference = roi

    if reduce_reference is not None:

        # For non-coactivation-based approaches, transpose the data matrix
        transpose = (method == 'coactivation')

        if isinstance(reduce_reference, string_types):

            # Number of components can't exceed feature count or cluster count
            n_feat = len(reference.data.T if transpose else reference.data)
            n_components = min(n_components, n_feat)
            if n_clusters is not None:
                n_components = min(n_components, n_clusters)

            reduce_reference = {
                'pca': sk_decomp.PCA,
                'ica': sk_decomp.FastICA
            }[reduce_reference](n_components)
        reference = reference.transform(reduce_reference, transpose=transpose)

    if method == 'coactivation':
        distances = pairwise_distances(roi.data, reference.data,
                                       metric=distance_metric)
    else:
        distances = reference.data

    # TODO: add additional clustering methods
    if isinstance(clustering_algorithm, string_types):
        clustering_algorithm = {
            'kmeans': sk_cluster.KMeans,
            'minik': sk_cluster.MiniBatchKMeans
        }[clustering_algorithm](n_clusters, **clustering_kwargs)

    labels = clustering_algorithm.fit_predict(distances) + 1.

    header = roi.masker.get_header()
    header['cal_max'] = labels.max()
    header['cal_min'] = labels.min()
    voxel_labels = roi.masker.unmask(labels)
    img = nifti1.Nifti1Image(voxel_labels, None, header)

    if output_dir is not None:
        if not exists(output_dir):
            makedirs(output_dir)
        if filename is None:
            filename = 'cluster_labels_k%d.nii.gz' % n_clusters
        outfile = join(output_dir, filename)
        img.to_filename(outfile)

        # Write coactivation images
        if coactivation_images:
            for l in np.unique(voxel_labels):
                roi_mask = np.copy(voxel_labels)
                roi_mask[roi_mask != l] = 0
                ids = dataset.get_studies(
                    mask=roi_mask, activation_threshold=coactivation_threshold)
                ma = meta.MetaAnalysis(dataset, ids)
                ma.save_results(output_dir=join(output_dir, 'coactivation'),
                                prefix='cluster_%d_coactivation' % l)
    else:
        return img
<commit_msg>update docstring and fix minimum constraint<commit_after>from copy import deepcopy
import numpy as np
from six import string_types
from sklearn import decomposition as sk_decomp
from sklearn import cluster as sk_cluster
from sklearn.metrics import pairwise_distances
from os.path import exists, join
from os import makedirs
from nibabel import nifti1
from neurosynth.analysis import meta

class Clusterable(object):

    '''
    Args:
        dataset: The Dataset instance to extract data from.
        mask: A mask defining the voxels to cluster. Can be a filename,
            nibabel image, or numpy array (see Masker.mask() for details).
        features (str or list): Optional string or list of strings specifying
            any feature names to use for study selection. E.g., passing
            ['emotion', 'reward'] would retain for analysis only those studies
            associated with the features emotion or reward at a frequency
            greater than feature_threshold.
        feature_threshold (float): The threshold to use when selecting studies
            on the basis of features.
        min_voxels (int): Minimum number of active voxels a study
            must report in order to be retained in the dataset. By default,
            all studies are used.
        min_studies (int): Minimum number of studies a voxel must be
            active in in order to be retained in analysis. By default, all
            voxels are used.
    '''

    def __init__(self, dataset, mask=None, features=None,
                 feature_threshold=None, min_voxels=None, min_studies=None):

        self.dataset = dataset
        self.masker = deepcopy(dataset.masker)

        # Condition study inclusion on specific features
        if features is not None:
            ids = dataset.get_studies(features=features,
                                      frequency_threshold=feature_threshold)
            data = dataset.get_image_data(ids, dense=False)
        else:
            data = dataset.image_table.data

        # Trim data based on minimum number of voxels or studies
        if min_studies is not None:
            av = self.masker.unmask(
                data.sum(1) >= min_studies, output='vector')
            self.masker.add({'voxels': av})

        if min_voxels is not None:
            data = data[:, np.array(data.sum(0) >= min_voxels).squeeze()]

        if mask is not None:
            self.masker.add({'roi': mask})

        self.data = data[self.masker.get_mask(['voxels', 'roi']), :].toarray()

    def transform(self, transformer, transpose=False):
        ''' Apply a transformation to the Clusterable instance. Accepts any
        scikit-learn-style class that implements a fit_transform() method. '''
        data = self.data.T if transpose else self.data
        data = transformer.fit_transform(data)
        self.data = data.T if transpose else data
        return self

def magic(dataset, method='coactivation', roi_mask=None,
          coactivation_mask=None, features=None, feature_threshold=0.05,
          min_voxels_per_study=None, min_studies_per_voxel=None,
          reduce_reference='pca', n_components=100,
          distance_metric='correlation', clustering_algorithm='kmeans',
          n_clusters=5, clustering_kwargs={}, output_dir=None, filename=None,
          coactivation_images=False, coactivation_threshold=0.1):
    ''' Execute a full clustering analysis pipeline.
    Args:
        dataset: a Dataset instance to extract all data from.
        method (str): the overall clustering approach to use. Valid options:
            'coactivation' (default): Clusters voxel within the ROI mask based
                on shared pattern of coactivation with the rest of the brain.
            'studies': Treat each study as a feature in an n-dimensional space.
                I.e., voxels will be assigned to the same cluster if they tend
                to be co-reported in similar studies.
            'features': Voxels will be assigned to the same cluster if they
                tend to have similar feature vectors (i.e., the studies that
                activate those voxels tend to use similar terms).
        roi_mask: A string, nibabel image, or numpy array providing an
            inclusion mask of voxels to cluster. If None, the default mask
            in the Dataset instance is used (typically, all in-brain voxels).
        coactivation_mask: If method='coactivation', this mask defines the
            voxels to use when generating the pairwise distance matrix. For
            example, if a PFC mask is passed, all voxels in the roi_mask will
            be clustered based on how similar their patterns of coactivation
            with PFC voxels are. Can be a str, nibabel image, or numpy array.
        features (str or list): Optional string or list of strings specifying
            any feature names to use for study selection. E.g., passing
            ['emotion', 'reward'] would retain for analysis only those studies
            associated with the features emotion or reward at a frequency
            greater than feature_threshold.
        feature_threshold (float): The threshold to use when selecting studies
            on the basis of features.
        min_voxels_per_study (int): Minimum number of active voxels a study
            must report in order to be retained in the dataset. By default,
            all studies are used.
        min_studies_per_voxel (int): Minimum number of studies a voxel must be
            active in in order to be retained in analysis. By default, all
            voxels are used.
        reduce_reference (str, scikit-learn object or None): The dimensionality
            reduction algorithm to apply to the feature space prior to the
            computation of pairwise distances. If a string is passed (either
            'pca' or 'ica'), n_components must be specified. If None, no
            dimensionality reduction will be applied. Otherwise, must be a
            scikit-learn-style object that exposes a transform() method.
        n_components (int): Number of components to extract during the
            dimensionality reduction step. Only used if reduce_reference is
            a string.
        distance_metric (str): The distance metric to use when computing
            pairwise distances on the to-be-clustered voxels. Can be any of the
            metrics supported by sklearn.metrics.pairwise_distances.
        clustering_algorithm (str or scikit-learn object): the clustering
            algorithm to use. If a string, must be one of 'kmeans' or 'minik'.
            Otherwise, any sklearn class that exposes a fit_predict() method.
        n_clusters (int): If clustering_algorithm is a string, the number of
            clusters to extract.
        clustering_kwargs (dict): Additional keywords to pass to the clustering
            object.
        output_dir (str): The directory to write results to. If None (default),
            returns the cluster label image rather than saving to disk.
        filename (str): Name of cluster label image file. Defaults to
            cluster_labels_k{k}.nii.gz, where k is the number of clusters.
        coactivation_images (bool): If True, saves a meta-analytic coactivation
            map for every ROI in the resulting cluster map.
        coactivation_threshold (float or int): If coactivation_images is True,
            this is the threshold used to define whether or not a study is
            considered to activation within a cluster ROI. Integer values are
            interpreted as minimum number of voxels within the ROI; floats
            are interpreted as the proportion of voxels. Defaults to 0.1 (i.e.,
            10% of all voxels within ROI must be active).
    '''

    roi = Clusterable(dataset, roi_mask, min_voxels=min_voxels_per_study,
                      min_studies=min_studies_per_voxel, features=features,
                      feature_threshold=feature_threshold)

    if method == 'coactivation':
        reference = Clusterable(dataset, coactivation_mask,
                                min_voxels=min_voxels_per_study,
                                min_studies=min_studies_per_voxel,
                                features=features,
                                feature_threshold=feature_threshold)
    elif method == 'features':
        reference = deepcopy(roi)
        feature_data = dataset.feature_table.data
        n_studies = len(feature_data)
        reference.data = reference.data.dot(feature_data.values) / n_studies
    elif method == 'studies':
        reference = roi

    if reduce_reference is not None:

        if isinstance(reduce_reference, string_types):

            # Number of components can't exceed feature count or cluster count
            n_feat = reference.data.shape[1]
            n_components = min(n_components, n_feat)

            reduce_reference = {
                'pca': sk_decomp.PCA,
                'ica': sk_decomp.FastICA
            }[reduce_reference](n_components)

        # For non-coactivation-based approaches, transpose the data matrix
        transpose = (method == 'coactivation')
        reference = reference.transform(reduce_reference, transpose=transpose)

    if method == 'coactivation':
        distances = pairwise_distances(roi.data, reference.data,
                                       metric=distance_metric)
    else:
        distances = reference.data

    # TODO: add additional clustering methods
    if isinstance(clustering_algorithm, string_types):
        clustering_algorithm = {
            'kmeans': sk_cluster.KMeans,
            'minik': sk_cluster.MiniBatchKMeans
        }[clustering_algorithm](n_clusters, **clustering_kwargs)

    labels = clustering_algorithm.fit_predict(distances) + 1.
    header = roi.masker.get_header()
    header['cal_max'] = labels.max()
    header['cal_min'] = labels.min()
    voxel_labels = roi.masker.unmask(labels)
    img = nifti1.Nifti1Image(voxel_labels, None, header)

    if output_dir is not None:
        if not exists(output_dir):
            makedirs(output_dir)
        if filename is None:
            filename = 'cluster_labels_k%d.nii.gz' % n_clusters
        outfile = join(output_dir, filename)
        img.to_filename(outfile)

        # Write coactivation images
        if coactivation_images:
            for l in np.unique(voxel_labels):
                roi_mask = np.copy(voxel_labels)
                roi_mask[roi_mask != l] = 0
                ids = dataset.get_studies(
                    mask=roi_mask, activation_threshold=coactivation_threshold)
                ma = meta.MetaAnalysis(dataset, ids)
                ma.save_results(output_dir=join(output_dir, 'coactivation'),
                                prefix='cluster_%d_coactivation' % l)
    else:
        return img
<|endoftext|>


<commit_before># -*- coding: utf-8 -

import re

from time import time
from random import randint
from hashlib import md5
from sys import version_info
from datetime import datetime
from redis import StrictRedis

PY3K = version_info[0] == 3
EMAIL_REGEXP = re.compile(r"^[a-z0-9]+[_a-z0-9-]*(\.[_a-z0-9-]+)*@[a-z0-9]+[\.a-z0-9-]*(\.[a-z]{2,4})$")

def intid ():
	""" Return pseudo-unique decimal id. """
	return int((time() - 1374000000) * 100) * 100 + randint(0, 99)

def hexid ():
	""" Return pseudo-unique hexadecimal id. """
	return '%x' % intid()

class Key (object):
	def __init__ (self, key, db=None):
		self._key = key
		self._exists = None
		self._db = db or conf.db

	def exists (self):
		""" Check if model key exists. """

		if self._exists is None:
			self._exists = self._db.exists(self._key)

		return self._exists

	def delete (self, pipe=None):
		_pipe = self.getpipe(pipe)

		if self._exists is False:
			return

		_pipe.delete(self._key)

		if pipe is None and len(_pipe):
			_pipe.execute()

	def getpipe (self, pipe=None):
		return self._db.pipeline(transaction=True) if pipe is None else pipe

class Hash (Key):
	def __init__ (self, key, db=None):
		super(Hash, self).__init__(key, db)

		self._diff = dict()
		self._data = None

	def __len__ (self):
		return len(self.export())

	def __contains__ (self, name):
		if name in self._diff:
			return True

		self.load()
		return name in self._data

	def __getitem__ (self, name):
		if name in self._diff:
			return self._diff[name]

		self.load()
		return self._data[name] if name in self._data else None

	def __setitem__ (self, name, value):
		if self.loaded() and name in self._data and self._data[name] == value:
			pass

		else:
			self._diff[name] = value

	def __delitem__ (self, name):
		self._diff[name] = None

	def get (self, name, default=None):
		return self[name] if name in self else default

	def pop (self, name, default=None):
		val = self.get(name, default)
		del self[name]
		return val

	def export (self):
		self.load()
		data = self._data.copy()
		data.update(self._diff)

		return dict([(k, v) for (k, v) in data.items() if v is not None])

	def load (self):
		""" Load data into hash if needed. """

		if self.loaded():
			return

		self._data = dict()

		if self._exists is False:
			return

		for k, v in self._db.hgetall(self._key).items():
			if PY3K:
				k = k.decode(encoding='UTF-8')
				v = v.decode(encoding='UTF-8')

			self._data[k] = v

		self._exists = bool(len(self._data))

	def loaded (self):
		""" Check if model data is loaded. """
		return self._data is not None

	def unload (self):
		""" Unload model data. """
		self._data = None

	def getdiff (self):
		return self._diff.copy()

	def getorigin (self):
		self.load()
		return self._data.copy()

	def save (self, pipe=None):
		""" Save model (optionally within given parent pipe). """

		if not len(self._diff):
			return

		_pipe = self.getpipe(pipe)

		if self._exists is not False: # Remove hash item.
			for k, v in self._diff.items():
				if v is None:
					_pipe.hdel(self._key, k)

		# Cleanup diff from None values.
		self._diff = dict([i for i in self._diff.items() if i[1] is not None])

		if not len(self._diff):
			return

		if self.loaded():
			self._data.update(self._diff)

		_pipe.hmset(self._key, self._diff)
		self._exists = True
		self._diff = dict()

		if pipe is None and len(_pipe):
			_pipe.execute()

	def delete (self, pipe=None):
		""" Delete model (optionally within given parent pipe). """

		super(Hash, self).delete(pipe)

		self._diff = dict()
		self._data = dict()
		self._exists = False

	def revert (self):
		""" Revert local changes. """
		self._diff = dict()

class Field (object):
	def __init__ (self, field, index=False, unique=False):
		self.index = bool(index)
		self.unique = bool(unique)
		self.field = field

	def __get__ (self, model, owner):
		self.owner = owner
		return self if model is None else model[self.field]

	def __set__ (self, model, value):
		model[self.field] = value

	def find (self, val):
		assert self.index or self.unique
		key = index_key(self.owner.getprefix(), self.field, val)
		ids = self.owner.getdb().smembers(key)
		return [self.owner(model_id) for model_id in ids]

	def choice (self, val, count=1):
		""" Return *count* random model(s) from find() result. """

		assert self.index or self.unique
		key = index_key(self.owner.getprefix(), self.field, val)
		ids = self.owner.getdb().srandmember(key, count)

		return None if not len(ids) else \
			[self.owner(model_id) for model_id in ids]

	def save_index (self, model, pipe=None):
		key = index_key(model.getprefix(), self.field, model[self.field])

		if self.unique:
			model_id = bytes(model._id, 'utf-8') if PY3K else model._id
			ids = model.getdb().smembers(key)
			ids.discard(model_id)

			if len(ids):
				raise Exception('Duplicate key error')

		prev_idx_key = self._prev_idx_key(model)

		if prev_idx_key != key:
			pipe.srem(prev_idx_key, model._id)
			pipe.sadd(key, model._id)

	def del_index (self, model, pipe=None):
		# Get previous index value.
		pipe.srem(self._prev_idx_key(model), model._id)

	def _prev_idx_key (self, model):
		""" Get previous value index key. """

		if model.loaded() and self.field in model._data:
			prev_val = model._data[self.field]

		else:
			prev_val = model.getdb().hget(model.getkey(), self.field)

			if PY3K and prev_val is not None:
				prev_val = prev_val.decode('utf-8')

		return index_key(model.getprefix(), self.field, prev_val)

def index_key (prefix, name, value):
	return ':'.join((prefix, name, str(value)))

class String (Field):
	def __init__ (self, minlen=None, maxlen=None, **kw):
		super(String, self).__init__(**kw)

		if minlen is not None and maxlen is not None:
			assert minlen < maxlen

		self.minlen = minlen
		self.maxlen = maxlen

	def __get__ (self, model, owner):
		val = super(String, self).__get__(model, owner)
		return val.decode('utf-8') if PY3K and type(val) is bytes else val

	def __set__ (self, model, value):
		if value is not None:
			value = str(value)

			if self.minlen is not None and len(value) < self.minlen:
				raise Exception('Minimal length check failed')

			if self.maxlen is not None and len(value) > self.maxlen:
				raise Exception('Maximum length check failed')

		model[self.field] = value

class Email (Field):
	def __set__ (self, model, value):
		if value is not None and EMAIL_REGEXP.match(value) == None:
			raise Exception('Email validation failed')

		return super(Email, self).__set__(model, value)

class Integer (Field):
	def __init__ (self, minval=None, maxval=None, **kw):
		super(Integer, self).__init__(**kw)

		if minval is not None and maxval is not None:
			assert minval < maxval

		self.minval = minval
		self.maxval = maxval

	def __get__ (self, model, owner):
		self.owner = owner
		return self if model is None else int(model[self.field])

	def __set__ (self, model, value):
		if value is not None:
			value = int(value)

			if self.minval is not None and value < self.minval:
				raise Exception('Minimal value check failed')

			if self.maxval is not None and value > self.maxval:
				raise Exception('Maximum value check failed')

		model[self.field] = value

class DateTime (Field):
	def __get__ (self, model, owner):
		self.owner = owner

		if model is None:
			return self

		return None if model[self.field] is None \
			else datetime.fromtimestamp(int(model[self.field]))

	def __set__ (self, model, value):
		if value is not None:
			if type(value) is datetime:
				value = value.strftime('%s')

			value = int(value)

		model[self.field] = value

class MD5Pass (String):
	def __set__ (self, model, value):
		super(MD5Pass, self).__set__(model, value)

		if value is not None:
			val = model[self.field]

			if PY3K:
				val = val.encode('utf-8')

			model[self.field] = md5(val).hexdigest()

class Reference (Field):
	def __init__ (self, cls, **kw):
		super(Reference, self).__init__(**kw)
		assert issubclass(cls, Model)
		self._cls = cls

	def __get__ (self, model, owner):
		self.owner = owner

		if model is None:
			return self

		model_id = model[self.field]
		return None if model_id is None else self._cls(model_id)

	def __set__ (self, model, parent):
		if parent is not None:
			assert parent.__class__ is self._cls
			model[self.field] = parent._id

		else:
			model[self.field] = None

	def find (self, val):
		if isinstance(val, Model):
			val = val._id

		return super(Reference, self).find(val)

	def choice (self, val):
		if isinstance(val, Model):
			val = val._id

		return super(Reference, self).choice(val)

class MetaModel (type):
	def __new__ (mcs, name, bases, dct):
		cls = super(MetaModel, mcs).__new__(mcs, name, bases, dct)
		cls._objects = dict() # id -> model objects registry.
		return cls

	def __call__ (cls, model_id, *args, **kw):
		if PY3K and type(model_id) is bytes:
			model_id = model_id.decode('utf-8')

		else:
			model_id = str(model_id)

		if model_id not in cls._objects:
			cls._objects[model_id] = object.__new__(cls, *args, **kw)
			cls._objects[model_id].__init__(model_id)

		return cls._objects[model_id]

class conf (object):
	""" Configuration storage and model decorator. """

	db = StrictRedis()

	def __init__ (self, prefix=None, db=None):
		self._prefix = prefix
		self._db = db

	def __call__ (self, cls):
		if self._db is not None:
			cls._db = self._db

		if self._prefix is not None:
			Model._cls2prefix[cls] = self._prefix

		return cls

if PY3K:
	exec('class BaseModel (Hash, metaclass=MetaModel): pass')

else:
	exec('class BaseModel (Hash): __metaclass__ = MetaModel')

class Model (BaseModel):
	_cls2prefix = dict()

	def __init__ (self, model_id):
		self._id = model_id
		key = ':'.join((self.getprefix(), self._id))
		super(Model, self).__init__(key, self.getdb())

	@classmethod
	def getdb (cls):
		try:
			return cls._db

		except:
			return conf.db

	def getfields (self):
		""" Return name -> field dict of registered fields. """
		props = self.__class__.__dict__.items()
		return dict([(k, v) for (k, v) in props if isinstance(v, Field)])

	def getid (self):
		return self._id

	def getkey (self):
		return self._key

	@classmethod
	def getprefix (cls):
		return Model._cls2prefix[cls] \
			if cls in Model._cls2prefix \
				else cls.__name__.lower()

	def delete (self, pipe=None):
		_pipe = self.getpipe(pipe)

		for field in self.getfields().values():
			if field.index or field.unique:
				field.del_index(self, _pipe)

		super(Model, self).delete(_pipe)

		if pipe is None and len(_pipe):
			_pipe.execute()

	def save (self, pipe=None):
		_pipe = self.getpipe(pipe)
		diff = self.getdiff()

		fields = [f for f in self.getfields().values() \
			if f.field in diff and (f.index or f.unique)]

		for field in fields:
			field.save_index(self, _pipe)

		super(Model, self).save(_pipe)

		if pipe is None and len(_pipe):
			_pipe.execute()

	@classmethod
	def save_all (cls, pipe=None):
		""" Save all known models. Deleted models ignored by empty diff. """

		if cls is not Model:
			_pipe = cls.getdb().pipeline(transaction=True) \
				if pipe is None else pipe

			for model in cls._objects.values():
				model.save(_pipe)

			if pipe is None and len(_pipe):
				_pipe.execute()

		for child in cls.__subclasses__():
			child.save_all()

	def free (self):
		del self.__class__._objects[self._id]

	@classmethod
	def free_all (cls):
		""" Cleanup models registry. """

		cls._objects = dict()

		for child in cls.__subclasses__():
			child.free_all()

class FlaskRedisca (object):
	def __init__ (self, app=None, autosave=False):
		self.autosave = autosave

		if app is not None:
			self.init_app(app)

	def init_app (self, app):
		self.app = app

		conf.db = StrictRedis(**self.app.config['REDISCA'])
		self.app.before_request(self.before_request)
		self.app.teardown_request(self.after_request)

	def before_request (self):
		pass

	def after_request (self, exc):
		if self.autosave:
			Model.save_all()

		Model.free_all()
<commit_msg>Model.getfields decorated with classmethod<commit_after># -*- coding: utf-8 -

import re

from time import time
from random import randint
from hashlib import md5
from sys import version_info
from datetime import datetime
from redis import StrictRedis

PY3K = version_info[0] == 3
EMAIL_REGEXP = re.compile(r"^[a-z0-9]+[_a-z0-9-]*(\.[_a-z0-9-]+)*@[a-z0-9]+[\.a-z0-9-]*(\.[a-z]{2,4})$")

def intid ():
	""" Return pseudo-unique decimal id. """
	return int((time() - 1374000000) * 100) * 100 + randint(0, 99)

def hexid ():
	""" Return pseudo-unique hexadecimal id. """
	return '%x' % intid()

class Key (object):
	def __init__ (self, key, db=None):
		self._key = key
		self._exists = None
		self._db = db or conf.db

	def exists (self):
		""" Check if model key exists. """

		if self._exists is None:
			self._exists = self._db.exists(self._key)

		return self._exists

	def delete (self, pipe=None):
		_pipe = self.getpipe(pipe)

		if self._exists is False:
			return

		_pipe.delete(self._key)

		if pipe is None and len(_pipe):
			_pipe.execute()

	def getpipe (self, pipe=None):
		return self._db.pipeline(transaction=True) if pipe is None else pipe

class Hash (Key):
	def __init__ (self, key, db=None):
		super(Hash, self).__init__(key, db)

		self._diff = dict()
		self._data = None

	def __len__ (self):
		return len(self.export())

	def __contains__ (self, name):
		if name in self._diff:
			return True

		self.load()
		return name in self._data

	def __getitem__ (self, name):
		if name in self._diff:
			return self._diff[name]

		self.load()
		return self._data[name] if name in self._data else None

	def __setitem__ (self, name, value):
		if self.loaded() and name in self._data and self._data[name] == value:
			pass

		else:
			self._diff[name] = value

	def __delitem__ (self, name):
		self._diff[name] = None

	def get (self, name, default=None):
		return self[name] if name in self else default

	def pop (self, name, default=None):
		val = self.get(name, default)
		del self[name]
		return val

	def export (self):
		self.load()
		data = self._data.copy()
		data.update(self._diff)

		return dict([(k, v) for (k, v) in data.items() if v is not None])

	def load (self):
		""" Load data into hash if needed. """

		if self.loaded():
			return

		self._data = dict()

		if self._exists is False:
			return

		for k, v in self._db.hgetall(self._key).items():
			if PY3K:
				k = k.decode(encoding='UTF-8')
				v = v.decode(encoding='UTF-8')

			self._data[k] = v

		self._exists = bool(len(self._data))

	def loaded (self):
		""" Check if model data is loaded. """
		return self._data is not None

	def unload (self):
		""" Unload model data. """
		self._data = None

	def getdiff (self):
		return self._diff.copy()

	def getorigin (self):
		self.load()
		return self._data.copy()

	def save (self, pipe=None):
		""" Save model (optionally within given parent pipe). """

		if not len(self._diff):
			return

		_pipe = self.getpipe(pipe)

		if self._exists is not False: # Remove hash item.
			for k, v in self._diff.items():
				if v is None:
					_pipe.hdel(self._key, k)

		# Cleanup diff from None values.
		self._diff = dict([i for i in self._diff.items() if i[1] is not None])

		if not len(self._diff):
			return

		if self.loaded():
			self._data.update(self._diff)

		_pipe.hmset(self._key, self._diff)
		self._exists = True
		self._diff = dict()

		if pipe is None and len(_pipe):
			_pipe.execute()

	def delete (self, pipe=None):
		""" Delete model (optionally within given parent pipe). """

		super(Hash, self).delete(pipe)

		self._diff = dict()
		self._data = dict()
		self._exists = False

	def revert (self):
		""" Revert local changes. """
		self._diff = dict()

class Field (object):
	def __init__ (self, field, index=False, unique=False):
		self.index = bool(index)
		self.unique = bool(unique)
		self.field = field

	def __get__ (self, model, owner):
		self.owner = owner
		return self if model is None else model[self.field]

	def __set__ (self, model, value):
		model[self.field] = value

	def find (self, val):
		assert self.index or self.unique
		key = index_key(self.owner.getprefix(), self.field, val)
		ids = self.owner.getdb().smembers(key)
		return [self.owner(model_id) for model_id in ids]

	def choice (self, val, count=1):
		""" Return *count* random model(s) from find() result. """

		assert self.index or self.unique
		key = index_key(self.owner.getprefix(), self.field, val)
		ids = self.owner.getdb().srandmember(key, count)

		return None if not len(ids) else \
			[self.owner(model_id) for model_id in ids]

	def save_index (self, model, pipe=None):
		key = index_key(model.getprefix(), self.field, model[self.field])

		if self.unique:
			model_id = bytes(model._id, 'utf-8') if PY3K else model._id
			ids = model.getdb().smembers(key)
			ids.discard(model_id)

			if len(ids):
				raise Exception('Duplicate key error')

		prev_idx_key = self._prev_idx_key(model)

		if prev_idx_key != key:
			pipe.srem(prev_idx_key, model._id)
			pipe.sadd(key, model._id)

	def del_index (self, model, pipe=None):
		# Get previous index value.
		pipe.srem(self._prev_idx_key(model), model._id)

	def _prev_idx_key (self, model):
		""" Get previous value index key. """

		if model.loaded() and self.field in model._data:
			prev_val = model._data[self.field]

		else:
			prev_val = model.getdb().hget(model.getkey(), self.field)

			if PY3K and prev_val is not None:
				prev_val = prev_val.decode('utf-8')

		return index_key(model.getprefix(), self.field, prev_val)

def index_key (prefix, name, value):
	return ':'.join((prefix, name, str(value)))

class String (Field):
	def __init__ (self, minlen=None, maxlen=None, **kw):
		super(String, self).__init__(**kw)

		if minlen is not None and maxlen is not None:
			assert minlen < maxlen

		self.minlen = minlen
		self.maxlen = maxlen

	def __get__ (self, model, owner):
		val = super(String, self).__get__(model, owner)
		return val.decode('utf-8') if PY3K and type(val) is bytes else val

	def __set__ (self, model, value):
		if value is not None:
			value = str(value)

			if self.minlen is not None and len(value) < self.minlen:
				raise Exception('Minimal length check failed')

			if self.maxlen is not None and len(value) > self.maxlen:
				raise Exception('Maximum length check failed')

		model[self.field] = value

class Email (Field):
	def __set__ (self, model, value):
		if value is not None and EMAIL_REGEXP.match(value) == None:
			raise Exception('Email validation failed')

		return super(Email, self).__set__(model, value)

class Integer (Field):
	def __init__ (self, minval=None, maxval=None, **kw):
		super(Integer, self).__init__(**kw)

		if minval is not None and maxval is not None:
			assert minval < maxval

		self.minval = minval
		self.maxval = maxval

	def __get__ (self, model, owner):
		self.owner = owner
		return self if model is None else int(model[self.field])

	def __set__ (self, model, value):
		if value is not None:
			value = int(value)

			if self.minval is not None and value < self.minval:
				raise Exception('Minimal value check failed')

			if self.maxval is not None and value > self.maxval:
				raise Exception('Maximum value check failed')

		model[self.field] = value

class DateTime (Field):
	def __get__ (self, model, owner):
		self.owner = owner

		if model is None:
			return self

		return None if model[self.field] is None \
			else datetime.fromtimestamp(int(model[self.field]))

	def __set__ (self, model, value):
		if value is not None:
			if type(value) is datetime:
				value = value.strftime('%s')

			value = int(value)

		model[self.field] = value

class MD5Pass (String):
	def __set__ (self, model, value):
		super(MD5Pass, self).__set__(model, value)

		if value is not None:
			val = model[self.field]

			if PY3K:
				val = val.encode('utf-8')

			model[self.field] = md5(val).hexdigest()

class Reference (Field):
	def __init__ (self, cls, **kw):
		super(Reference, self).__init__(**kw)
		assert issubclass(cls, Model)
		self._cls = cls

	def __get__ (self, model, owner):
		self.owner = owner

		if model is None:
			return self

		model_id = model[self.field]
		return None if model_id is None else self._cls(model_id)

	def __set__ (self, model, parent):
		if parent is not None:
			assert parent.__class__ is self._cls
			model[self.field] = parent._id

		else:
			model[self.field] = None

	def find (self, val):
		if isinstance(val, Model):
			val = val._id

		return super(Reference, self).find(val)

	def choice (self, val):
		if isinstance(val, Model):
			val = val._id

		return super(Reference, self).choice(val)

class MetaModel (type):
	def __new__ (mcs, name, bases, dct):
		cls = super(MetaModel, mcs).__new__(mcs, name, bases, dct)
		cls._objects = dict() # id -> model objects registry.
		return cls

	def __call__ (cls, model_id, *args, **kw):
		if PY3K and type(model_id) is bytes:
			model_id = model_id.decode('utf-8')

		else:
			model_id = str(model_id)

		if model_id not in cls._objects:
			cls._objects[model_id] = object.__new__(cls, *args, **kw)
			cls._objects[model_id].__init__(model_id)

		return cls._objects[model_id]

class conf (object):
	""" Configuration storage and model decorator. """

	db = StrictRedis()

	def __init__ (self, prefix=None, db=None):
		self._prefix = prefix
		self._db = db

	def __call__ (self, cls):
		if self._db is not None:
			cls._db = self._db

		if self._prefix is not None:
			Model._cls2prefix[cls] = self._prefix

		return cls

if PY3K:
	exec('class BaseModel (Hash, metaclass=MetaModel): pass')

else:
	exec('class BaseModel (Hash): __metaclass__ = MetaModel')

class Model (BaseModel):
	_cls2prefix = dict()

	def __init__ (self, model_id):
		self._id = model_id
		key = ':'.join((self.getprefix(), self._id))
		super(Model, self).__init__(key, self.getdb())

	@classmethod
	def getdb (cls):
		try:
			return cls._db

		except:
			return conf.db

	@classmethod
	def getfields (cls):
		""" Return name -> field dict of registered fields. """
		props = cls.__dict__.items()
		return dict([(k, v) for (k, v) in props if isinstance(v, Field)])

	def getid (self):
		return self._id

	def getkey (self):
		return self._key

	@classmethod
	def getprefix (cls):
		return Model._cls2prefix[cls] \
			if cls in Model._cls2prefix \
				else cls.__name__.lower()

	def delete (self, pipe=None):
		_pipe = self.getpipe(pipe)

		for field in self.getfields().values():
			if field.index or field.unique:
				field.del_index(self, _pipe)

		super(Model, self).delete(_pipe)

		if pipe is None and len(_pipe):
			_pipe.execute()

	def save (self, pipe=None):
		_pipe = self.getpipe(pipe)
		diff = self.getdiff()

		fields = [f for f in self.getfields().values() \
			if f.field in diff and (f.index or f.unique)]

		for field in fields:
			field.save_index(self, _pipe)

		super(Model, self).save(_pipe)

		if pipe is None and len(_pipe):
			_pipe.execute()

	@classmethod
	def save_all (cls, pipe=None):
		""" Save all known models. Deleted models ignored by empty diff. """

		if cls is not Model:
			_pipe = cls.getdb().pipeline(transaction=True) \
				if pipe is None else pipe

			for model in cls._objects.values():
				model.save(_pipe)

			if pipe is None and len(_pipe):
				_pipe.execute()

		for child in cls.__subclasses__():
			child.save_all()

	def free (self):
		del self.__class__._objects[self._id]

	@classmethod
	def free_all (cls):
		""" Cleanup models registry. """

		cls._objects = dict()

		for child in cls.__subclasses__():
			child.free_all()

class FlaskRedisca (object):
	def __init__ (self, app=None, autosave=False):
		self.autosave = autosave

		if app is not None:
			self.init_app(app)

	def init_app (self, app):
		self.app = app

		conf.db = StrictRedis(**self.app.config['REDISCA'])
		self.app.before_request(self.before_request)
		self.app.teardown_request(self.after_request)

	def before_request (self):
		pass

	def after_request (self, exc):
		if self.autosave:
			Model.save_all()

		Model.free_all()
<|endoftext|>