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

import re
import netaddr
from textwrap import dedent
from sqlalchemy.schema import UniqueConstraint
from sqlalchemy.orm import validates

from pyfarm.flaskapp import db
from pyfarm.ext.config.core.loader import Loader
from pyfarm.ext.config.enum import AgentState
from pyfarm.ext.system.network import IP_NONNETWORK
from pyfarm.models.mixins import StateValidationMixin
from pyfarm.models.core import (
    DBCFG, TABLE_AGENT, TABLE_AGENT_TAGS, TABLE_AGENT_SOFTWARE, IDColumn)

REGEX_HOSTNAME = re.compile("^(([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*"
                            "[a-zA-Z0-9])\.)*([A-Za-z0-9]|[A-Za-z0-9]"
                            "[A-Za-z0-9\-]*[A-Za-z0-9])$")

class AgentTaggingMixin(object):
    """
    Mixin used which provides some common structures to
    :class:`.AgentTagsModel` and :class:`.AgentSoftwareModel`
    """
    @validates("tag", "software")
    def validate_string_column(self, key, value):
        """
        Ensures `value` is a string or something that can be converted
        to a string.
        """
        if isinstance(value, (int, long)):
            value = str(value)
        elif not isinstance(value, basestring):
            raise ValueError("expected a string for `%s`" % key)

        return value

class AgentTagsModel(db.Model, AgentTaggingMixin):
    """
    Table model used to store tags for an agent.

    .. note::
        This table enforces two forms of uniqueness.  The :attr:`id` column
        must be unique and the combination of these columns must also be
        unique to limit the frequency of duplicate data:

            * :attr:`_agentid`
            * :attr:`tag`

    .. autoattribute:: _agentid
    """
    __tablename__ = TABLE_AGENT_TAGS
    __table_args__ = (UniqueConstraint("_agentid", "tag"), )
    id = IDColumn()
    _agentid = db.Column(db.Integer, db.ForeignKey("%s.id" % TABLE_AGENT),
                         doc=dedent("""
                         The foreign key which stores :attr:`AgentModel.id`"""))
    tag = db.Column(db.String,
                    doc=dedent("""
                    A string value to tag an :class:`.AgentModel`. Generally
                    this value is used for grouping like resources together
                    on the network but could also be used by jobs as a sort of
                    requirement."""))

class AgentSoftwareModel(db.Model, AgentTaggingMixin):
    """
    Stores information about an the software installed on
    an agent.

    .. note::
        This table enforces two forms of uniqueness.  The :attr:`id` column
        must be unique and the combination of these columns must also be
        unique to limit the frequency of duplicate data:

            * :attr:`_agentid`
            * :attr:`version`
            * :attr:`software`

    .. autoattribute:: _agentid
    """
    __tablename__ = TABLE_AGENT_SOFTWARE
    __table_args__ = (UniqueConstraint("_agentid", "version", "software"), )
    id = IDColumn()
    _agentid = db.Column(db.Integer, db.ForeignKey("%s.id" % TABLE_AGENT),
                         doc=dedent("""
                         The foreign key which stores :attr:`AgentModel.id`"""))
    software = db.Column(db.String, nullable=False,
                         doc=dedent("""
                         The name of the software installed.  No normalization
                         is performed prior to being stored in the database"""))
    version = db.Column(db.String, default="any", nullable=False,
                        doc=dedent("""
                        The version of the software installed on a host.  This
                        value does not follow any special formatting rules
                        because the format depends on the 3rd party."""))

class AgentModel(db.Model, StateValidationMixin):
    """
    Stores information about an agent include its network address,
    state, allocation configuration, etc.

    .. note::
        This table enforces two forms of uniqueness.  The :attr:`id` column
        must be unique and the combination of these columns must also be
        unique to limit the frequency of duplicate data:

            * :attr:`hostname`
            * :attr:`ip`
            * :attr:`subnet`
            * :attr:`port`

    """
    __tablename__ = TABLE_AGENT
    __table_args__ = (UniqueConstraint("hostname", "ip", "subnet", "port"), )
    STATE_ENUM = AgentState()
    STATE_DEFAULT = STATE_ENUM.ONLINE
    id = IDColumn()

    # host state information
    enabled = db.Column(db.Boolean, default=True, nullable=False,
                        doc=dedent("""
                        Tells the queue and various other parts of PyFarm if
                        this host is considered part of the pool.  Setting
                        this to True will prevent:

                        * new work from being sent to an agent
                        * failed job types from rerunning
                        * batch commands being picked up over the pub/sub
                          channel"""))
    state = db.Column(db.Integer, default=STATE_DEFAULT, nullable=False,
                      doc=dedent("""
                      Stores the current state of the host.  This value can be
                      changed either by a master telling the host to do
                      something with a task or from the host via REST api.

                      .. csv-table:: **Values (from enum.yml:AgentState)**
                          :header: Integer, Description
                          :widths: 10, 50

                          16,Offline - host is unreachable
                          17,Online - ready to receive work
                          18,Disabled - same as online but cannot receive work
                          19,Running - currently processing work"""))
    hostname = db.Column(db.String, nullable=False,
                         doc=dedent("""
                         The hostname we should use to talk to this host.
                         Preferably this value will be the fully qualified
                         name instead of the base hostname alone."""))
    ip = db.Column(db.String(15), nullable=False,
                   doc=dedent("""
                   The IPv4 network address this host resides on.  This is
                   'best guess' address using factors described in
                   :meth:`pyfarm.system.network.NetworkInfo.ip`"""))
    subnet = db.Column(db.String(15), nullable=False,
                       doc=dedent("""
                       The subnet associated with the agent at the time we
                       discovered :attr:`ip`"""))
    ram = db.Column(db.Integer, nullable=False,
                    doc=dedent("""
                    The amount of ram (in megabytes) installed on the agent.
                    This value is provided by
                    :attr:`pyfarm.system.memory.MemoryInfo.TOTAL_RAM`"""))
    cpus = db.Column(db.Integer, nullable=False,
                     doc=dedent("""
                     The number of cpus installed on the agent.  This value
                     is provided by
                     :attr:`pyfarm.system.processor.ProcessorInfo.NUM_CPUS`
                     """))
    port = db.Column(db.Integer, nullable=False,
                     doc=dedent("""
                     The port the agent is currently running on"""))

    # Max allocation of the two primary resources which `1.0` is 100%
    # allocation.  For `cpu_allocation` 100% allocation typically means
    # one task per cpu.
    ram_allocation = db.Column(db.Float, nullable=False,
                               default=DBCFG.get("agent.ram_allocation", .8),
                               doc=dedent("""
                               The amount of ram the agent is allowed to
                               allocate towards work.  A value of `1.0` would
                               mean to let the agent use all of the memory
                               installed on the system when assigning work.

                               **configured by**: `agent.ram_allocation`"""))
    cpu_allocation = db.Column(db.Float, nullable=False,
                               default=DBCFG.get("agent.cpu_allocation", 1.0),
                               doc=dedent("""
                               The total amount of cpu space an agent is
                               allowed to process work in.  A value of `1.0`
                               would mean an agent can handle as much work
                               as the system could handle given the
                               requirements of a task.  For example if an agent
                               has 8 cpus, cpu_allocation is .5, and a task
                               requires 4 cpus then only that task will run
                               on the system.

                               **configured by**: `agent.cpu_allocation`"""))

    # relationships
    tasks = db.relationship("TaskModel", backref="agent", lazy="dynamic",
                            doc=dedent("""
                            Relationship between an :class:`AgentModel`
                            and any :class:`pyfarm.models.TaskModel`
                            objects"""))
    tags = db.relationship("AgentTagsModel", backref="agent", lazy="dynamic",
                           doc=dedent("""
                           Relationship between an :class:`AgentModel`
                           and any :class:`pyfarm.models.AgentTagsModel`
                           objects"""))
    software = db.relation("AgentSoftwareModel", backref="agent",
                           lazy="dynamic", doc=dedent("""
                           Relationship between an :class:`AgentModel`
                           and any :class:`pyfarm.models.AgentSoftwareModel`
                           objects"""))

    @validates("hostname")
    def validate_hostname(self, key, value):
        """
        Ensures that the hostname provided by `value` matches a regular
        expression that expresses what a valid hostname is.
        """
        # ensure hostname does not contain characters we can't use
        if not REGEX_HOSTNAME.match(value):
            raise ValueError("%s is not valid for %s" % (value, key))

        return value

    @validates("ip", "subnet")
    def validate_address(self, key, value):
        """
        Ensures the :attr:`ip` and :attr:`subnet` are valid.  For ip addresses
        this will make sure `value` is:

            * not a hostmask
            * not link local (:rfc:`3927`)
            * not used for multicast (:rfc:`1112`)
            * not a netmask (:rfc:`4632`)
            * not reserved (:rfc:`6052`)
            * a private address (:rfc:`1918`)

        This method will also value subnet masks to ensure their format is
        valid.
        """
        try:
            ip = netaddr.IPAddress(value)

        except ValueError, e:
            raise ValueError(
                "%s is not a valid address format: %s" % (value, e))

        if key == "ip":
            valid = all([
                not ip.is_hostmask(), not ip.is_link_local(),
                not ip.is_loopback(), not ip.is_multicast(),
                not ip.is_netmask(), ip.is_private(),
                not ip.is_reserved()
            ])
            if not valid:
                raise ValueError("%s it not a private ip address" % value)

        elif key == "subnet":
            valid = all([
                not ip.is_hostmask(), not ip.is_link_local(),
                not ip.is_loopback(), not ip.is_multicast(),
                ip.is_netmask(), not ip.is_private(),
                ip.is_reserved()
            ])
            if not valid:
                raise ValueError("%s is not valid subnet" % value)

        return value

    @validates("ram", "cpus", "port")
    def validate_resource(self, key, value):
        """
        Ensure the `value` provided for `key` is within an expected range as
        specified in `agent.yml`
        """
        if value is None:
            return value

        min_value = DBCFG.get("agent.min_%s" % key)
        max_value = DBCFG.get("agent.max_%s" % key)

        # quick sanity check of the incoming config
        assert isinstance(min_value, int), "db.min_%s must be an integer" % key
        assert isinstance(max_value, int), "db.max_%s must be an integer" % key
        assert min_value >= 1, "db.min_%s must be > 0" % key
        assert max_value >= 1, "db.max_%s must be > 0" % key

        # check the provided input
        if min_value > value or value > max_value:
            msg = "value for `%s` must be between " % key
            msg += "%s and %s" % (min_value, max_value)
            raise ValueError(msg)

        return value

class AgentTag(AgentTagsModel):
    """
    Provides :meth:`__init__` for :class:`.AgentTagsModel` so the model can
    be instanced with initial values.
    """
    def __init__(self, agent, tag):
        if isinstance(agent, Agent):
            agentid = agent.id
        else:
            agentid = agent

        self._agentid = agentid
        self.tag = tag

class AgentSoftware(AgentSoftwareModel):
    """
    Provides :meth:`__init__` for :class:`AgentSoftwareModel` so the model can
    be instanced with initial values.
    """
    def __init__(self, agent, software, version="any"):
        if isinstance(agent, Agent):
            agentid = agent.id
        else:
            agentid = agent

        self._agentid = agentid
        self.software = software
        self.version = version

class Agent(AgentModel):
    """
    Provides :meth:`__init__` for :class:`AgentModel` so the model can
    be instanced with initial values.
    """
    def __init__(self, hostname, ip, subnet, state=None, enabled=None,
                 ram=None, cpus=None, port=None, ram_allocation=None,
                 cpu_allocation=None):
        self.hostname = hostname
        self.ip = ip
        self.subnet = subnet

        if state is not None:
            self.state = state

        if enabled is not None:
            self.enabled = enabled

        if ram is not None:
            self.ram = ram

        if cpus is not None:
            self.cpus = cpus

        if port is not None:
            self.port = port

        if ram is not None:
            self.ram_allocation = ram_allocation

        if cpu_allocation is not None:
            self.cpu_allocation = cpu_allocation<commit_msg>should be checking against AgentModel<commit_after># No shebang line, this module is meant to be imported
#
# Copyright 2013 Oliver Palmer
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import netaddr
from textwrap import dedent
from sqlalchemy.schema import UniqueConstraint
from sqlalchemy.orm import validates

from pyfarm.flaskapp import db
from pyfarm.ext.config.core.loader import Loader
from pyfarm.ext.config.enum import AgentState
from pyfarm.ext.system.network import IP_NONNETWORK
from pyfarm.models.mixins import StateValidationMixin
from pyfarm.models.core import (
    DBCFG, TABLE_AGENT, TABLE_AGENT_TAGS, TABLE_AGENT_SOFTWARE, IDColumn)

REGEX_HOSTNAME = re.compile("^(([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*"
                            "[a-zA-Z0-9])\.)*([A-Za-z0-9]|[A-Za-z0-9]"
                            "[A-Za-z0-9\-]*[A-Za-z0-9])$")

class AgentTaggingMixin(object):
    """
    Mixin used which provides some common structures to
    :class:`.AgentTagsModel` and :class:`.AgentSoftwareModel`
    """
    @validates("tag", "software")
    def validate_string_column(self, key, value):
        """
        Ensures `value` is a string or something that can be converted
        to a string.
        """
        if isinstance(value, (int, long)):
            value = str(value)
        elif not isinstance(value, basestring):
            raise ValueError("expected a string for `%s`" % key)

        return value

class AgentTagsModel(db.Model, AgentTaggingMixin):
    """
    Table model used to store tags for an agent.

    .. note::
        This table enforces two forms of uniqueness.  The :attr:`id` column
        must be unique and the combination of these columns must also be
        unique to limit the frequency of duplicate data:

            * :attr:`_agentid`
            * :attr:`tag`

    .. autoattribute:: _agentid
    """
    __tablename__ = TABLE_AGENT_TAGS
    __table_args__ = (UniqueConstraint("_agentid", "tag"), )
    id = IDColumn()
    _agentid = db.Column(db.Integer, db.ForeignKey("%s.id" % TABLE_AGENT),
                         doc=dedent("""
                         The foreign key which stores :attr:`AgentModel.id`"""))
    tag = db.Column(db.String,
                    doc=dedent("""
                    A string value to tag an :class:`.AgentModel`. Generally
                    this value is used for grouping like resources together
                    on the network but could also be used by jobs as a sort of
                    requirement."""))

class AgentSoftwareModel(db.Model, AgentTaggingMixin):
    """
    Stores information about an the software installed on
    an agent.

    .. note::
        This table enforces two forms of uniqueness.  The :attr:`id` column
        must be unique and the combination of these columns must also be
        unique to limit the frequency of duplicate data:

            * :attr:`_agentid`
            * :attr:`version`
            * :attr:`software`

    .. autoattribute:: _agentid
    """
    __tablename__ = TABLE_AGENT_SOFTWARE
    __table_args__ = (UniqueConstraint("_agentid", "version", "software"), )
    id = IDColumn()
    _agentid = db.Column(db.Integer, db.ForeignKey("%s.id" % TABLE_AGENT),
                         doc=dedent("""
                         The foreign key which stores :attr:`AgentModel.id`"""))
    software = db.Column(db.String, nullable=False,
                         doc=dedent("""
                         The name of the software installed.  No normalization
                         is performed prior to being stored in the database"""))
    version = db.Column(db.String, default="any", nullable=False,
                        doc=dedent("""
                        The version of the software installed on a host.  This
                        value does not follow any special formatting rules
                        because the format depends on the 3rd party."""))

class AgentModel(db.Model, StateValidationMixin):
    """
    Stores information about an agent include its network address,
    state, allocation configuration, etc.

    .. note::
        This table enforces two forms of uniqueness.  The :attr:`id` column
        must be unique and the combination of these columns must also be
        unique to limit the frequency of duplicate data:

            * :attr:`hostname`
            * :attr:`ip`
            * :attr:`subnet`
            * :attr:`port`

    """
    __tablename__ = TABLE_AGENT
    __table_args__ = (UniqueConstraint("hostname", "ip", "subnet", "port"), )
    STATE_ENUM = AgentState()
    STATE_DEFAULT = STATE_ENUM.ONLINE
    id = IDColumn()

    # host state information
    enabled = db.Column(db.Boolean, default=True, nullable=False,
                        doc=dedent("""
                        Tells the queue and various other parts of PyFarm if
                        this host is considered part of the pool.  Setting
                        this to True will prevent:

                        * new work from being sent to an agent
                        * failed job types from rerunning
                        * batch commands being picked up over the pub/sub
                          channel"""))
    state = db.Column(db.Integer, default=STATE_DEFAULT, nullable=False,
                      doc=dedent("""
                      Stores the current state of the host.  This value can be
                      changed either by a master telling the host to do
                      something with a task or from the host via REST api.

                      .. csv-table:: **Values (from enum.yml:AgentState)**
                          :header: Integer, Description
                          :widths: 10, 50

                          16,Offline - host is unreachable
                          17,Online - ready to receive work
                          18,Disabled - same as online but cannot receive work
                          19,Running - currently processing work"""))
    hostname = db.Column(db.String, nullable=False,
                         doc=dedent("""
                         The hostname we should use to talk to this host.
                         Preferably this value will be the fully qualified
                         name instead of the base hostname alone."""))
    ip = db.Column(db.String(15), nullable=False,
                   doc=dedent("""
                   The IPv4 network address this host resides on.  This is
                   'best guess' address using factors described in
                   :meth:`pyfarm.system.network.NetworkInfo.ip`"""))
    subnet = db.Column(db.String(15), nullable=False,
                       doc=dedent("""
                       The subnet associated with the agent at the time we
                       discovered :attr:`ip`"""))
    ram = db.Column(db.Integer, nullable=False,
                    doc=dedent("""
                    The amount of ram (in megabytes) installed on the agent.
                    This value is provided by
                    :attr:`pyfarm.system.memory.MemoryInfo.TOTAL_RAM`"""))
    cpus = db.Column(db.Integer, nullable=False,
                     doc=dedent("""
                     The number of cpus installed on the agent.  This value
                     is provided by
                     :attr:`pyfarm.system.processor.ProcessorInfo.NUM_CPUS`
                     """))
    port = db.Column(db.Integer, nullable=False,
                     doc=dedent("""
                     The port the agent is currently running on"""))

    # Max allocation of the two primary resources which `1.0` is 100%
    # allocation.  For `cpu_allocation` 100% allocation typically means
    # one task per cpu.
    ram_allocation = db.Column(db.Float, nullable=False,
                               default=DBCFG.get("agent.ram_allocation", .8),
                               doc=dedent("""
                               The amount of ram the agent is allowed to
                               allocate towards work.  A value of `1.0` would
                               mean to let the agent use all of the memory
                               installed on the system when assigning work.

                               **configured by**: `agent.ram_allocation`"""))
    cpu_allocation = db.Column(db.Float, nullable=False,
                               default=DBCFG.get("agent.cpu_allocation", 1.0),
                               doc=dedent("""
                               The total amount of cpu space an agent is
                               allowed to process work in.  A value of `1.0`
                               would mean an agent can handle as much work
                               as the system could handle given the
                               requirements of a task.  For example if an agent
                               has 8 cpus, cpu_allocation is .5, and a task
                               requires 4 cpus then only that task will run
                               on the system.

                               **configured by**: `agent.cpu_allocation`"""))

    # relationships
    tasks = db.relationship("TaskModel", backref="agent", lazy="dynamic",
                            doc=dedent("""
                            Relationship between an :class:`AgentModel`
                            and any :class:`pyfarm.models.TaskModel`
                            objects"""))
    tags = db.relationship("AgentTagsModel", backref="agent", lazy="dynamic",
                           doc=dedent("""
                           Relationship between an :class:`AgentModel`
                           and any :class:`pyfarm.models.AgentTagsModel`
                           objects"""))
    software = db.relation("AgentSoftwareModel", backref="agent",
                           lazy="dynamic", doc=dedent("""
                           Relationship between an :class:`AgentModel`
                           and any :class:`pyfarm.models.AgentSoftwareModel`
                           objects"""))

    @validates("hostname")
    def validate_hostname(self, key, value):
        """
        Ensures that the hostname provided by `value` matches a regular
        expression that expresses what a valid hostname is.
        """
        # ensure hostname does not contain characters we can't use
        if not REGEX_HOSTNAME.match(value):
            raise ValueError("%s is not valid for %s" % (value, key))

        return value

    @validates("ip", "subnet")
    def validate_address(self, key, value):
        """
        Ensures the :attr:`ip` and :attr:`subnet` are valid.  For ip addresses
        this will make sure `value` is:

            * not a hostmask
            * not link local (:rfc:`3927`)
            * not used for multicast (:rfc:`1112`)
            * not a netmask (:rfc:`4632`)
            * not reserved (:rfc:`6052`)
            * a private address (:rfc:`1918`)

        This method will also value subnet masks to ensure their format is
        valid.
        """
        try:
            ip = netaddr.IPAddress(value)

        except ValueError, e:
            raise ValueError(
                "%s is not a valid address format: %s" % (value, e))

        if key == "ip":
            valid = all([
                not ip.is_hostmask(), not ip.is_link_local(),
                not ip.is_loopback(), not ip.is_multicast(),
                not ip.is_netmask(), ip.is_private(),
                not ip.is_reserved()
            ])
            if not valid:
                raise ValueError("%s it not a private ip address" % value)

        elif key == "subnet":
            valid = all([
                not ip.is_hostmask(), not ip.is_link_local(),
                not ip.is_loopback(), not ip.is_multicast(),
                ip.is_netmask(), not ip.is_private(),
                ip.is_reserved()
            ])
            if not valid:
                raise ValueError("%s is not valid subnet" % value)

        return value

    @validates("ram", "cpus", "port")
    def validate_resource(self, key, value):
        """
        Ensure the `value` provided for `key` is within an expected range as
        specified in `agent.yml`
        """
        if value is None:
            return value

        min_value = DBCFG.get("agent.min_%s" % key)
        max_value = DBCFG.get("agent.max_%s" % key)

        # quick sanity check of the incoming config
        assert isinstance(min_value, int), "db.min_%s must be an integer" % key
        assert isinstance(max_value, int), "db.max_%s must be an integer" % key
        assert min_value >= 1, "db.min_%s must be > 0" % key
        assert max_value >= 1, "db.max_%s must be > 0" % key

        # check the provided input
        if min_value > value or value > max_value:
            msg = "value for `%s` must be between " % key
            msg += "%s and %s" % (min_value, max_value)
            raise ValueError(msg)

        return value

class AgentTag(AgentTagsModel):
    """
    Provides :meth:`__init__` for :class:`.AgentTagsModel` so the model can
    be instanced with initial values.
    """
    def __init__(self, agent, tag):
        if isinstance(agent, AgentModel):
            agentid = agent.id
        else:
            agentid = agent

        self._agentid = agentid
        self.tag = tag

class AgentSoftware(AgentSoftwareModel):
    """
    Provides :meth:`__init__` for :class:`AgentSoftwareModel` so the model can
    be instanced with initial values.
    """
    def __init__(self, agent, software, version="any"):
        if isinstance(agent, AgentModel):
            agentid = agent.id
        else:
            agentid = agent

        self._agentid = agentid
        self.software = software
        self.version = version

class Agent(AgentModel):
    """
    Provides :meth:`__init__` for :class:`AgentModel` so the model can
    be instanced with initial values.
    """
    def __init__(self, hostname, ip, subnet, state=None, enabled=None,
                 ram=None, cpus=None, port=None, ram_allocation=None,
                 cpu_allocation=None):
        self.hostname = hostname
        self.ip = ip
        self.subnet = subnet

        if state is not None:
            self.state = state

        if enabled is not None:
            self.enabled = enabled

        if ram is not None:
            self.ram = ram

        if cpus is not None:
            self.cpus = cpus

        if port is not None:
            self.port = port

        if ram is not None:
            self.ram_allocation = ram_allocation

        if cpu_allocation is not None:
            self.cpu_allocation = cpu_allocation<|endoftext|>


<commit_before>""" Load and parse configuration file
"""
import ConfigParser
import os
from string import Template

def load_config(config_file=None):
    """Loads a configuration file.

    By default it assumes ~/.ngi_config/ngi_reports.conf
    """
    try:
        if not config_file:
            config_file = os.path.join(os.environ.get('HOME'), '.ngi_config', 'ngi_reports.conf')
            if not os.path.exists(config_file):
                config_file = os.path.join(os.environ.get("NGI_REPORTS_CONFIG"))
        config = ConfigParser.SafeConfigParser({'ngi_node': 'unknown'})
        with open(config_file) as f:
            config.readfp(f)
        return config
    except IOError:
        raise IOError(("There was a problem loading the configuration file. "
                "Please make sure that ~/.ngi_config/ngi_reports.conf exists "
                "or env vairable 'NGI_REPORTS_CONFIG' is set with path to conf "
                "file and set with read permissions"))

def expand_path(input_path, substitutions):
    """Use python's string templates to replace substitution patterns in
        the input path for those available in the substitution dict
    """
    return Template(input_path).substitute(substitutions)<commit_msg>fix typo, tack @sylvinite<commit_after>""" Load and parse configuration file
"""
import ConfigParser
import os
from string import Template

def load_config(config_file=None):
    """Loads a configuration file.

    By default it assumes ~/.ngi_config/ngi_reports.conf
    """
    try:
        if not config_file:
            config_file = os.path.join(os.environ.get('HOME'), '.ngi_config', 'ngi_reports.conf')
            if not os.path.exists(config_file):
                config_file = os.path.join(os.environ.get("NGI_REPORTS_CONFIG"))
        config = ConfigParser.SafeConfigParser({'ngi_node': 'unknown'})
        with open(config_file) as f:
            config.readfp(f)
        return config
    except IOError:
        raise IOError(("There was a problem loading the configuration file. "
                "Please make sure that ~/.ngi_config/ngi_reports.conf exists "
                "or env variable 'NGI_REPORTS_CONFIG' is set with path to conf "
                "file and set with read permissions"))

def expand_path(input_path, substitutions):
    """Use python's string templates to replace substitution patterns in
        the input path for those available in the substitution dict
    """
    return Template(input_path).substitute(substitutions)<|endoftext|>


<commit_before>'''
	Server Density
	www.serverdensity.com
	----
	A web based server resource monitoring application

	Licensed under Simplified BSD License (see LICENSE)
	(C) Boxed Ice 2009 all rights reserved
'''

# General config
DEBUG_MODE = 0
CHECK_FREQUENCY = 60

VERSION = '1.0.0b3'

# Core modules
import ConfigParser
import logging
import sched
import time
import sys

# Check we're not using an old version of Python. We need 2.4 above because some modules (like subprocess)
# were only introduced in 2.4.
if int(sys.version_info[1]) <= 4:
	print 'You are using an outdated version of Python. Please update to v2.4 or above (v3 is not supported).'
	sys.exit(2)

# Custom modules
from checks import checks
from daemon import Daemon

# Config handling
try:
	config = ConfigParser.ConfigParser()
	config.read('config.cfg')
	SD_URL = config.get('Main', 'sd_url')
	AGENT_KEY = config.get('Main', 'agent_key')

except ConfigParser.NoSectionError, e:
	print 'Config file not found or incorrectly formatted'
	quit()

except ConfigParser.ParsingError, e:
	print 'Config file not found or incorrectly formatted'
	quit()

# Override the generic daemon class to run our checks
class agent(Daemon):

	def run(self):
		agentLogger = logging.getLogger('agent')
		agentLogger.debug('Creating checks instance')

		# Checks instance
		c = checks(SD_URL, AGENT_KEY, CHECK_FREQUENCY, VERSION, DEBUG_MODE)

		# Schedule the checks
		agentLogger.debug('Scheduling checks every ' + str(CHECK_FREQUENCY) + ' seconds')
		s = sched.scheduler(time.time, time.sleep)
		s.enter(CHECK_FREQUENCY, 1, c.doChecks, (s,))
		s.run()

# Control of daemon
if __name__ == '__main__':
	# Logging
	if DEBUG_MODE:
		logging.basicConfig(filename='/tmp/sd-agent.log', filemode='w', level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

	mainLogger = logging.getLogger('main')
	mainLogger.debug('Agent called')

	# Daemon instance from agent class
	daemon = agent('/tmp/sd-agent.pid')

	# Control options
	if len(sys.argv) == 2:
		if 'start' == sys.argv[1]:
			mainLogger.debug('Start daemon')
			daemon.start()

		elif 'stop' == sys.argv[1]:
			mainLogger.debug('Stop daemon')
			daemon.stop()

		elif 'restart' == sys.argv[1]:
			mainLogger.debug('Restart daemon')
			daemon.restart()

		elif 'update' == sys.argv[1]:
			mainLogger.debug('Updating agent')

			import httplib
			import platform
			import urllib2

			print 'Checking if there is a new version';

			# Get the latest version info
			try:
				mainLogger.debug('Update: checking for update')

				request = urllib2.urlopen('http://www.serverdensity.com/agentupdate/')
				response = request.read()

			except urllib2.HTTPError, e:
				print 'Unable to get latest version info - HTTPError = ' + str(e.reason)

			except urllib2.URLError, e:
				print 'Unable to get latest version info - URLError = ' + str(e.reason)

			except httplib.HTTPException, e:
				print 'Unable to get latest version info - HTTPException'

			except Exception, e:
				import traceback
				print 'Unable to get latest version info - Exception = ' + traceback.format_exc()

			mainLogger.debug('Update: importing json/minjson')

			# We need to return the data using JSON. As of Python 2.6+, there is a core JSON
			# module. We have a 2.4/2.5 compatible lib included with the agent but if we're
			# on 2.6 or above, we should use the core module which will be faster
			pythonVersion = platform.python_version_tuple()

			# Decode the JSON
			if int(pythonVersion[1]) >= 6: # Don't bother checking major version since we only support v2 anyway
				import json

				mainLogger.debug('Update: decoding JSON (json)')

				updateInfo = json.loads(response)

			else:
				import minjson

				mainLogger.debug('Update: decoding JSON minjson')

				updateInfo = minjson.safeRead(response)

			# Do the version check
			if updateInfo['version'] != VERSION:
				import urllib

				print 'A new version is available. Downloading...'

				# Loop through the new files and download each, overwriting the existing one
				for file in updateInfo['files']:
					mainLogger.debug('Update: downloading ' + file)

					print file

					downloadedFile = urllib.urlretrieve('http://www.serverdensity.com/downloads/sd-agent/' + file, file)

				mainLogger.debug('Update: done')

				print 'Update completed. Please restart the agent.'

			else:
				print 'The agent is already up to date'

		else:
			print 'Unknown command'
			sys.exit(2)
		sys.exit(0)

	else:
		print 'usage: %s start|stop|restart|update' % sys.argv[0]
		sys.exit(2)<commit_msg>AddedMD5 checksum checking to update mechanism<commit_after>'''
	Server Density
	www.serverdensity.com
	----
	A web based server resource monitoring application

	Licensed under Simplified BSD License (see LICENSE)
	(C) Boxed Ice 2009 all rights reserved
'''

# General config
DEBUG_MODE = 0
CHECK_FREQUENCY = 60

VERSION = '1.0.0b4'

# Core modules
import ConfigParser
import logging
import sched
import time
import sys

# Check we're not using an old version of Python. We need 2.4 above because some modules (like subprocess)
# were only introduced in 2.4.
if int(sys.version_info[1]) <= 4:
	print 'You are using an outdated version of Python. Please update to v2.4 or above (v3 is not supported).'
	sys.exit(2)

# Custom modules
from checks import checks
from daemon import Daemon

# Config handling
try:
	config = ConfigParser.ConfigParser()
	config.read('config.cfg')
	SD_URL = config.get('Main', 'sd_url')
	AGENT_KEY = config.get('Main', 'agent_key')

except ConfigParser.NoSectionError, e:
	print 'Config file not found or incorrectly formatted'
	quit()

except ConfigParser.ParsingError, e:
	print 'Config file not found or incorrectly formatted'
	quit()

# Override the generic daemon class to run our checks
class agent(Daemon):

	def run(self):
		agentLogger = logging.getLogger('agent')
		agentLogger.debug('Creating checks instance')

		# Checks instance
		c = checks(SD_URL, AGENT_KEY, CHECK_FREQUENCY, VERSION, DEBUG_MODE)

		# Schedule the checks
		agentLogger.debug('Scheduling checks every ' + str(CHECK_FREQUENCY) + ' seconds')
		s = sched.scheduler(time.time, time.sleep)
		s.enter(CHECK_FREQUENCY, 1, c.doChecks, (s,))
		s.run()

# Control of daemon
if __name__ == '__main__':
	# Logging
	if DEBUG_MODE:
		logging.basicConfig(filename='/tmp/sd-agent.log', filemode='w', level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

	mainLogger = logging.getLogger('main')
	mainLogger.debug('Agent called')

	# Daemon instance from agent class
	daemon = agent('/tmp/sd-agent.pid')

	# Control options
	if len(sys.argv) == 2:
		if 'start' == sys.argv[1]:
			mainLogger.debug('Start daemon')
			daemon.start()

		elif 'stop' == sys.argv[1]:
			mainLogger.debug('Stop daemon')
			daemon.stop()

		elif 'restart' == sys.argv[1]:
			mainLogger.debug('Restart daemon')
			daemon.restart()

		elif 'update' == sys.argv[1]:
			mainLogger.debug('Updating agent')

			import httplib
			import platform
			import urllib2

			print 'Checking if there is a new version';

			# Get the latest version info
			try:
				mainLogger.debug('Update: checking for update')

				request = urllib2.urlopen('http://www.serverdensity.com/agentupdate/')
				response = request.read()

			except urllib2.HTTPError, e:
				print 'Unable to get latest version info - HTTPError = ' + str(e.reason)

			except urllib2.URLError, e:
				print 'Unable to get latest version info - URLError = ' + str(e.reason)

			except httplib.HTTPException, e:
				print 'Unable to get latest version info - HTTPException'

			except Exception, e:
				import traceback
				print 'Unable to get latest version info - Exception = ' + traceback.format_exc()

			mainLogger.debug('Update: importing json/minjson')

			# We need to return the data using JSON. As of Python 2.6+, there is a core JSON
			# module. We have a 2.4/2.5 compatible lib included with the agent but if we're
			# on 2.6 or above, we should use the core module which will be faster
			pythonVersion = platform.python_version_tuple()

			# Decode the JSON
			if int(pythonVersion[1]) >= 6: # Don't bother checking major version since we only support v2 anyway
				import json

				mainLogger.debug('Update: decoding JSON (json)')

				try:
					updateInfo = json.loads(response)
				except Exception, e:
					print 'Unable to get latest version info. Try again later.'
					quit()

			else:
				import minjson

				mainLogger.debug('Update: decoding JSON minjson')

				try:
					updateInfo = minjson.safeRead(response)
				except Exception, e:
					print 'Unable to get latest version info. Try again later.'
					quit()

			# Do the version check
			if updateInfo['version'] != VERSION:
				import md5
				import urllib

				print 'A new version is available.'

				def downloadFile(agentFile, recursed = False):
					mainLogger.debug('Update: downloading ' + agentFile['name'])
					print 'Downloading ' + agentFile['name']

					downloadedFile = urllib.urlretrieve('http://www.serverdensity.com/downloads/sd-agent/' + agentFile['name'])

					# Do md5 check to make sure the file downloaded properly
					checksum = md5.new()
					f = file(downloadedFile[0], 'rb')

					# Although the files are small, we can't guarantee the available memory nor that there
					# won't be large files in the future, so read the file in small parts (1kb at time)
					while True:
						part = f.read(1024)

						if not part:
							break # end of file

						checksum.update(part)

					f.close()

					# Do we have a match?
					if checksum.hexdigest() == agentFile['md5']:
						return downloadedFile[0]

					else:
						# Try once more
						if recursed == False:
							downloadFile(agentFile, True)

						else:
							print agentFile['name'] + ' did not match its checksum - it is corrupted. This may be caused by network issues so please try again in a moment.'
							quit()

				# Loop through the new files and call the download function
				for agentFile in updateInfo['files']:
					agentFile['tempFile'] = downloadFile(agentFile)

				# If we got to here then everything worked out fine. However, all the files are still in temporary locations so we need to move them
				# This is to stop an update breaking a working agent if the update fails halfway through
				import os

				for agentFile in updateInfo['files']:
					mainLogger.debug('Update: updating ' + agentFile['name'])
					print 'Updating ' + agentFile['name']

					if os.path.exists(agentFile['name']):
						os.remove(agentFile['name'])

					os.rename(agentFile['tempFile'], agentFile['name'])

				mainLogger.debug('Update: done')

				print 'Update completed. Please restart the agent (python agent.py restart).'

			else:
				print 'The agent is already up to date'

		else:
			print 'Unknown command'
			sys.exit(2)
		sys.exit(0)

	else:
		print 'usage: %s start|stop|restart|update' % sys.argv[0]
		sys.exit(2)<|endoftext|>


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

import os
import re
import subprocess
import sys
import threading
import time
import tempfile
import platform

from nimp.utilities.logging import *

if os.name is "nt":
    from nimp.utilities.windows_utilities import *

#-------------------------------------------------------------------------------
def _default_log_callback(line, default_log_function):
    default_log_function(line)

#-------------------------------------------------------------------------------
def capture_process_output(directory, command, input = None):
    process = subprocess.Popen(command,
                               cwd    = directory,
                               stdout = subprocess.PIPE,
                               stderr = subprocess.PIPE,
                               stdin  = subprocess.PIPE)
    if input is not None:
        input = input.encode("cp437")
    output, error = process.communicate(input)

    # Bonjour toi qui te tappe un bug parce que j'ai fait le connard et mis
    # l'encodage en dur.
    return process.wait(), output.decode("cp437"), error.decode("cp437")

#-------------------------------------------------------------------------------
def call_process(directory, command, log_callback = _default_log_callback):
    log_verbose("Running {0} in directory {1}", command, directory)

    if os.name is "nt":
        ods_logger = OutputDebugStringLogger()

    process = subprocess.Popen(command,
                               stdin                = None,
                               bufsize              = 0,
                               stdout               = subprocess.PIPE,
                               stderr               = subprocess.PIPE,
                               cwd                  = directory)
    if os.name is "nt":
        ods_logger.attach(process.pid)
        ods_logger.start()

    def log_output(log_function, pipe):
        output_buffer = ""
        while process.poll() == None:
            try:
                for line in iter(pipe.readline, ''):
                    line  = line.decode("cp850")

                    if line == '':
                        break

                    line  = line.replace("{", "{{").replace("}", "}}")
                    line  = line.rstrip('\r\n')

                    log_callback(line, log_function)
            except ValueError:
                return

    log_thread_args = [ (log_verbose, process.stdout), (log_error, process.stderr) ]

    if os.name is "nt":
        log_thread_args += [(log_verbose, ods_logger.output)]

    log_threads     = [ threading.Thread(target = log_output, args = args) for args in log_thread_args ]

    for thread in log_threads:
        thread.start()

    process_return = process.wait()

    if os.name == "nt":
        ods_logger.stop()

    for thread in log_threads:
        thread.join()

    log_verbose("Program returned with code {0}", process_return)
    return process_return
<commit_msg>utilities: work around MSYS2 path substitution.<commit_after># -*- coding: utf-8 -*-

import os
import re
import subprocess
import sys
import threading
import time
import tempfile
import platform

from nimp.utilities.logging import *

if os.name is "nt":
    from nimp.utilities.windows_utilities import *

#-------------------------------------------------------------------------------
def _default_log_callback(line, default_log_function):
    default_log_function(line)

#-------------------------------------------------------------------------------
def capture_process_output(directory, command, input = None):
    process = subprocess.Popen(command,
                               cwd    = directory,
                               stdout = subprocess.PIPE,
                               stderr = subprocess.PIPE,
                               stdin  = subprocess.PIPE)
    if input is not None:
        input = input.encode("cp437")
    output, error = process.communicate(input)

    # Bonjour toi qui te tappe un bug parce que j'ai fait le connard et mis
    # l'encodage en dur.
    return process.wait(), output.decode("cp437"), error.decode("cp437")

#-------------------------------------------------------------------------------
def call_process(directory, command, log_callback = _default_log_callback):
    log_verbose("Running {0} in directory {1}", command, directory)

    if os.name is "nt":
        ods_logger = OutputDebugStringLogger()

    # If we’re running under MSYS, leading slashes in command line arguments
    # will be treated as a path, so we need to escape them
    if os.environ['MSYSTEM'] == 'MSYS':
        command = [re.sub('^/', '//', x) for x in command]

    process = subprocess.Popen(command,
                               stdin                = None,
                               bufsize              = 0,
                               stdout               = subprocess.PIPE,
                               stderr               = subprocess.PIPE,
                               cwd                  = directory)
    if os.name is "nt":
        ods_logger.attach(process.pid)
        ods_logger.start()

    def log_output(log_function, pipe):
        output_buffer = ""
        while process.poll() == None:
            try:
                for line in iter(pipe.readline, ''):
                    line  = line.decode("cp850")

                    if line == '':
                        break

                    line  = line.replace("{", "{{").replace("}", "}}")
                    line  = line.rstrip('\r\n')

                    log_callback(line, log_function)
            except ValueError:
                return

    log_thread_args = [ (log_verbose, process.stdout), (log_error, process.stderr) ]

    if os.name is "nt":
        log_thread_args += [(log_verbose, ods_logger.output)]

    log_threads     = [ threading.Thread(target = log_output, args = args) for args in log_thread_args ]

    for thread in log_threads:
        thread.start()

    process_return = process.wait()

    if os.name == "nt":
        ods_logger.stop()

    for thread in log_threads:
        thread.join()

    log_verbose("Program returned with code {0}", process_return)
    return process_return
<|endoftext|>


<commit_before># Copyright (c) 2014, German Neuroinformatics Node (G-Node)
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted under the terms of the BSD License. See
# LICENSE file in the root of the Project.

import unittest
import sys

from nix import *

class TestDataArray(unittest.TestCase):

    def setUp(self):
        self.file  = File.open("unittest.h5", FileMode.Overwrite)
        self.block = self.file.create_block("test block", "recordingsession")
        self.array = self.block.create_data_array("test array", "signal", DataType.Double, (100, ))
        self.other = self.block.create_data_array("other array", "signal", DataType.Double, (100, ))

    def tearDown(self):
        del self.file.blocks[self.block.id]
        self.file.close()

    def test_data_array_eq(self):
        assert(self.array == self.array)
        assert(not self.array == self.other)
        assert(not self.array == None)

    def test_data_array_id(self):
        assert(self.array.id is not None)

    def test_data_array_name(self):
        assert(self.array.name is not None)

    def test_data_array_type(self):
        def set_none():
            self.array.type = None

        assert(self.array.type is not None)
        self.assertRaises(Exception, set_none)

        self.array.type = "foo type"
        assert(self.array.type == "foo type")

    def test_data_array_definition(self):
        assert(self.array.definition is None)

        self.array.definition = "definition"
        assert(self.array.definition == "definition")

        self.array.definition = None
        assert(self.array.definition is None)

    def test_data_array_timestamps(self):
        created_at = self.array.created_at
        assert(created_at > 0)

        updated_at = self.array.updated_at
        assert(updated_at > 0)

        self.array.force_created_at(1403530068)
        assert(self.array.created_at == 1403530068)

    def test_data_array_label(self):
        assert(self.array.label is None)

        self.array.label = "label"
        assert(self.array.label == "label")

        self.array.label = None
        assert(self.array.label is None)

    def test_data_array_unit(self):
        assert(self.array.unit is None)

        self.array.unit = "mV"
        assert(self.array.unit == "mV")

        self.array.unit = None
        assert(self.array.unit is None)

    def test_data_array_exp_origin(self):
        assert(self.array.expansion_origin is None)

        self.array.expansion_origin = 10.2
        assert(self.array.expansion_origin == 10.2)

        self.array.expansion_origin = None
        assert(self.array.expansion_origin is None)

    def test_data_array_coefficients(self):
        assert(self.array.polynom_coefficients == ())

        self.array.polynom_coefficients = (1.1, 2.2)
        assert(self.array.polynom_coefficients == (1.1, 2.2))

        # TODO delete does not work

    def test_data_array_data(self):
        import numpy as np

        assert(self.array.polynom_coefficients == ())

        data = np.array([float(i) for i in range(100)])
        dout = np.empty_like(data)
        self.array.data.write_direct(data)
        assert(self.array.data.dtype == np.dtype(float))
        self.array.data.read_direct(dout)
        assert(np.array_equal(data, dout))
        dout = np.array(self.array.data)
        assert(np.array_equal(data, dout))
        assert(self.array.data_extent == data.shape)
        assert(self.array.data_extent == self.array.data.shape)
        assert(self.array.data.size == data.size)

        assert(len(self.array.data) == len(data))

        #indexing support in 1-d arrays
        self.assertRaises(IndexError, lambda : self.array.data[1:4:5])
        self.assertRaises(IndexError, lambda : self.array.data[[1,3,]])

        dout = np.array([self.array.data[i] for i in range(100)])
        assert(np.array_equal(data, dout))

        dout = self.array.data[...]
        assert(np.array_equal(data, dout))

        #indexed writing (1-d)
        data = np.array([float(-i) for i in range(100)])
        self.array.data[()] = data
        assert(np.array_equal(self.array.data[...], data))

        self.array.data[...] = [float(-i) for i in range(100)]
        assert(np.array_equal(self.array.data[()], data))

        self.array.data[0] = 42
        assert(self.array.data[0] == 42.0)

        #changing shape via data_extent property
        self.array.data_extent = (200, )
        assert(self.array.data_extent == (200, ))

        # TODO delete does not work
        data = np.eye(123)
        a1 = self.block.create_data_array("double array", "signal", DataType.Double, (123, 123))
        dset = a1.data
        dset.write_direct(data)
        dout = np.empty_like(data)
        dset.read_direct(dout)
        assert(np.array_equal(data, dout))

        #indexing support in 2-d arrays
        self.assertRaises(IndexError, lambda : self.array.data[[], [1,2]])

        dout = dset[12]
        assert(dout.shape == data[12].shape)
        assert(np.array_equal(dout, data[12]))
        assert(np.array_equal(dset[()], data))
        assert(np.array_equal(dset[...], data))
        assert(np.array_equal(dset[12, ...], data[12, ...]))
        assert(np.array_equal(dset[..., 12], data[..., 12]))
        assert(np.array_equal(dset[1:], data[1:]))
        assert(np.array_equal(dset[:1], data[:1]))
        assert(np.array_equal(dset[1:10, 1:10], data[1:10, 1:10]))

        a3 = self.block.create_data_array("int identity array", "signal", DataType.Int32, (123, 123))
        assert(a3.data_extent == (123, 123))
        assert(a3.data.dtype == np.dtype('i4'))

        data = np.random.rand(3, 4, 5)
        a4 = self.block.create_data_array("3d array", "signal", DataType.Double, (3, 4, 5))
        dset = a4.data
        dset.write_direct(data)
        assert(dset.shape == data.shape)
        assert(len(dset) == len(data))
        assert(dset.size == data.size)
        assert(np.array_equal(dset[2, ...], data[2, ...]))
        assert(np.array_equal(dset[..., 3], data[..., 3]))
        assert(np.array_equal(dset[2, ..., 3], data[2, ..., 3]))
        assert(np.array_equal(dset[1:2, ..., 3:5], data[1:2, ..., 3:5]))

        #indexed writing (n-d)
        d2 = np.random.rand(2,2)
        dset[1, 0:2, 0:2] = d2
        assert(np.array_equal(dset[1, 0:2, 0:2], d2))

        # test for the size check in DataSet.__len__
        # by simulating a system with a really smal int
        savemaxsize = sys.maxsize
        sys.maxsize = len(dset) - 1
        self.assertRaises(OverflowError, lambda : len(dset))
        sys.maxsize = savemaxsize

    def test_data_array_dimensions(self):
        assert(len(self.array.dimensions) == 0)

        setd    = self.array.append_set_dimension()
        ranged  = self.array.append_range_dimension(range(10))
        sampled = self.array.append_sampled_dimension(0.1)

        assert(len(self.array.dimensions) == 3)

        self.assertRaises(TypeError, lambda : self.array.dimensions["notexist"])
        self.assertRaises(KeyError, lambda : self.array.dimensions[-4])
        self.assertRaises(KeyError, lambda : self.array.dimensions[3])

        assert(isinstance(str(self.array.dimensions), basestring))
        assert(isinstance(repr(self.array.dimensions), basestring))

        dims   = list(self.array.dimensions)
        for i in range(3):
            assert(dims[i].index == self.array.dimensions[i].index)
            assert(dims[i].dimension_type == self.array.dimensions[i].dimension_type)

            assert(self.array.dimensions[i].index == self.array.dimensions[i - 3].index)

        del self.array.dimensions[2]
        del self.array.dimensions[1]
        del self.array.dimensions[0]

        assert(len(self.array.dimensions) == 0)

    def test_data_array_sources(self):
        source1 = self.block.create_source("source1", "channel")
        source2 = self.block.create_source("source2", "electrode")

        assert(len(self.array.sources) == 0)

        self.array.sources.append(source1)
        self.array.sources.append(source2)

        self.assertRaises(TypeError, lambda : self.array.sources.append(100))

        assert(len(self.array.sources) == 2)
        assert(source1 in self.array.sources)
        assert(source2 in self.array.sources)

        del self.array.sources[source2]
        assert(self.array.sources[0] == source1)

        del self.array.sources[source1]
        assert(len(self.array.sources) == 0)
<commit_msg>[test] DataArray: test dtype transmogrification<commit_after># Copyright (c) 2014, German Neuroinformatics Node (G-Node)
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted under the terms of the BSD License. See
# LICENSE file in the root of the Project.

import unittest
import sys
import numpy as np

from nix import *

class TestDataArray(unittest.TestCase):

    def setUp(self):
        self.file  = File.open("unittest.h5", FileMode.Overwrite)
        self.block = self.file.create_block("test block", "recordingsession")
        self.array = self.block.create_data_array("test array", "signal", DataType.Double, (100, ))
        self.other = self.block.create_data_array("other array", "signal", DataType.Double, (100, ))

    def tearDown(self):
        del self.file.blocks[self.block.id]
        self.file.close()

    def test_data_array_eq(self):
        assert(self.array == self.array)
        assert(not self.array == self.other)
        assert(not self.array == None)

    def test_data_array_id(self):
        assert(self.array.id is not None)

    def test_data_array_name(self):
        assert(self.array.name is not None)

    def test_data_array_type(self):
        def set_none():
            self.array.type = None

        assert(self.array.type is not None)
        self.assertRaises(Exception, set_none)

        self.array.type = "foo type"
        assert(self.array.type == "foo type")

    def test_data_array_definition(self):
        assert(self.array.definition is None)

        self.array.definition = "definition"
        assert(self.array.definition == "definition")

        self.array.definition = None
        assert(self.array.definition is None)

    def test_data_array_timestamps(self):
        created_at = self.array.created_at
        assert(created_at > 0)

        updated_at = self.array.updated_at
        assert(updated_at > 0)

        self.array.force_created_at(1403530068)
        assert(self.array.created_at == 1403530068)

    def test_data_array_label(self):
        assert(self.array.label is None)

        self.array.label = "label"
        assert(self.array.label == "label")

        self.array.label = None
        assert(self.array.label is None)

    def test_data_array_unit(self):
        assert(self.array.unit is None)

        self.array.unit = "mV"
        assert(self.array.unit == "mV")

        self.array.unit = None
        assert(self.array.unit is None)

    def test_data_array_exp_origin(self):
        assert(self.array.expansion_origin is None)

        self.array.expansion_origin = 10.2
        assert(self.array.expansion_origin == 10.2)

        self.array.expansion_origin = None
        assert(self.array.expansion_origin is None)

    def test_data_array_coefficients(self):
        assert(self.array.polynom_coefficients == ())

        self.array.polynom_coefficients = (1.1, 2.2)
        assert(self.array.polynom_coefficients == (1.1, 2.2))

        # TODO delete does not work

    def test_data_array_data(self):

        assert(self.array.polynom_coefficients == ())

        data = np.array([float(i) for i in range(100)])
        dout = np.empty_like(data)
        self.array.data.write_direct(data)
        assert(self.array.data.dtype == np.dtype(float))
        self.array.data.read_direct(dout)
        assert(np.array_equal(data, dout))
        dout = np.array(self.array.data)
        assert(np.array_equal(data, dout))
        assert(self.array.data_extent == data.shape)
        assert(self.array.data_extent == self.array.data.shape)
        assert(self.array.data.size == data.size)

        assert(len(self.array.data) == len(data))

        #indexing support in 1-d arrays
        self.assertRaises(IndexError, lambda : self.array.data[1:4:5])
        self.assertRaises(IndexError, lambda : self.array.data[[1,3,]])

        dout = np.array([self.array.data[i] for i in range(100)])
        assert(np.array_equal(data, dout))

        dout = self.array.data[...]
        assert(np.array_equal(data, dout))

        #indexed writing (1-d)
        data = np.array([float(-i) for i in range(100)])
        self.array.data[()] = data
        assert(np.array_equal(self.array.data[...], data))

        self.array.data[...] = [float(-i) for i in range(100)]
        assert(np.array_equal(self.array.data[()], data))

        self.array.data[0] = 42
        assert(self.array.data[0] == 42.0)

        #changing shape via data_extent property
        self.array.data_extent = (200, )
        assert(self.array.data_extent == (200, ))

        # TODO delete does not work
        data = np.eye(123)
        a1 = self.block.create_data_array("double array", "signal", DataType.Double, (123, 123))
        dset = a1.data
        dset.write_direct(data)
        dout = np.empty_like(data)
        dset.read_direct(dout)
        assert(np.array_equal(data, dout))

        #indexing support in 2-d arrays
        self.assertRaises(IndexError, lambda : self.array.data[[], [1,2]])

        dout = dset[12]
        assert(dout.shape == data[12].shape)
        assert(np.array_equal(dout, data[12]))
        assert(np.array_equal(dset[()], data))
        assert(np.array_equal(dset[...], data))
        assert(np.array_equal(dset[12, ...], data[12, ...]))
        assert(np.array_equal(dset[..., 12], data[..., 12]))
        assert(np.array_equal(dset[1:], data[1:]))
        assert(np.array_equal(dset[:1], data[:1]))
        assert(np.array_equal(dset[1:10, 1:10], data[1:10, 1:10]))

        a3 = self.block.create_data_array("int identity array", "signal", DataType.Int32, (123, 123))
        assert(a3.data_extent == (123, 123))
        assert(a3.data.dtype == np.dtype('i4'))

        data = np.random.rand(3, 4, 5)
        a4 = self.block.create_data_array("3d array", "signal", DataType.Double, (3, 4, 5))
        dset = a4.data
        dset.write_direct(data)
        assert(dset.shape == data.shape)
        assert(len(dset) == len(data))
        assert(dset.size == data.size)
        assert(np.array_equal(dset[2, ...], data[2, ...]))
        assert(np.array_equal(dset[..., 3], data[..., 3]))
        assert(np.array_equal(dset[2, ..., 3], data[2, ..., 3]))
        assert(np.array_equal(dset[1:2, ..., 3:5], data[1:2, ..., 3:5]))

        #indexed writing (n-d)
        d2 = np.random.rand(2,2)
        dset[1, 0:2, 0:2] = d2
        assert(np.array_equal(dset[1, 0:2, 0:2], d2))

        # test for the size check in DataSet.__len__
        # by simulating a system with a really smal int
        savemaxsize = sys.maxsize
        sys.maxsize = len(dset) - 1
        self.assertRaises(OverflowError, lambda : len(dset))
        sys.maxsize = savemaxsize

    def test_data_array_dtype(self):
        da = self.block.create_data_array('dtype_f8', 'b', 'f8', (10, 10))
        assert(da.data.dtype == np.dtype('f8'))

        da = self.block.create_data_array('dtype_i16', 'b', np.int16, (10, 10))
        assert(da.data.dtype == np.int16)

        da = self.block.create_data_array('dtype_int', 'b', int, (10, 10))
        assert(da.data.dtype == np.dtype(int))

        da = self.block.create_data_array('dtype_ndouble', 'b', DataType.Double, (10, 10))
        assert(da.data.dtype == np.dtype('f8'))

    def test_data_array_dimensions(self):
        assert(len(self.array.dimensions) == 0)

        setd    = self.array.append_set_dimension()
        ranged  = self.array.append_range_dimension(range(10))
        sampled = self.array.append_sampled_dimension(0.1)

        assert(len(self.array.dimensions) == 3)

        self.assertRaises(TypeError, lambda : self.array.dimensions["notexist"])
        self.assertRaises(KeyError, lambda : self.array.dimensions[-4])
        self.assertRaises(KeyError, lambda : self.array.dimensions[3])

        assert(isinstance(str(self.array.dimensions), basestring))
        assert(isinstance(repr(self.array.dimensions), basestring))

        dims   = list(self.array.dimensions)
        for i in range(3):
            assert(dims[i].index == self.array.dimensions[i].index)
            assert(dims[i].dimension_type == self.array.dimensions[i].dimension_type)

            assert(self.array.dimensions[i].index == self.array.dimensions[i - 3].index)

        del self.array.dimensions[2]
        del self.array.dimensions[1]
        del self.array.dimensions[0]

        assert(len(self.array.dimensions) == 0)

    def test_data_array_sources(self):
        source1 = self.block.create_source("source1", "channel")
        source2 = self.block.create_source("source2", "electrode")

        assert(len(self.array.sources) == 0)

        self.array.sources.append(source1)
        self.array.sources.append(source2)

        self.assertRaises(TypeError, lambda : self.array.sources.append(100))

        assert(len(self.array.sources) == 2)
        assert(source1 in self.array.sources)
        assert(source2 in self.array.sources)

        del self.array.sources[source2]
        assert(self.array.sources[0] == source1)

        del self.array.sources[source1]
        assert(len(self.array.sources) == 0)
<|endoftext|>


<commit_before>from collections import defaultdict

from django import forms
from django.conf import settings
from django.conf.urls import url
from django.contrib import admin, messages
from django.contrib.admin import forms as admin_forms
from django.contrib.auth import admin as auth_admin, get_user_model
from django.contrib.auth.models import Group
from django.core.exceptions import ValidationError
from django.urls import reverse
from django.utils.html import format_html_join
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext_lazy as _
from rest_framework import permissions as rf_permissions
from reversion.admin import VersionAdmin

from nodeconductor.core import models
from nodeconductor.core.authentication import can_access_admin_site

def get_admin_url(obj):
    return reverse('admin:%s_%s_change' % (obj._meta.app_label, obj._meta.model_name), args=[obj.id])

def render_to_readonly(value):
    return "
{0}
".format(value)

class ReadonlyTextWidget(forms.TextInput):
    def format_value(self, value):
        return value

    def render(self, name, value, attrs=None, renderer=None):
        return render_to_readonly(self.format_value(value))

class OptionalChoiceField(forms.ChoiceField):
    def __init__(self, choices=(), *args, **kwargs):
        empty = [('', '---------')]
        choices = empty + sorted(choices, key=lambda (code, label): label)
        super(OptionalChoiceField, self).__init__(choices, *args, **kwargs)

class UserCreationForm(auth_admin.UserCreationForm):
    class Meta(object):
        model = get_user_model()
        fields = ("username",)

    # overwritten to support custom User model
    def clean_username(self):
        # Since User.username is unique, this check is redundant,
        # but it sets a nicer error message than the ORM. See #13147.
        username = self.cleaned_data["username"]
        try:
            get_user_model()._default_manager.get(username=username)
        except get_user_model().DoesNotExist:
            return username
        raise forms.ValidationError(
            self.error_messages['duplicate_username'],
            code='duplicate_username',
        )

class UserChangeForm(auth_admin.UserChangeForm):
    class Meta(object):
        model = get_user_model()
        fields = '__all__'

    def __init__(self, *args, **kwargs):
        super(UserChangeForm, self).__init__(*args, **kwargs)
        competences = [(key, key) for key in settings.NODECONDUCTOR.get('USER_COMPETENCE_LIST', [])]
        self.fields['preferred_language'] = OptionalChoiceField(choices=settings.LANGUAGES, required=False)
        self.fields['competence'] = OptionalChoiceField(choices=competences, required=False)

    def clean_civil_number(self):
        # Empty string should be converted to None.
        # Otherwise uniqueness constraint is violated.
        # See also: http://stackoverflow.com/a/1400046/175349
        civil_number = self.cleaned_data.get('civil_number')
        if civil_number:
            return civil_number.strip()
        return None

class UserAdmin(auth_admin.UserAdmin):
    list_display = ('username', 'uuid', 'email', 'full_name', 'native_name', 'is_active', 'is_staff', 'is_support')
    search_fields = ('username', 'uuid', 'full_name', 'native_name', 'email', 'civil_number')
    list_filter = ('is_active', 'is_staff', 'is_support')
    fieldsets = (
        (None, {'fields': ('username', 'password', 'registration_method')}),
        (_('Personal info'), {'fields': (
            'civil_number', 'full_name', 'native_name', 'email',
            'preferred_language', 'competence', 'phone_number'
        )}),
        (_('Organization'), {'fields': ('organization', 'organization_approved')}),
        (_('Permissions'), {'fields': ('is_active', 'is_staff', 'is_support', 'customer_roles', 'project_roles')}),
        (_('Important dates'), {'fields': ('last_login', 'date_joined', 'agreement_date')}),
    )
    readonly_fields = ('registration_method', 'agreement_date', 'customer_roles', 'project_roles')
    form = UserChangeForm
    add_form = UserCreationForm

    def customer_roles(self, instance):
        from nodeconductor.structure.models import CustomerPermission
        customer_permissions = CustomerPermission.objects.filter(user=instance, is_active=True).order_by('customer')

        return format_html_join(
            mark_safe('
'),
            '{}',
            ((str(permission),) for permission in customer_permissions),
        ) or mark_safe("%s" % _('User has no roles in any organization.'))

    customer_roles.short_description = _('Roles in organizations')

    def project_roles(self, instance):
        from nodeconductor.structure.models import ProjectPermission
        permissions = ProjectPermission.objects.filter(user=instance, is_active=True).order_by('project')

        return format_html_join(
            mark_safe('
'),
            '{}',
            ((str(line),) for line in permissions),
        ) or mark_safe("%s" % _('User has no roles in any project.'))

    project_roles.short_description = _('Roles in projects')

class SshPublicKeyAdmin(admin.ModelAdmin):
    list_display = ('user', 'name', 'fingerprint')
    search_fields = ('user', 'name', 'fingerprint')
    readonly_fields = ('user', 'name', 'fingerprint', 'public_key')

class CustomAdminAuthenticationForm(admin_forms.AdminAuthenticationForm):
    error_messages = {
        'invalid_login': _("Please enter the correct %(username)s and password "
                           "for a staff or a support account. Note that both fields may be "
                           "case-sensitive."),
    }

    def confirm_login_allowed(self, user):
        if not can_access_admin_site(user):
            return super(CustomAdminAuthenticationForm, self).confirm_login_allowed(user)

class CustomAdminSite(admin.AdminSite):
    site_title = _('Waldur MasterMind admin')
    site_header = _('Waldur MasterMind administration')
    index_title = _('Waldur MasterMind administration')
    login_form = CustomAdminAuthenticationForm

    def has_permission(self, request):
        is_safe = request.method in rf_permissions.SAFE_METHODS
        return can_access_admin_site(request.user) and (is_safe or request.user.is_staff)

    @classmethod
    def clone_default(cls):
        instance = cls()
        instance._registry = admin.site._registry.copy()
        instance._actions = admin.site._actions.copy()
        instance._global_actions = admin.site._global_actions.copy()
        return instance

admin_site = CustomAdminSite.clone_default()
admin.site = admin_site
admin.site.register(models.User, UserAdmin)
admin.site.register(models.SshPublicKey, SshPublicKeyAdmin)

# TODO: Extract common classes to admin_utils module and remove hack.
# This hack is needed because admin is imported several times.
# Please note that admin module should NOT be imported by other apps.
if admin.site.is_registered(Group):
    admin.site.unregister(Group)

class ReversionAdmin(VersionAdmin):
    def add_view(self, request, form_url='', extra_context=None):
        # Revision creation is ignored in this method because it has to be implemented in model.save method
        return super(VersionAdmin, self).add_view(request, form_url, extra_context)

    def change_view(self, request, object_id, form_url='', extra_context=None):
        # Revision creation is ignored in this method because it has to be implemented in model.save method
        return super(VersionAdmin, self).change_view(request, object_id, form_url, extra_context)

class ExecutorAdminAction(object):
    """ Add executor as action to admin model.

    Usage example:
        class PullSecurityGroups(ExecutorAdminAction):
            executor = executors.TenantPullSecurityGroupsExecutor  # define executor
            short_description = 'Pull security groups'  # description for admin page

            def validate(self, tenant):
                if tenant.state != Tenant.States.OK:
                    raise ValidationError('Tenant has to be in state OK to pull security groups.')

        pull_security_groups = PullSecurityGroups()  # this action could be registered as admin action

    """
    executor = NotImplemented

    def __call__(self, admin_class, request, queryset):
        errors = defaultdict(list)
        successfully_executed = []
        for instance in queryset:
            try:
                self.validate(instance)
            except ValidationError as e:
                errors[str(e)].append(instance)
            else:
                self.executor.execute(instance)
                successfully_executed.append(instance)

        if successfully_executed:
            message = _('Operation was successfully scheduled for %(count)d instances: %(names)s') % dict(
                count=len(successfully_executed),
                names=', '.join([str(i) for i in successfully_executed])
            )
            admin_class.message_user(request, message)

        for error, instances in errors.items():
            message = _('Failed to schedule operation for %(count)d instances: %(names)s. Error: %(message)s') % dict(
                count=len(instances),
                names=', '.join([str(i) for i in instances]),
                message=error,
            )
            admin_class.message_user(request, message, level=messages.ERROR)

    def validate(self, instance):
        """ Raise validation error if action cannot be performed for given instance """
        pass

class ExtraActionsMixin(object):
    """
    Allows to add extra actions to admin list page.
    """
    change_list_template = 'admin/core/change_list.html'

    def get_extra_actions(self):
        raise NotImplementedError('Method "get_extra_actions" should be implemented in ExtraActionsMixin.')

    def get_urls(self):
        """
        Inject extra action URLs.
        """
        urls = []

        for action in self.get_extra_actions():
            regex = r'^{}/$'.format(self._get_action_href(action))
            view = self.admin_site.admin_view(action)
            urls.append(url(regex, view))

        return urls + super(ExtraActionsMixin, self).get_urls()

    def changelist_view(self, request, extra_context=None):
        """
        Inject extra links into template context.
        """
        links = []

        for action in self.get_extra_actions():
            links.append({
                'label': self._get_action_label(action),
                'href': self._get_action_href(action)
            })

        extra_context = extra_context or {}
        extra_context['extra_links'] = links

        return super(ExtraActionsMixin, self).changelist_view(
            request, extra_context=extra_context,
        )

    def _get_action_href(self, action):
        return action.__name__

    def _get_action_label(self, action):
        return getattr(action, 'name', action.__name__.replace('_', ' ').capitalize())
<commit_msg>Make customer_roles and project_roles consistent in admin<commit_after>from collections import defaultdict

from django import forms
from django.conf import settings
from django.conf.urls import url
from django.contrib import admin, messages
from django.contrib.admin import forms as admin_forms
from django.contrib.auth import admin as auth_admin, get_user_model
from django.contrib.auth.models import Group
from django.core.exceptions import ValidationError
from django.urls import reverse
from django.utils.html import format_html_join
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext_lazy as _
from rest_framework import permissions as rf_permissions
from reversion.admin import VersionAdmin

from nodeconductor.core import models
from nodeconductor.core.authentication import can_access_admin_site

def get_admin_url(obj):
    return reverse('admin:%s_%s_change' % (obj._meta.app_label, obj._meta.model_name), args=[obj.id])

def render_to_readonly(value):
    return "
{0}
".format(value)

class ReadonlyTextWidget(forms.TextInput):
    def format_value(self, value):
        return value

    def render(self, name, value, attrs=None, renderer=None):
        return render_to_readonly(self.format_value(value))

class OptionalChoiceField(forms.ChoiceField):
    def __init__(self, choices=(), *args, **kwargs):
        empty = [('', '---------')]
        choices = empty + sorted(choices, key=lambda (code, label): label)
        super(OptionalChoiceField, self).__init__(choices, *args, **kwargs)

class UserCreationForm(auth_admin.UserCreationForm):
    class Meta(object):
        model = get_user_model()
        fields = ("username",)

    # overwritten to support custom User model
    def clean_username(self):
        # Since User.username is unique, this check is redundant,
        # but it sets a nicer error message than the ORM. See #13147.
        username = self.cleaned_data["username"]
        try:
            get_user_model()._default_manager.get(username=username)
        except get_user_model().DoesNotExist:
            return username
        raise forms.ValidationError(
            self.error_messages['duplicate_username'],
            code='duplicate_username',
        )

class UserChangeForm(auth_admin.UserChangeForm):
    class Meta(object):
        model = get_user_model()
        fields = '__all__'

    def __init__(self, *args, **kwargs):
        super(UserChangeForm, self).__init__(*args, **kwargs)
        competences = [(key, key) for key in settings.NODECONDUCTOR.get('USER_COMPETENCE_LIST', [])]
        self.fields['preferred_language'] = OptionalChoiceField(choices=settings.LANGUAGES, required=False)
        self.fields['competence'] = OptionalChoiceField(choices=competences, required=False)

    def clean_civil_number(self):
        # Empty string should be converted to None.
        # Otherwise uniqueness constraint is violated.
        # See also: http://stackoverflow.com/a/1400046/175349
        civil_number = self.cleaned_data.get('civil_number')
        if civil_number:
            return civil_number.strip()
        return None

class UserAdmin(auth_admin.UserAdmin):
    list_display = ('username', 'uuid', 'email', 'full_name', 'native_name', 'is_active', 'is_staff', 'is_support')
    search_fields = ('username', 'uuid', 'full_name', 'native_name', 'email', 'civil_number')
    list_filter = ('is_active', 'is_staff', 'is_support')
    fieldsets = (
        (None, {'fields': ('username', 'password', 'registration_method')}),
        (_('Personal info'), {'fields': (
            'civil_number', 'full_name', 'native_name', 'email',
            'preferred_language', 'competence', 'phone_number'
        )}),
        (_('Organization'), {'fields': ('organization', 'organization_approved')}),
        (_('Permissions'), {'fields': ('is_active', 'is_staff', 'is_support', 'customer_roles', 'project_roles')}),
        (_('Important dates'), {'fields': ('last_login', 'date_joined', 'agreement_date')}),
    )
    readonly_fields = ('registration_method', 'agreement_date', 'customer_roles', 'project_roles')
    form = UserChangeForm
    add_form = UserCreationForm

    def customer_roles(self, instance):
        from nodeconductor.structure.models import CustomerPermission
        permissions = CustomerPermission.objects.filter(user=instance, is_active=True).order_by('customer')

        return format_html_join(
            mark_safe('
'),
            '{}',
            ((str(permission),) for permission in permissions),
        ) or mark_safe("%s" % _('User has no roles in any organization.'))

    customer_roles.short_description = _('Roles in organizations')

    def project_roles(self, instance):
        from nodeconductor.structure.models import ProjectPermission
        permissions = ProjectPermission.objects.filter(user=instance, is_active=True).order_by('project')

        return format_html_join(
            mark_safe('
'),
            '{}',
            ((str(permission),) for permission in permissions),
        ) or mark_safe("%s" % _('User has no roles in any project.'))

    project_roles.short_description = _('Roles in projects')

class SshPublicKeyAdmin(admin.ModelAdmin):
    list_display = ('user', 'name', 'fingerprint')
    search_fields = ('user', 'name', 'fingerprint')
    readonly_fields = ('user', 'name', 'fingerprint', 'public_key')

class CustomAdminAuthenticationForm(admin_forms.AdminAuthenticationForm):
    error_messages = {
        'invalid_login': _("Please enter the correct %(username)s and password "
                           "for a staff or a support account. Note that both fields may be "
                           "case-sensitive."),
    }

    def confirm_login_allowed(self, user):
        if not can_access_admin_site(user):
            return super(CustomAdminAuthenticationForm, self).confirm_login_allowed(user)

class CustomAdminSite(admin.AdminSite):
    site_title = _('Waldur MasterMind admin')
    site_header = _('Waldur MasterMind administration')
    index_title = _('Waldur MasterMind administration')
    login_form = CustomAdminAuthenticationForm

    def has_permission(self, request):
        is_safe = request.method in rf_permissions.SAFE_METHODS
        return can_access_admin_site(request.user) and (is_safe or request.user.is_staff)

    @classmethod
    def clone_default(cls):
        instance = cls()
        instance._registry = admin.site._registry.copy()
        instance._actions = admin.site._actions.copy()
        instance._global_actions = admin.site._global_actions.copy()
        return instance

admin_site = CustomAdminSite.clone_default()
admin.site = admin_site
admin.site.register(models.User, UserAdmin)
admin.site.register(models.SshPublicKey, SshPublicKeyAdmin)

# TODO: Extract common classes to admin_utils module and remove hack.
# This hack is needed because admin is imported several times.
# Please note that admin module should NOT be imported by other apps.
if admin.site.is_registered(Group):
    admin.site.unregister(Group)

class ReversionAdmin(VersionAdmin):
    def add_view(self, request, form_url='', extra_context=None):
        # Revision creation is ignored in this method because it has to be implemented in model.save method
        return super(VersionAdmin, self).add_view(request, form_url, extra_context)

    def change_view(self, request, object_id, form_url='', extra_context=None):
        # Revision creation is ignored in this method because it has to be implemented in model.save method
        return super(VersionAdmin, self).change_view(request, object_id, form_url, extra_context)

class ExecutorAdminAction(object):
    """ Add executor as action to admin model.

    Usage example:
        class PullSecurityGroups(ExecutorAdminAction):
            executor = executors.TenantPullSecurityGroupsExecutor  # define executor
            short_description = 'Pull security groups'  # description for admin page

            def validate(self, tenant):
                if tenant.state != Tenant.States.OK:
                    raise ValidationError('Tenant has to be in state OK to pull security groups.')

        pull_security_groups = PullSecurityGroups()  # this action could be registered as admin action

    """
    executor = NotImplemented

    def __call__(self, admin_class, request, queryset):
        errors = defaultdict(list)
        successfully_executed = []
        for instance in queryset:
            try:
                self.validate(instance)
            except ValidationError as e:
                errors[str(e)].append(instance)
            else:
                self.executor.execute(instance)
                successfully_executed.append(instance)

        if successfully_executed:
            message = _('Operation was successfully scheduled for %(count)d instances: %(names)s') % dict(
                count=len(successfully_executed),
                names=', '.join([str(i) for i in successfully_executed])
            )
            admin_class.message_user(request, message)

        for error, instances in errors.items():
            message = _('Failed to schedule operation for %(count)d instances: %(names)s. Error: %(message)s') % dict(
                count=len(instances),
                names=', '.join([str(i) for i in instances]),
                message=error,
            )
            admin_class.message_user(request, message, level=messages.ERROR)

    def validate(self, instance):
        """ Raise validation error if action cannot be performed for given instance """
        pass

class ExtraActionsMixin(object):
    """
    Allows to add extra actions to admin list page.
    """
    change_list_template = 'admin/core/change_list.html'

    def get_extra_actions(self):
        raise NotImplementedError('Method "get_extra_actions" should be implemented in ExtraActionsMixin.')

    def get_urls(self):
        """
        Inject extra action URLs.
        """
        urls = []

        for action in self.get_extra_actions():
            regex = r'^{}/$'.format(self._get_action_href(action))
            view = self.admin_site.admin_view(action)
            urls.append(url(regex, view))

        return urls + super(ExtraActionsMixin, self).get_urls()

    def changelist_view(self, request, extra_context=None):
        """
        Inject extra links into template context.
        """
        links = []

        for action in self.get_extra_actions():
            links.append({
                'label': self._get_action_label(action),
                'href': self._get_action_href(action)
            })

        extra_context = extra_context or {}
        extra_context['extra_links'] = links

        return super(ExtraActionsMixin, self).changelist_view(
            request, extra_context=extra_context,
        )

    def _get_action_href(self, action):
        return action.__name__

    def _get_action_label(self, action):
        return getattr(action, 'name', action.__name__.replace('_', ' ').capitalize())
<|endoftext|>