<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# LICENSE
#
# Copyright (C) 2010-2018 GEM Foundation, G. Weatherill, M. Pagani,
# D. Monelli.
#
# The Hazard Modeller's Toolkit is free software: you can redistribute
# it and/or modify it under the terms of the GNU Affero General Public
# License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>
#
# DISCLAIMER
#
# The software Hazard Modeller's Toolkit (openquake.hmtk) provided herein
# is released as a prototype implementation on behalf of
# scientists and engineers working within the GEM Foundation (Global
# Earthquake Model).
#
# It is distributed for the purpose of open collaboration and in the
# hope that it will be useful to the scientific, engineering, disaster
# risk and software design communities.
#
# The software is NOT distributed as part of GEM's OpenQuake suite
# (https://www.globalquakemodel.org/tools-products) and must be considered as a
# separate entity. The software provided herein is designed and implemented
# by scientific staff. It is not developed to the design standards, nor
# subject to same level of critical review by professional software
# developers, as GEM's OpenQuake software suite.
#
# Feedback and contribution to the software is welcome, and can be
# directed to the hazard scientific staff of the GEM Model Facility
# (hazard@globalquakemodel.org).
#
# The Hazard Modeller's Toolkit (openquake.hmtk) is therefore distributed WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# The GEM Foundation, and the authors of the software, assume no
# liability for use of the software.

'''
Module implements :class: openquake.hmtk.sources.source_model.mtkSourceModel, the
general class to describe a set of seismogenic sources
'''
from openquake.hazardlib.tom import PoissonTOM
from openquake.hazardlib.sourcewriter import write_source_model
from openquake.hmtk.sources.point_source import mtkPointSource
from openquake.hmtk.sources.area_source import mtkAreaSource
from openquake.hmtk.sources.simple_fault_source import mtkSimpleFaultSource
from openquake.hmtk.sources.complex_fault_source import mtkComplexFaultSource

class mtkSourceModel(object):
    '''
    Object to describe a seismogenic source model (composite of
    multiple sources with mixed typologies)

    :param str id:
        Identifier for the source model
    :param str name:
        Source model name
    :param list sources:
        List of seismogenic sources
    '''

    def __init__(self, identifier=None, name=None, sources=None):
        self.id = identifier
        self.name = name
        if isinstance(sources, list):
            self.sources = sources
        else:
            if sources:
                raise ValueError('Sources must be input as list!')
            self.sources = []

    def __iter__(self):
        return iter(self.sources)

    def get_number_sources(self):
        '''
        Returns the number of sources in the model
        '''
        return len(self.sources)

    def serialise_to_nrml(self, filename, use_defaults=False):
        '''
        Writes the source model to a nrml source model file given by the
        filename

        :param str filename:
            Path to output file

        :param bool use_defaults:
            Boolean to indicate whether to use default values (True) or not.
            If set to False, ValueErrors will be raised when an essential
            attribute is missing.
        '''
        source_model = self.convert_to_oqhazardlib(
            PoissonTOM(1.0), 2.0, 2.0, 10.0, use_defaults=use_defaults)
        write_source_model(filename, source_model, name=self.name)

    def convert_to_oqhazardlib(
            self, tom, simple_mesh_spacing=1.0,
            complex_mesh_spacing=2.0, area_discretisation=10.0,
            use_defaults=False):
        """
        Converts the source model to an iterator of sources of :class:
        openquake.hazardlib.source.base.BaseSeismicSource
        """
        oq_source_model = []
        for source in self.sources:
            if isinstance(source, mtkAreaSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    simple_mesh_spacing,
                    area_discretisation,
                    use_defaults))
            elif isinstance(source, mtkPointSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    simple_mesh_spacing,
                    use_defaults))
            elif isinstance(source, mtkSimpleFaultSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    simple_mesh_spacing,
                    use_defaults))
            elif isinstance(source, mtkComplexFaultSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    complex_mesh_spacing,
                    use_defaults))
            else:
                raise ValueError('Source type not recognised!')
        return oq_source_model
<commit_msg>Adds container __len__ to mtkSourceModel<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# LICENSE
#
# Copyright (C) 2010-2018 GEM Foundation, G. Weatherill, M. Pagani,
# D. Monelli.
#
# The Hazard Modeller's Toolkit is free software: you can redistribute
# it and/or modify it under the terms of the GNU Affero General Public
# License as published by the Free Software Foundation, either version
# 3 of the License, or (at your option) any later version.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>
#
# DISCLAIMER
#
# The software Hazard Modeller's Toolkit (openquake.hmtk) provided herein
# is released as a prototype implementation on behalf of
# scientists and engineers working within the GEM Foundation (Global
# Earthquake Model).
#
# It is distributed for the purpose of open collaboration and in the
# hope that it will be useful to the scientific, engineering, disaster
# risk and software design communities.
#
# The software is NOT distributed as part of GEM's OpenQuake suite
# (https://www.globalquakemodel.org/tools-products) and must be considered as a
# separate entity. The software provided herein is designed and implemented
# by scientific staff. It is not developed to the design standards, nor
# subject to same level of critical review by professional software
# developers, as GEM's OpenQuake software suite.
#
# Feedback and contribution to the software is welcome, and can be
# directed to the hazard scientific staff of the GEM Model Facility
# (hazard@globalquakemodel.org).
#
# The Hazard Modeller's Toolkit (openquake.hmtk) is therefore distributed WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# The GEM Foundation, and the authors of the software, assume no
# liability for use of the software.

'''
Module implements :class: openquake.hmtk.sources.source_model.mtkSourceModel, the
general class to describe a set of seismogenic sources
'''
from openquake.hazardlib.tom import PoissonTOM
from openquake.hazardlib.sourcewriter import write_source_model
from openquake.hmtk.sources.point_source import mtkPointSource
from openquake.hmtk.sources.area_source import mtkAreaSource
from openquake.hmtk.sources.simple_fault_source import mtkSimpleFaultSource
from openquake.hmtk.sources.complex_fault_source import mtkComplexFaultSource

class mtkSourceModel(object):
    '''
    Object to describe a seismogenic source model (composite of
    multiple sources with mixed typologies)

    :param str id:
        Identifier for the source model
    :param str name:
        Source model name
    :param list sources:
        List of seismogenic sources
    '''

    def __init__(self, identifier=None, name=None, sources=None):
        self.id = identifier
        self.name = name
        if isinstance(sources, list):
            self.sources = sources
        else:
            if sources:
                raise ValueError('Sources must be input as list!')
            self.sources = []

    def __iter__(self):
        for source in self.sources:
            yield source

    def __len__(self):
        return len(sources)

    def get_number_sources(self):
        '''
        Returns the number of sources in the model
        '''
        return len(self.sources)

    def serialise_to_nrml(self, filename, use_defaults=False):
        '''
        Writes the source model to a nrml source model file given by the
        filename

        :param str filename:
            Path to output file

        :param bool use_defaults:
            Boolean to indicate whether to use default values (True) or not.
            If set to False, ValueErrors will be raised when an essential
            attribute is missing.
        '''
        source_model = self.convert_to_oqhazardlib(
            PoissonTOM(1.0), 2.0, 2.0, 10.0, use_defaults=use_defaults)
        write_source_model(filename, source_model, name=self.name)

    def convert_to_oqhazardlib(
            self, tom, simple_mesh_spacing=1.0,
            complex_mesh_spacing=2.0, area_discretisation=10.0,
            use_defaults=False):
        """
        Converts the source model to an iterator of sources of :class:
        openquake.hazardlib.source.base.BaseSeismicSource
        """
        oq_source_model = []
        for source in self.sources:
            if isinstance(source, mtkAreaSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    simple_mesh_spacing,
                    area_discretisation,
                    use_defaults))
            elif isinstance(source, mtkPointSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    simple_mesh_spacing,
                    use_defaults))
            elif isinstance(source, mtkSimpleFaultSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    simple_mesh_spacing,
                    use_defaults))
            elif isinstance(source, mtkComplexFaultSource):
                oq_source_model.append(source.create_oqhazardlib_source(
                    tom,
                    complex_mesh_spacing,
                    use_defaults))
            else:
                raise ValueError('Source type not recognised!')
        return oq_source_model
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

"""
This module defines functions to compute loss ratio curves
using the classical psha based approach.
"""

from scipy import sqrt, stats, log, exp
from numpy import empty, linspace
from numpy import array, concatenate
from numpy import subtract, mean

from openquake import shapes
from openquake.risk.common import loop, collect
from openquake.utils.general import MemoizeMutable

STEPS_PER_INTERVAL = 5

def compute_loss_ratio_curve(vuln_function, hazard_curve):
    """Compute a loss ratio curve for a specific hazard curve (e.g., site),
    by applying a given vulnerability function.

    A loss ratio curve is a function that has loss ratios as X values
    and PoEs (Probabilities of Exceendance) as Y values.

    This is the main (and only) public function of this module.

    :param vuln_function: the vulnerability function used
        to compute the curve.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    :param hazard_curve: the hazard curve used to compute the curve.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    """

    lrem = _compute_lrem(vuln_function)
    lrem_po = _compute_lrem_po(vuln_function, lrem, hazard_curve)
    loss_ratios = _generate_loss_ratios(vuln_function)

    return shapes.Curve(zip(loss_ratios, lrem_po.sum(axis=1)))

def _compute_lrem_po(vuln_function, lrem, hazard_curve):
    """Compute the LREM * PoOs (Probability of Occurence) matrix.

    :param vuln_function: the vulnerability function used
        to compute the matrix.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    :param hazard_curve: the hazard curve used to compute the matrix.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    :param lrem: the LREM used to compute the matrix.
    :type lrem: 2-dimensional :py:class:`numpy.ndarray`
    """

    lrem = array(lrem)
    lrem_po = empty(lrem.shape)
    imls = _compute_imls(vuln_function)

    if hazard_curve:
        pos = _convert_pes_to_pos(hazard_curve, imls)
        for idx, po in enumerate(pos):
            lrem_po[:, idx] = lrem[:, idx] * po

    return lrem_po

def _generate_loss_ratios(vuln_function):
    """Generate the set of loss ratios used to compute the LREM
    (Loss Ratio Exceedance Matrix).

    :param vuln_function: the vulnerability function where the
        loss ratios are taken from.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    """

    # we manually add 0.0 as first loss ratio and 1.0 as last loss ratio
    loss_ratios = concatenate(
        (array([0.0]), vuln_function.means, array([1.0])))

    return _split_loss_ratios(loss_ratios)

@MemoizeMutable
def _compute_lrem(vuln_function, distribution=None):
    """Compute the LREM (Loss Ratio Exceedance Matrix).

    :param vuln_function: the vulnerability function used
        to compute the LREM.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    """

    if distribution is None:
        # this is so we can memoize the thing
        distribution = stats.lognorm

    loss_ratios = _generate_loss_ratios(vuln_function)
    # LREM has number of rows equal to the number of loss ratios
    # and number of columns equal to the number if imls
    lrem = empty((loss_ratios.size, vuln_function.imls.size), float)

    for idx, value in enumerate(vuln_function):
        mean_val, cov = value[1:]

        stddev = cov * mean_val
        variance = stddev ** 2.0
        mu = log(mean_val ** 2.0 / sqrt(variance + mean_val ** 2.0))
        sigma = sqrt(log((variance / mean_val ** 2.0) + 1.0))

        for row, value in enumerate(lrem):
            lrem[row][idx] = distribution.sf(loss_ratios[row],
                    sigma, scale=exp(mu))

    return lrem

def _split_loss_ratios(loss_ratios, steps=None):
    """Split the loss ratios, producing a new set of loss ratios.

    :param loss_ratios: the loss ratios to be splitted.
    :type loss_ratios: list
    :param steps: the number of steps we make to go from one loss
        ratio to the next. For example, if we have [1.0, 2.0]:

        steps = 1 produces [1.0, 2.0]
        steps = 2 produces [1.0, 1.5, 2.0]
        steps = 3 produces [1.0, 1.33, 1.66, 2.0]
    :type steps: integer
    """

    if steps is None:
        steps = STEPS_PER_INTERVAL

    splitted_ratios = set()

    for interval in loop(array(loss_ratios), linspace, steps + 1):
        splitted_ratios.update(interval)

    return array(sorted(splitted_ratios))

def _compute_imls(vuln_function):
    """Compute the mean IMLs (Intensity Measure Level)
    for the given vulnerability function.

    :param vuln_function: the vulnerability function where
        the IMLs (Intensity Measure Level) are taken from.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    """

    imls = vuln_function.imls

    # "special" cases for lowest part and highest part of the curve
    lowest_iml_value = imls[0] - ((imls[1] - imls[0]) / 2)

    # if the calculated lowest_curve_value goes < 0 we have to force the 0
    # IMLs have to be >= 0
    if lowest_iml_value < 0:
        lowest_iml_value = 0

    highest_iml_value = imls[-1] + ((imls[-1] - imls[-2]) / 2)
    between_iml_values = collect(loop(imls, lambda x, y: mean([x, y])))

    return [lowest_iml_value] + between_iml_values + [highest_iml_value]

def _compute_pes_from_imls(hazard_curve, imls):
    """Return the PoEs (Probability of Exceendance) defined in the
    given hazard curve for each IML (Intensity Measure Level) passed.

    :param hazard_curve: the hazard curve used to extract the PoEs.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    :param imls: the IMLs (Intensity Measure Level) of the
        vulnerability function used to interpolate the hazard curve.
    :type imls: list
    """

    return array([hazard_curve.ordinate_for(iml) for iml in imls])

def _convert_pes_to_pos(hazard_curve, imls):
    """For each IML (Intensity Measure Level) compute the
    PoOs (Probability of Occurence) from the PoEs
    (Probability of Exceendance) defined in the given hazard curve.

    :param hazard_curve: the hazard curve used to compute the PoOs.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    :param imls: the IMLs (Intensity Measure Level) of the
        vulnerability function used to interpolate the hazard curve.
    :type imls: list
    """

    return collect(loop(_compute_pes_from_imls(hazard_curve, imls),
            lambda x, y: subtract(array(x), array(y))))
<commit_msg>using ordinate_for on all imls in a single pass<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

"""
This module defines functions to compute loss ratio curves
using the classical psha based approach.
"""

from scipy import sqrt, stats, log, exp
from numpy import empty, linspace
from numpy import array, concatenate
from numpy import subtract, mean

from openquake import shapes
from openquake.risk.common import loop, collect
from openquake.utils.general import MemoizeMutable

STEPS_PER_INTERVAL = 5

def compute_loss_ratio_curve(vuln_function, hazard_curve):
    """Compute a loss ratio curve for a specific hazard curve (e.g., site),
    by applying a given vulnerability function.

    A loss ratio curve is a function that has loss ratios as X values
    and PoEs (Probabilities of Exceendance) as Y values.

    This is the main (and only) public function of this module.

    :param vuln_function: the vulnerability function used
        to compute the curve.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    :param hazard_curve: the hazard curve used to compute the curve.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    """

    lrem = _compute_lrem(vuln_function)
    lrem_po = _compute_lrem_po(vuln_function, lrem, hazard_curve)
    loss_ratios = _generate_loss_ratios(vuln_function)

    return shapes.Curve(zip(loss_ratios, lrem_po.sum(axis=1)))

def _compute_lrem_po(vuln_function, lrem, hazard_curve):
    """Compute the LREM * PoOs (Probability of Occurence) matrix.

    :param vuln_function: the vulnerability function used
        to compute the matrix.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    :param hazard_curve: the hazard curve used to compute the matrix.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    :param lrem: the LREM used to compute the matrix.
    :type lrem: 2-dimensional :py:class:`numpy.ndarray`
    """

    lrem = array(lrem)
    lrem_po = empty(lrem.shape)
    imls = _compute_imls(vuln_function)

    if hazard_curve:
        pos = _convert_pes_to_pos(hazard_curve, imls)
        for idx, po in enumerate(pos):
            lrem_po[:, idx] = lrem[:, idx] * po

    return lrem_po

def _generate_loss_ratios(vuln_function):
    """Generate the set of loss ratios used to compute the LREM
    (Loss Ratio Exceedance Matrix).

    :param vuln_function: the vulnerability function where the
        loss ratios are taken from.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    """

    # we manually add 0.0 as first loss ratio and 1.0 as last loss ratio
    loss_ratios = concatenate(
        (array([0.0]), vuln_function.means, array([1.0])))

    return _split_loss_ratios(loss_ratios)

@MemoizeMutable
def _compute_lrem(vuln_function, distribution=None):
    """Compute the LREM (Loss Ratio Exceedance Matrix).

    :param vuln_function: the vulnerability function used
        to compute the LREM.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    """

    if distribution is None:
        # this is so we can memoize the thing
        distribution = stats.lognorm

    loss_ratios = _generate_loss_ratios(vuln_function)
    # LREM has number of rows equal to the number of loss ratios
    # and number of columns equal to the number if imls
    lrem = empty((loss_ratios.size, vuln_function.imls.size), float)

    for idx, value in enumerate(vuln_function):
        mean_val, cov = value[1:]

        stddev = cov * mean_val
        variance = stddev ** 2.0
        mu = log(mean_val ** 2.0 / sqrt(variance + mean_val ** 2.0))
        sigma = sqrt(log((variance / mean_val ** 2.0) + 1.0))

        for row, value in enumerate(lrem):
            lrem[row][idx] = distribution.sf(loss_ratios[row],
                    sigma, scale=exp(mu))

    return lrem

def _split_loss_ratios(loss_ratios, steps=None):
    """Split the loss ratios, producing a new set of loss ratios.

    :param loss_ratios: the loss ratios to be splitted.
    :type loss_ratios: list
    :param steps: the number of steps we make to go from one loss
        ratio to the next. For example, if we have [1.0, 2.0]:

        steps = 1 produces [1.0, 2.0]
        steps = 2 produces [1.0, 1.5, 2.0]
        steps = 3 produces [1.0, 1.33, 1.66, 2.0]
    :type steps: integer
    """

    if steps is None:
        steps = STEPS_PER_INTERVAL

    splitted_ratios = set()

    for interval in loop(array(loss_ratios), linspace, steps + 1):
        splitted_ratios.update(interval)

    return array(sorted(splitted_ratios))

def _compute_imls(vuln_function):
    """Compute the mean IMLs (Intensity Measure Level)
    for the given vulnerability function.

    :param vuln_function: the vulnerability function where
        the IMLs (Intensity Measure Level) are taken from.
    :type vuln_function: :py:class:`openquake.shapes.VulnerabilityFunction`
    """

    imls = vuln_function.imls

    # "special" cases for lowest part and highest part of the curve
    lowest_iml_value = imls[0] - ((imls[1] - imls[0]) / 2)

    # if the calculated lowest_curve_value goes < 0 we have to force the 0
    # IMLs have to be >= 0
    if lowest_iml_value < 0:
        lowest_iml_value = 0

    highest_iml_value = imls[-1] + ((imls[-1] - imls[-2]) / 2)
    between_iml_values = collect(loop(imls, lambda x, y: mean([x, y])))

    return [lowest_iml_value] + between_iml_values + [highest_iml_value]

def _compute_pes_from_imls(hazard_curve, imls):
    """Return the PoEs (Probability of Exceendance) defined in the
    given hazard curve for each IML (Intensity Measure Level) passed.

    :param hazard_curve: the hazard curve used to extract the PoEs.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    :param imls: the IMLs (Intensity Measure Level) of the
        vulnerability function used to interpolate the hazard curve.
    :type imls: list
    """

    return hazard_curve.ordinate_for(imls)

def _convert_pes_to_pos(hazard_curve, imls):
    """For each IML (Intensity Measure Level) compute the
    PoOs (Probability of Occurence) from the PoEs
    (Probability of Exceendance) defined in the given hazard curve.

    :param hazard_curve: the hazard curve used to compute the PoOs.
    :type hazard_curve: :py:class:`openquake.shapes.Curve`
    :param imls: the IMLs (Intensity Measure Level) of the
        vulnerability function used to interpolate the hazard curve.
    :type imls: list
    """

    return collect(loop(_compute_pes_from_imls(hazard_curve, imls),
            lambda x, y: subtract(array(x), array(y))))
<|endoftext|>


<commit_before>'''
As of zookeeper 3.4.0, the `mntr` admin command is provided for easy parsing of zookeeper stats.
This check first parses the `stat` admin command for a version number.
If the zookeeper version supports `mntr`, it is also parsed.

Duplicate information is being reported by both `mntr` and `stat` to keep backwards compatability.
Example:
    `stat` reports: zookeeper.latency.avg
    `mntr` reports: zookeeper.avg.latency
If available, make use of the data reported by `mntr` not `stat`.
The duplicate `stat` reports are only kept for backward compatability.

Besides the usual zookeeper state of `leader`, `follower`, `observer` and `standalone`,
this check will report three other states:

    `down`: the check cannot connect to zookeeper
    `inactive`: the zookeeper instance has lost connection to the cluster
    `unknown`: an unexpected error has occured in this check

States can be accessed through the gauge `zookeeper.instances.<state>,
through the set `zookeeper.instances`, or through the `mode:<state>` tag.

Parses the response from zookeeper's `stat` admin command, which looks like:

```
Zookeeper version: 3.2.2--1, built on 03/16/2010 07:31 GMT
Clients:
 /10.42.114.160:32634[1](queued=0,recved=12,sent=0)
 /10.37.137.74:21873[1](queued=0,recved=53613,sent=0)
 /10.37.137.74:21876[1](queued=0,recved=57436,sent=0)
 /10.115.77.32:32990[1](queued=0,recved=16,sent=0)
 /10.37.137.74:21891[1](queued=0,recved=55011,sent=0)
 /10.37.137.74:21797[1](queued=0,recved=19431,sent=0)

Latency min/avg/max: -10/0/20007
Received: 101032173
Sent: 0
Outstanding: 0
Zxid: 0x1034799c7
Mode: leader
Node count: 487
```

`stat` tested with Zookeeper versions 3.0.0 to 3.4.5

The following is an example of the `mntr` commands output:

```
zk_version  3.4.5-cdh4.4.0--1, built on 09/04/2013 01:46 GMT
zk_avg_latency  0
zk_max_latency  0
zk_min_latency  0
zk_packets_received 4
zk_packets_sent 3
zk_num_alive_connections    1
zk_outstanding_requests 0
zk_server_state standalone
zk_znode_count  4
zk_watch_count  0
zk_ephemerals_count 0
zk_approximate_data_size    27
zk_open_file_descriptor_count   29
zk_max_file_descriptor_count    4096
```

`mntr` tested with ZooKeeper 3.4.5
'''
# stdlib
from distutils.version import LooseVersion
import re
import socket
from StringIO import StringIO
import struct
import sys

# project
from checks import AgentCheck

class ZKConnectionFailure(Exception):
    """ Raised when we are unable to connect or get the output of a command. """
    pass

class ZookeeperCheck(AgentCheck):
    version_pattern = re.compile(r'Zookeeper version: ([^.]+)\.([^.]+)\.([^-]+)', flags=re.I)

    SOURCE_TYPE_NAME = 'zookeeper'

    STATUS_TYPES = [
        'leader',
        'follower',
        'observer',
        'standalone',
        'down',
        'inactive',
        'unknown',
    ]

    def check(self, instance):
        host = instance.get('host', 'localhost')
        port = int(instance.get('port', 2181))
        timeout = float(instance.get('timeout', 3.0))
        expected_mode = (instance.get('expected_mode') or '').strip()
        tags = instance.get('tags', [])
        cx_args = (host, port, timeout)
        sc_tags = ["host:{0}".format(host), "port:{0}".format(port)]
        hostname = socket.gethostname()

        zk_version = None # parse_stat will parse and set version string

        # Send a service check based on the `ruok` response.
        # Set instance status to down if not ok.
        try:
            ruok_out = self._send_command('ruok', *cx_args)
        except ZKConnectionFailure:
            # The server should not respond at all if it's not OK.
            status = AgentCheck.CRITICAL
            message = 'No response from `ruok` command'
            self.increment('zookeeper.timeouts')
            self.report_instance_mode(hostname, 'down', tags)
            raise
        else:
            ruok_out.seek(0)
            ruok = ruok_out.readline()
            if ruok == 'imok':
                status = AgentCheck.OK
            else:
                status = AgentCheck.WARNING
            message = u'Response from the server: %s' % ruok
        finally:
            self.service_check('zookeeper.ruok', status, message=message,
                    tags=sc_tags)

        # Read metrics from the `stat` output.
        try:
            stat_out = self._send_command('stat', *cx_args)
        except ZKConnectionFailure:
            self.increment('zookeeper.timeouts')
            self.report_instance_mode(hostname, 'down', tags)
            raise
        except:
            e = sys.exc_info()[1]
            print >> sys.stderr, "Error: %s" % e
            self.increment('zookeeper.datadog_client_exception')
            self.report_instance_mode(hostname, 'unknown', tags)
            raise
        else:
            # Parse the response
            metrics, new_tags, mode, zk_version = self.parse_stat(stat_out)

            # Write the data
            if mode != 'inactive':
                for metric, value in metrics:
                    self.gauge(metric, value, tags=tags + new_tags)
            self.report_instance_mode(hostname, mode, tags)

            if expected_mode:
                if mode == expected_mode:
                    status = AgentCheck.OK
                    message = u"Server is in %s mode" % mode
                else:
                    status = AgentCheck.CRITICAL
                    message = u"Server is in %s mode but check expects %s mode"\
                              % (mode, expected_mode)
                self.service_check('zookeeper.mode', status, message=message,
                                   tags=sc_tags)

        if zk_version and LooseVersion(zk_version) > LooseVersion("3.4.0"):
            try:
                mntr_out = self._send_command('mntr', *cx_args)
            except ZKConnectionFailure:
                self.increment('zookeeper.timeouts')
                self.report_instance_mode(hostname, 'down', tags)
                raise
            except:
                e = sys.exc_info()[1]
                print >> sys.stderr, "Error: %s" % e
                self.increment('zookeeper.datadog_client_exception')
                self.report_instance_mode(hostname, 'unknown', tags)
                raise
            else:
                metrics, mode = self.parse_mntr(mntr_out)
                mode_tag = "mode:%s" % mode
                if mode != 'inactive':
                    for name in metrics:
                        self.gauge(name, metrics[name], tags=tags + [mode_tag])

    def report_instance_mode(self, hostname, mode, tags):
        tags = tags + ['mode:%s' % mode]
        self.set('zookeeper.instances', hostname, tags=tags)

        gauges = {}
        for key in self.STATUS_TYPES:
            gauges[key] = 0

        if mode in gauges:
            gauges[mode] = 1
        else:
            gauges['unknown'] = 1

        for k in gauges:
            gauge_name = 'zookeeper.instances.%s' % k
            self.gauge(gauge_name, gauges[k])

    def _send_command(self, command, host, port, timeout):
        sock = socket.socket()
        sock.settimeout(timeout)
        buf = StringIO()
        chunk_size = 1024
        # try-finally and try-except to stay compatible with python 2.4
        try:
            try:
                # Connect to the zk client port and send the stat command
                sock.connect((host, port))
                sock.sendall(command)

                # Read the response into a StringIO buffer
                chunk = sock.recv(chunk_size)
                buf.write(chunk)
                num_reads = 1
                max_reads = 10000
                while chunk:
                    if num_reads > max_reads:
                        # Safeguard against an infinite loop
                        raise Exception("Read %s bytes before exceeding max reads of %s. "
                                        % (buf.tell(), max_reads))
                    chunk = sock.recv(chunk_size)
                    buf.write(chunk)
                    num_reads += 1
            except (socket.timeout, socket.error):
                raise ZKConnectionFailure()
        finally:
            sock.close()
        return buf

    def parse_stat(self, buf):
        ''' `buf` is a readable file-like object
            returns a tuple: (metrics, tags, mode, version)
        '''
        metrics = []
        buf.seek(0)

        # Check the version line to make sure we parse the rest of the
        # body correctly. Particularly, the Connections val was added in
        # >= 3.4.4.
        start_line = buf.readline()
        match = self.version_pattern.match(start_line)
        if match is None:
            return (None, None, "inactive", None)
            raise Exception("Could not parse version from stat command output: %s" % start_line)
        else:
            version_tuple = match.groups()
        has_connections_val = version_tuple >= ('3', '4', '4')
        version = "%s.%s.%s" % version_tuple

        # Clients:
        buf.readline() # skip the Clients: header
        connections = 0
        client_line = buf.readline().strip()
        if client_line:
            connections += 1
        while client_line:
            client_line = buf.readline().strip()
            if client_line:
                connections += 1

        # Latency min/avg/max: -10/0/20007
        _, value = buf.readline().split(':')
        l_min, l_avg, l_max = [int(v) for v in value.strip().split('/')]
        metrics.append(('zookeeper.latency.min', l_min))
        metrics.append(('zookeeper.latency.avg', l_avg))
        metrics.append(('zookeeper.latency.max', l_max))

        # Received: 101032173
        _, value = buf.readline().split(':')
        metrics.append(('zookeeper.bytes_received', long(value.strip())))

        # Sent: 1324
        _, value = buf.readline().split(':')
        metrics.append(('zookeeper.bytes_sent', long(value.strip())))

        if has_connections_val:
            # Connections: 1
            _, value = buf.readline().split(':')
            metrics.append(('zookeeper.connections', int(value.strip())))
        else:
            # If the zk version doesnt explicitly give the Connections val,
            # use the value we computed from the client list.
            metrics.append(('zookeeper.connections', connections))

        # Outstanding: 0
        _, value = buf.readline().split(':')
        # Fixme: This metric name is wrong. It should be removed in a major version of the agent
        # See https://github.com/DataDog/dd-agent/issues/1383
        metrics.append(('zookeeper.bytes_outstanding', long(value.strip())))
        metrics.append(('zookeeper.outstanding_requests', long(value.strip())))

        # Zxid: 0x1034799c7
        _, value = buf.readline().split(':')
        # Parse as a 64 bit hex int
        zxid = long(value.strip(), 16)
        # convert to bytes
        zxid_bytes = struct.pack('>q', zxid)
        # the higher order 4 bytes is the epoch
        (zxid_epoch,) = struct.unpack('>i', zxid_bytes[0:4])
        # the lower order 4 bytes is the count
        (zxid_count,) = struct.unpack('>i', zxid_bytes[4:8])

        metrics.append(('zookeeper.zxid.epoch', zxid_epoch))
        metrics.append(('zookeeper.zxid.count', zxid_count))

        # Mode: leader
        _, value = buf.readline().split(':')
        mode = value.strip().lower()
        tags = [u'mode:' + mode]

        # Node count: 487
        _, value = buf.readline().split(':')
        metrics.append(('zookeeper.nodes', long(value.strip())))

        return metrics, tags, mode, version

    def parse_mntr(self, buf):
        ''' `buf` is a readable file-like object
            Returns: a tuple (metrics, mode)
            if mode == 'inactive', metrics will be None
        '''

        buf.seek(0)
        first = buf.readline() # first is version string or error
        if first == 'This ZooKeeper instance is not currently serving requests':
            return (None, 'inactive')

        metrics = {}

        for line in buf:
            try:
                key, value = line.split()
                name = self._normalize_metric_label(key)
                metrics[name] = value
            except:
                e = sys.exc_info()[1]
                raise Exception("Data not in 'key value' format, error: %s" % e)

        # mode is a string {'standalone', 'leader', 'follower', 'observer'}
        mode = metrics.pop('zookeeper.server.state').lower()

        for key in metrics: # everything else is an int
            metrics[key] = int(metrics[key])

        return (metrics, mode)

    def _normalize_metric_label(self, key):
        if re.match('zk', key):
            key = key.replace('zk', 'zookeeper', 1)
        return key.replace('_', '.')
<commit_msg>pythonic exceptions and proper hostname usage<commit_after>'''
As of zookeeper 3.4.0, the `mntr` admin command is provided for easy parsing of zookeeper stats.
This check first parses the `stat` admin command for a version number.
If the zookeeper version supports `mntr`, it is also parsed.

Duplicate information is being reported by both `mntr` and `stat` to keep backwards compatability.
Example:
    `stat` reports: zookeeper.latency.avg
    `mntr` reports: zookeeper.avg.latency
If available, make use of the data reported by `mntr` not `stat`.
The duplicate `stat` reports are only kept for backward compatability.

Besides the usual zookeeper state of `leader`, `follower`, `observer` and `standalone`,
this check will report three other states:

    `down`: the check cannot connect to zookeeper
    `inactive`: the zookeeper instance has lost connection to the cluster
    `unknown`: an unexpected error has occured in this check

States can be accessed through the gauge `zookeeper.instances.<state>,
through the set `zookeeper.instances`, or through the `mode:<state>` tag.

Parses the response from zookeeper's `stat` admin command, which looks like:

```
Zookeeper version: 3.2.2--1, built on 03/16/2010 07:31 GMT
Clients:
 /10.42.114.160:32634[1](queued=0,recved=12,sent=0)
 /10.37.137.74:21873[1](queued=0,recved=53613,sent=0)
 /10.37.137.74:21876[1](queued=0,recved=57436,sent=0)
 /10.115.77.32:32990[1](queued=0,recved=16,sent=0)
 /10.37.137.74:21891[1](queued=0,recved=55011,sent=0)
 /10.37.137.74:21797[1](queued=0,recved=19431,sent=0)

Latency min/avg/max: -10/0/20007
Received: 101032173
Sent: 0
Outstanding: 0
Zxid: 0x1034799c7
Mode: leader
Node count: 487
```

`stat` tested with Zookeeper versions 3.0.0 to 3.4.5

The following is an example of the `mntr` commands output:

```
zk_version  3.4.5-cdh4.4.0--1, built on 09/04/2013 01:46 GMT
zk_avg_latency  0
zk_max_latency  0
zk_min_latency  0
zk_packets_received 4
zk_packets_sent 3
zk_num_alive_connections    1
zk_outstanding_requests 0
zk_server_state standalone
zk_znode_count  4
zk_watch_count  0
zk_ephemerals_count 0
zk_approximate_data_size    27
zk_open_file_descriptor_count   29
zk_max_file_descriptor_count    4096
```

`mntr` tested with ZooKeeper 3.4.5
'''
# stdlib
from distutils.version import LooseVersion
import re
import socket
from StringIO import StringIO
import struct

# project
from checks import AgentCheck
from util import get_hostname

class ZKConnectionFailure(Exception):
    """ Raised when we are unable to connect or get the output of a command. """
    pass

class ZookeeperCheck(AgentCheck):
    version_pattern = re.compile(r'Zookeeper version: ([^.]+)\.([^.]+)\.([^-]+)', flags=re.I)

    SOURCE_TYPE_NAME = 'zookeeper'

    STATUS_TYPES = [
        'leader',
        'follower',
        'observer',
        'standalone',
        'down',
        'inactive',
        'unknown',
    ]

    def check(self, instance):
        host = instance.get('host', 'localhost')
        port = int(instance.get('port', 2181))
        timeout = float(instance.get('timeout', 3.0))
        expected_mode = (instance.get('expected_mode') or '').strip()
        tags = instance.get('tags', [])
        cx_args = (host, port, timeout)
        sc_tags = ["host:{0}".format(host), "port:{0}".format(port)]
        hostname = get_hostname(self.agentConfig)

        zk_version = None # parse_stat will parse and set version string

        # Send a service check based on the `ruok` response.
        # Set instance status to down if not ok.
        try:
            ruok_out = self._send_command('ruok', *cx_args)
        except ZKConnectionFailure:
            # The server should not respond at all if it's not OK.
            status = AgentCheck.CRITICAL
            message = 'No response from `ruok` command'
            self.increment('zookeeper.timeouts')
            self.report_instance_mode(hostname, 'down', tags)
            raise
        else:
            ruok_out.seek(0)
            ruok = ruok_out.readline()
            if ruok == 'imok':
                status = AgentCheck.OK
            else:
                status = AgentCheck.WARNING
            message = u'Response from the server: %s' % ruok
        finally:
            self.service_check('zookeeper.ruok', status, message=message,
                    tags=sc_tags)

        # Read metrics from the `stat` output.
        try:
            stat_out = self._send_command('stat', *cx_args)
        except ZKConnectionFailure:
            self.increment('zookeeper.timeouts')
            self.report_instance_mode(hostname, 'down', tags)
            raise
        except Exception as e:
            self.warning(e)
            self.increment('zookeeper.datadog_client_exception')
            self.report_instance_mode(hostname, 'unknown', tags)
            raise
        else:
            # Parse the response
            metrics, new_tags, mode, zk_version = self.parse_stat(stat_out)

            # Write the data
            if mode != 'inactive':
                for metric, value in metrics:
                    self.gauge(metric, value, tags=tags + new_tags)
            self.report_instance_mode(hostname, mode, tags)

            if expected_mode:
                if mode == expected_mode:
                    status = AgentCheck.OK
                    message = u"Server is in %s mode" % mode
                else:
                    status = AgentCheck.CRITICAL
                    message = u"Server is in %s mode but check expects %s mode"\
                              % (mode, expected_mode)
                self.service_check('zookeeper.mode', status, message=message,
                                   tags=sc_tags)

        if zk_version and LooseVersion(zk_version) > LooseVersion("3.4.0"):
            try:
                mntr_out = self._send_command('mntr', *cx_args)
            except ZKConnectionFailure:
                self.increment('zookeeper.timeouts')
                self.report_instance_mode(hostname, 'down', tags)
                raise
            except Exception as e:
                self.warning(e)
                self.increment('zookeeper.datadog_client_exception')
                self.report_instance_mode(hostname, 'unknown', tags)
                raise
            else:
                metrics, mode = self.parse_mntr(mntr_out)
                mode_tag = "mode:%s" % mode
                if mode != 'inactive':
                    for name in metrics:
                        self.gauge(name, metrics[name], tags=tags + [mode_tag])

    def report_instance_mode(self, hostname, mode, tags):
        tags = tags + ['mode:%s' % mode]
        self.set('zookeeper.instances', hostname, tags=tags)

        gauges = {}
        for key in self.STATUS_TYPES:
            gauges[key] = 0

        if mode in gauges:
            gauges[mode] = 1
        else:
            gauges['unknown'] = 1

        for k in gauges:
            gauge_name = 'zookeeper.instances.%s' % k
            self.gauge(gauge_name, gauges[k])

    def _send_command(self, command, host, port, timeout):
        sock = socket.socket()
        sock.settimeout(timeout)
        buf = StringIO()
        chunk_size = 1024
        # try-finally and try-except to stay compatible with python 2.4
        try:
            try:
                # Connect to the zk client port and send the stat command
                sock.connect((host, port))
                sock.sendall(command)

                # Read the response into a StringIO buffer
                chunk = sock.recv(chunk_size)
                buf.write(chunk)
                num_reads = 1
                max_reads = 10000
                while chunk:
                    if num_reads > max_reads:
                        # Safeguard against an infinite loop
                        raise Exception("Read %s bytes before exceeding max reads of %s. "
                                        % (buf.tell(), max_reads))
                    chunk = sock.recv(chunk_size)
                    buf.write(chunk)
                    num_reads += 1
            except (socket.timeout, socket.error):
                raise ZKConnectionFailure()
        finally:
            sock.close()
        return buf

    def parse_stat(self, buf):
        ''' `buf` is a readable file-like object
            returns a tuple: (metrics, tags, mode, version)
        '''
        metrics = []
        buf.seek(0)

        # Check the version line to make sure we parse the rest of the
        # body correctly. Particularly, the Connections val was added in
        # >= 3.4.4.
        start_line = buf.readline()
        match = self.version_pattern.match(start_line)
        if match is None:
            return (None, None, "inactive", None)
            raise Exception("Could not parse version from stat command output: %s" % start_line)
        else:
            version_tuple = match.groups()
        has_connections_val = version_tuple >= ('3', '4', '4')
        version = "%s.%s.%s" % version_tuple

        # Clients:
        buf.readline() # skip the Clients: header
        connections = 0
        client_line = buf.readline().strip()
        if client_line:
            connections += 1
        while client_line:
            client_line = buf.readline().strip()
            if client_line:
                connections += 1

        # Latency min/avg/max: -10/0/20007
        _, value = buf.readline().split(':')
        l_min, l_avg, l_max = [int(v) for v in value.strip().split('/')]
        metrics.append(('zookeeper.latency.min', l_min))
        metrics.append(('zookeeper.latency.avg', l_avg))
        metrics.append(('zookeeper.latency.max', l_max))

        # Received: 101032173
        _, value = buf.readline().split(':')
        metrics.append(('zookeeper.bytes_received', long(value.strip())))

        # Sent: 1324
        _, value = buf.readline().split(':')
        metrics.append(('zookeeper.bytes_sent', long(value.strip())))

        if has_connections_val:
            # Connections: 1
            _, value = buf.readline().split(':')
            metrics.append(('zookeeper.connections', int(value.strip())))
        else:
            # If the zk version doesnt explicitly give the Connections val,
            # use the value we computed from the client list.
            metrics.append(('zookeeper.connections', connections))

        # Outstanding: 0
        _, value = buf.readline().split(':')
        # Fixme: This metric name is wrong. It should be removed in a major version of the agent
        # See https://github.com/DataDog/dd-agent/issues/1383
        metrics.append(('zookeeper.bytes_outstanding', long(value.strip())))
        metrics.append(('zookeeper.outstanding_requests', long(value.strip())))

        # Zxid: 0x1034799c7
        _, value = buf.readline().split(':')
        # Parse as a 64 bit hex int
        zxid = long(value.strip(), 16)
        # convert to bytes
        zxid_bytes = struct.pack('>q', zxid)
        # the higher order 4 bytes is the epoch
        (zxid_epoch,) = struct.unpack('>i', zxid_bytes[0:4])
        # the lower order 4 bytes is the count
        (zxid_count,) = struct.unpack('>i', zxid_bytes[4:8])

        metrics.append(('zookeeper.zxid.epoch', zxid_epoch))
        metrics.append(('zookeeper.zxid.count', zxid_count))

        # Mode: leader
        _, value = buf.readline().split(':')
        mode = value.strip().lower()
        tags = [u'mode:' + mode]

        # Node count: 487
        _, value = buf.readline().split(':')
        metrics.append(('zookeeper.nodes', long(value.strip())))

        return metrics, tags, mode, version

    def parse_mntr(self, buf):
        ''' `buf` is a readable file-like object
            Returns: a tuple (metrics, mode)
            if mode == 'inactive', metrics will be None
        '''

        buf.seek(0)
        first = buf.readline() # first is version string or error
        if first == 'This ZooKeeper instance is not currently serving requests':
            return (None, 'inactive')

        metrics = {}

        for line in buf:
            try:
                key, value = line.split()
                name = self._normalize_metric_label(key)
                metrics[name] = value
            except Exception as e:
                raise Exception("Data not in 'key value' format, error: %s" % e)

        # mode is a string {'standalone', 'leader', 'follower', 'observer'}
        mode = metrics.pop('zookeeper.server.state').lower()

        for key in metrics: # everything else is an int
            metrics[key] = int(metrics[key])

        return (metrics, mode)

    def _normalize_metric_label(self, key):
        if re.match('zk', key):
            key = key.replace('zk', 'zookeeper', 1)
        return key.replace('_', '.')
<|endoftext|>


<commit_before>"""
OPAL Models!
"""
import collections
import json

from django.conf import settings
from django.db import models
from django.contrib import auth
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic
from django.dispatch import receiver
import reversion

from opal.utils import stringport, camelcase_to_underscore
from opal.utils.fields import ForeignKeyOrFreeText
from opal import exceptions, managers

# Imported models from module.

from opal.models.mixins import UpdatesFromDictMixin

options = stringport(settings.OPAL_OPTIONS_MODULE)

from django.contrib.auth.models import User

class UserProfile(models.Model):
    """
    Profile for our user
    """
    user                  = models.ForeignKey(User, unique=True)
    force_password_change = models.BooleanField(default=True)
    can_extract           = models.BooleanField(default=False)
    readonly              = models.BooleanField(default=False)

class Filter(models.Model):
    """
    Saved filters for users extracting data.
    """
    user     = models.ForeignKey(User)
    name     = models.CharField(max_length=200)
    criteria = models.TextField()

    def to_dict(self):
        return dict(id=self.pk, name=self.name, criteria=json.loads(self.criteria))

    def update_from_dict(self, data):
        self.criteria = json.dumps(data['criteria'])
        self.name = data['name']
        self.save()

class Patient(models.Model):
    def __unicode__(self):
        demographics = self.demographics_set.get()
        return '%s | %s' % (demographics.hospital_number, demographics.name)

    def create_episode(self):
        if self.get_active_episode() is None:
            return self.episode_set.create()
        else:
            raise exceptions.APIError('Patient %s already has active episode' % self)

    def get_active_episode(self):
        for episode in self.episode_set.order_by('id').reverse():
            if episode.is_active():
                return episode
        return None

    def to_dict(self, user):
        active_episode = self.get_active_episode()
        d = {
            'id': self.id,
            'episodes': {episode.id: episode.to_dict(user) for episode in self.episode_set.all()},
            'active_episode_id': active_episode.id if active_episode else None,
            }

        for model in PatientSubrecord.__subclasses__():
            subrecords = model.objects.filter(patient_id=self.id)
            d[model.get_api_name()] = [subrecord.to_dict(user) for subrecord in subrecords]
        return d

    def update_from_demographics_dict(self, demographics_data, user):
        demographics = self.demographics_set.get()
        demographics.update_from_dict(demographics_data, user)

class Episode(UpdatesFromDictMixin, models.Model):
    """
    An individual episode of care.

    A patient may have many episodes of care, but this maps to one occasion
    on which they found themselves on "The List".
    """
    patient = models.ForeignKey(Patient)
    active  = models.BooleanField(default=False)
    date_of_admission = models.DateField(null=True, blank=True)
    # TODO rename to date_of_discharge?
    discharge_date = models.DateField(null=True, blank=True)
    consistency_token = models.CharField(max_length=8)

    objects = managers.EpisodeManager()

    def __unicode__(self):
        demographics = self.patient.demographics_set.get()

        return '%s | %s | %s' % (demographics.hospital_number,
                                 demographics.name,
                                 self.date_of_admission)

    def is_active(self):
        # TODO Depreciate this.
        #
        # This is only here for API compatability.
        # Don't use me!
        return self.active

    def set_tag_names(self, tag_names, user):
        """
        1. Blitz dangling tags not in our current dict.
        2. Add new tags.
        3. Make sure that we set the Active boolean appropriately
        4. There is no step 4.
        """
        original_tag_names = self.get_tag_names(user)

        for tag_name in original_tag_names:
            if tag_name not in tag_names:
                params = {'team__name': tag_name}
                if tag_name == 'mine':
                    params['user'] = user
                self.tagging_set.get(**params).delete()

        for tag_name in tag_names:
            if tag_name not in original_tag_names:
                params = {'team': Team.objects.get(name=tag_name)}
                if tag_name == 'mine':
                    params['user'] = user
                self.tagging_set.create(**params)

        if len(tag_names) < 1:
            self.active = False
        elif tag_names == ['mine']:
            self.active = False
        elif not self.active:
            self.active = True
        self.save()

    def tagging_dict(self, user):
        if self.tagging_set.count() == 0:
            return [{}]
        td = [{t.team.name: True for t in self.tagging_set.exclude(team__name='mine')}]
        if self.tagging_set.filter(team__name='mine', user=user).count() > 0:
            td[0]['mine'] = True
        td[0]['id'] = self.id
        return td

    def get_tag_names(self, user):
        return [t.team.name for t in self.tagging_set.all() if t.user in (None, user)]

    def to_dict(self, user, shallow=False, with_context=False):
        """
        Serialisation to JSON for Episodes
        """
        d = {
            'id'               : self.id,
            'active'           : self.active,
            'date_of_admission': self.date_of_admission,
            'discharge_date'   : self.discharge_date,
            'consistency_token': self.consistency_token
            }
        if shallow:
            return d

        for model in PatientSubrecord.__subclasses__():
            subrecords = model.objects.filter(patient_id=self.patient.id)
            d[model.get_api_name()] = [subrecord.to_dict(user)
                                       for subrecord in subrecords]
        for model in EpisodeSubrecord.__subclasses__():
            subrecords = model.objects.filter(episode_id=self.id)
            d[model.get_api_name()] = [subrecord.to_dict(user)
                                       for subrecord in subrecords]

        d['tagging'] = self.tagging_dict(user)
        d['prev_episodes'] = []
        d['next_episodes'] = []

        if self.date_of_admission:
            eset = self.patient.episode_set

            d['prev_episodes'] = [
                e.to_dict(user, shallow=True)
                for e in
                eset.filter(date_of_admission__lt=self.date_of_admission).order_by('date_of_admission')]
            d['next_episodes'] = [
                e.to_dict(user, shallow=True)
                for e in
                eset.filter(date_of_admission__gt=self.date_of_admission).order_by('date_of_admission')]

        return d

    def update_from_location_dict(self, location_data, user):
        # TODO Completely depreciate this.
        location = self.location_set.get()
        location.update_from_dict(location_data, user)

class ContactNumber(models.Model):
    name = models.CharField(max_length=255)
    number = models.CharField(max_length=255)

    def __unicode__(self):
        return u'{0}: {1}'.format(self.name, self.number)

class Team(models.Model):
    name = models.CharField(max_length=250, help_text="This should only have letters and underscores")
    title = models.CharField(max_length=250)
    parent = models.ForeignKey('self', blank=True, null=True)
    active = models.BooleanField(default=True)
    order = models.IntegerField(blank=True, null=True)
    useful_numbers = models.ManyToManyField(ContactNumber, blank=True)

    def __unicode__(self):
        return self.title

    # TODO depreciate this and refactor accordingly
    @classmethod
    def to_TAGS(klass):
        from opal.utils import Tag
        teams = klass.objects.filter(active=True).order_by('order')
        tags = collections.OrderedDict()
        for team in teams:
            if not team.parent:
                tags[team.name] = Tag(team.name, team.title, [])
        for team in teams:
            if team.parent:
                tags[team.parent.name].subtags.append(Tag(team.name, team.title, None))
        return tags.values()

class Subrecord(UpdatesFromDictMixin, models.Model):
    consistency_token = models.CharField(max_length=8)

    _is_singleton = False

    class Meta:
        abstract = True

    def __unicode__(self):
        return u'{0}: {1}'.format(self.get_api_name(), self.id)

    @classmethod
    def get_api_name(cls):
        return camelcase_to_underscore(cls._meta.object_name)

    @classmethod
    def get_display_name(cls):
        if hasattr(cls, '_title'):
            return cls._title
        else:
            return cls._meta.object_name

    @classmethod
    def build_field_schema(cls):
        field_schema = []
        for fieldname in cls._get_fieldnames_to_serialize():
            if fieldname in ['id', 'patient_id', 'episode_id']:
                continue
            elif fieldname.endswith('_fk_id'):
                continue
            elif fieldname.endswith('_ft'):
                continue

            getter = getattr(cls, 'get_field_type_for_' + fieldname, None)
            if getter is None:
                field = cls._get_field_type(fieldname)
                if field in [models.CharField, ForeignKeyOrFreeText]:
                    field_type = 'string'
                else:
                    field_type = camelcase_to_underscore(field.__name__[:-5])
            else:
                field_type = getter()
            lookup_list = None
            if cls._get_field_type(fieldname) == ForeignKeyOrFreeText:
                fld = getattr(cls, fieldname)
                lookup_list = camelcase_to_underscore(fld.foreign_model.__name__)
            field_schema.append({'name': fieldname,
                                 'type': field_type,
                                 'lookup_list': lookup_list})
        return field_schema

    def _to_dict(self, user, fieldnames):
        """
        Allow a subset of FIELDNAMES
        """
        d = {}
        for name in fieldnames:
            getter = getattr(self, 'get_' + name, None)
            if getter is not None:
                value = getter(user)
            else:
                value = getattr(self, name)
            d[name] = value

        return d

    def to_dict(self, user):
        return self._to_dict(user, self._get_fieldnames_to_serialize())

class PatientSubrecord(Subrecord):
    patient = models.ForeignKey(Patient)

    class Meta:
        abstract = True

class EpisodeSubrecord(Subrecord):
    episode = models.ForeignKey(Episode, null=True)  # TODO make null=False

    class Meta:
        abstract = True

class Tagging(models.Model):
    _is_singleton = True
    _title = 'Teams'

    team = models.ForeignKey(Team, blank=True, null=True)
    user = models.ForeignKey(auth.models.User, null=True)
    episode = models.ForeignKey(Episode, null=True) # TODO make null=False

    def __unicode__(self):
        if self.user is not None:
            return 'User: %s - %s' % (self.user.username, self.team.name)
        else:
            return self.team.name

    @staticmethod
    def get_api_name():
        return 'tagging'

    @staticmethod
    def get_display_name():
        return 'Teams'

    @staticmethod
    def build_field_schema():
        teams = [{'name': t.name, 'type':'boolean'} for t in Team.objects.filter(active=True)]
        return teams
#        return [dict(name='team__name', type='string')]

    @classmethod
    def historic_tags_for_episodes(cls, episodes):
        """
        Given a list of episodes, return a dict indexed by episode id
        that contains historic tags for those episodes.
        """
        teams = {t.id: t.name for t in Team.objects.all()}
        deleted = reversion.get_deleted(cls)
        historic = collections.defaultdict(dict)
        for d in deleted:
            data = json.loads(d.serialized_data)[0]['fields']
            if data['episode'] in episodes:
                try:
                    tag_name = teams[data['team']]
                except KeyError:
                    tag_name = data['tag_name']
                historic[data['episode']][tag_name] = True
        return historic

    @classmethod
    def historic_episodes_for_tag(cls, tag):
        """
        Given a TAG return a list of episodes that have historically been
        tagged with it.
        """
        teams = {t.id: t.name for t in Team.objects.all()}
        deleted = reversion.get_deleted(cls)
        eids = set()
        for d in deleted:
            data = json.loads(d.serialized_data)[0]['fields']
            try:
                tag_name = teams[data['team']]
            except KeyError:
                tag_name = data['tag_name']
            if tag_name == tag:
                eids.add(data['episode'])

        historic = Episode.objects.filter(id__in=eids)
        return historic

class Synonym(models.Model):
    name = models.CharField(max_length=255)
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    content_object = generic.GenericForeignKey('content_type', 'object_id')

    class Meta:
        unique_together = (('name', 'content_type'))

    def __unicode__(self):
        return self.name

class LocatedModel(models.Model):
    address_line1 = models.CharField("Address line 1", max_length = 45,
                                     blank=True, null=True)
    address_line2 = models.CharField("Address line 2", max_length = 45,
                                     blank=True, null=True)
    city = models.CharField(max_length = 50, blank = True)
    county = models.CharField("County", max_length = 40,
                              blank=True, null=True)
    post_code = models.CharField("Post Code", max_length = 10,
                                 blank=True, null=True)

    class Meta:
        abstract = True

class GP(LocatedModel):
    name = models.CharField(blank=True, null=True, max_length=255)
    tel1 = models.CharField(blank=True, null=True, max_length=50)
    tel2 = models.CharField(blank=True, null=True, max_length=50)

class CommunityNurse(LocatedModel):
    name = models.CharField(blank=True, null=True, max_length=255)
    tel1 = models.CharField(blank=True, null=True, max_length=50)
    tel2 = models.CharField(blank=True, null=True, max_length=50)

option_models = {}

model_names = options.model_names

for name in model_names:
    class_name = name.capitalize() # TODO handle camelcase properly
    bases = (models.Model,)
    attrs = {
        'name': models.CharField(max_length=255, unique=True),
        'synonyms': generic.GenericRelation('Synonym'),
        'Meta': type('Meta', (object,), {'ordering': ['name']}),
        '__unicode__': lambda self: self.name,
        '__module__': __name__,
    }
    option_models[name] = type(class_name, bases, attrs)

# TODO
@receiver(models.signals.post_save, sender=Patient)
def create_patient_singletons(sender, **kwargs):
    if kwargs['created']:
        patient = kwargs['instance']
        for subclass in PatientSubrecord.__subclasses__():
            if subclass._is_singleton:
                obj = subclass.objects.create(patient=patient)

@receiver(models.signals.post_save, sender=Episode)
def create_episode_singletons(sender, **kwargs):
    if kwargs['created']:
        episode = kwargs['instance']
        for subclass in EpisodeSubrecord.__subclasses__():
            if subclass._is_singleton:
                subclass.objects.create(episode=episode)
<commit_msg>Single tag as mine is active now -because we're not pulling everything<commit_after>"""
OPAL Models!
"""
import collections
import json

from django.conf import settings
from django.db import models
from django.contrib import auth
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic
from django.dispatch import receiver
import reversion

from opal.utils import stringport, camelcase_to_underscore
from opal.utils.fields import ForeignKeyOrFreeText
from opal import exceptions, managers

# Imported models from module.

from opal.models.mixins import UpdatesFromDictMixin

options = stringport(settings.OPAL_OPTIONS_MODULE)

from django.contrib.auth.models import User

class UserProfile(models.Model):
    """
    Profile for our user
    """
    user                  = models.ForeignKey(User, unique=True)
    force_password_change = models.BooleanField(default=True)
    can_extract           = models.BooleanField(default=False)
    readonly              = models.BooleanField(default=False)

class Filter(models.Model):
    """
    Saved filters for users extracting data.
    """
    user     = models.ForeignKey(User)
    name     = models.CharField(max_length=200)
    criteria = models.TextField()

    def to_dict(self):
        return dict(id=self.pk, name=self.name, criteria=json.loads(self.criteria))

    def update_from_dict(self, data):
        self.criteria = json.dumps(data['criteria'])
        self.name = data['name']
        self.save()

class Patient(models.Model):
    def __unicode__(self):
        demographics = self.demographics_set.get()
        return '%s | %s' % (demographics.hospital_number, demographics.name)

    def create_episode(self):
        if self.get_active_episode() is None:
            return self.episode_set.create()
        else:
            raise exceptions.APIError('Patient %s already has active episode' % self)

    def get_active_episode(self):
        for episode in self.episode_set.order_by('id').reverse():
            if episode.is_active():
                return episode
        return None

    def to_dict(self, user):
        active_episode = self.get_active_episode()
        d = {
            'id': self.id,
            'episodes': {episode.id: episode.to_dict(user) for episode in self.episode_set.all()},
            'active_episode_id': active_episode.id if active_episode else None,
            }

        for model in PatientSubrecord.__subclasses__():
            subrecords = model.objects.filter(patient_id=self.id)
            d[model.get_api_name()] = [subrecord.to_dict(user) for subrecord in subrecords]
        return d

    def update_from_demographics_dict(self, demographics_data, user):
        demographics = self.demographics_set.get()
        demographics.update_from_dict(demographics_data, user)

class Episode(UpdatesFromDictMixin, models.Model):
    """
    An individual episode of care.

    A patient may have many episodes of care, but this maps to one occasion
    on which they found themselves on "The List".
    """
    patient = models.ForeignKey(Patient)
    active  = models.BooleanField(default=False)
    date_of_admission = models.DateField(null=True, blank=True)
    # TODO rename to date_of_discharge?
    discharge_date = models.DateField(null=True, blank=True)
    consistency_token = models.CharField(max_length=8)

    objects = managers.EpisodeManager()

    def __unicode__(self):
        demographics = self.patient.demographics_set.get()

        return '%s | %s | %s' % (demographics.hospital_number,
                                 demographics.name,
                                 self.date_of_admission)

    def is_active(self):
        # TODO Depreciate this.
        #
        # This is only here for API compatability.
        # Don't use me!
        return self.active

    def set_tag_names(self, tag_names, user):
        """
        1. Blitz dangling tags not in our current dict.
        2. Add new tags.
        3. Make sure that we set the Active boolean appropriately
        4. There is no step 4.
        """
        original_tag_names = self.get_tag_names(user)

        for tag_name in original_tag_names:
            if tag_name not in tag_names:
                params = {'team__name': tag_name}
                if tag_name == 'mine':
                    params['user'] = user
                self.tagging_set.get(**params).delete()

        for tag_name in tag_names:
            if tag_name not in original_tag_names:
                params = {'team': Team.objects.get(name=tag_name)}
                if tag_name == 'mine':
                    params['user'] = user
                self.tagging_set.create(**params)

        if len(tag_names) < 1:
            self.active = False
        elif tag_names == ['mine']:
            self.active = True
        elif not self.active:
            self.active = True
        self.save()

    def tagging_dict(self, user):
        if self.tagging_set.count() == 0:
            return [{}]
        td = [{t.team.name: True for t in self.tagging_set.exclude(team__name='mine')}]
        if self.tagging_set.filter(team__name='mine', user=user).count() > 0:
            td[0]['mine'] = True
        td[0]['id'] = self.id
        return td

    def get_tag_names(self, user):
        return [t.team.name for t in self.tagging_set.all() if t.user in (None, user)]

    def to_dict(self, user, shallow=False, with_context=False):
        """
        Serialisation to JSON for Episodes
        """
        d = {
            'id'               : self.id,
            'active'           : self.active,
            'date_of_admission': self.date_of_admission,
            'discharge_date'   : self.discharge_date,
            'consistency_token': self.consistency_token
            }
        if shallow:
            return d

        for model in PatientSubrecord.__subclasses__():
            subrecords = model.objects.filter(patient_id=self.patient.id)
            d[model.get_api_name()] = [subrecord.to_dict(user)
                                       for subrecord in subrecords]
        for model in EpisodeSubrecord.__subclasses__():
            subrecords = model.objects.filter(episode_id=self.id)
            d[model.get_api_name()] = [subrecord.to_dict(user)
                                       for subrecord in subrecords]

        d['tagging'] = self.tagging_dict(user)
        d['prev_episodes'] = []
        d['next_episodes'] = []

        if self.date_of_admission:
            eset = self.patient.episode_set

            d['prev_episodes'] = [
                e.to_dict(user, shallow=True)
                for e in
                eset.filter(date_of_admission__lt=self.date_of_admission).order_by('date_of_admission')]
            d['next_episodes'] = [
                e.to_dict(user, shallow=True)
                for e in
                eset.filter(date_of_admission__gt=self.date_of_admission).order_by('date_of_admission')]

        return d

    def update_from_location_dict(self, location_data, user):
        # TODO Completely depreciate this.
        location = self.location_set.get()
        location.update_from_dict(location_data, user)

class ContactNumber(models.Model):
    name = models.CharField(max_length=255)
    number = models.CharField(max_length=255)

    def __unicode__(self):
        return u'{0}: {1}'.format(self.name, self.number)

class Team(models.Model):
    name = models.CharField(max_length=250, help_text="This should only have letters and underscores")
    title = models.CharField(max_length=250)
    parent = models.ForeignKey('self', blank=True, null=True)
    active = models.BooleanField(default=True)
    order = models.IntegerField(blank=True, null=True)
    useful_numbers = models.ManyToManyField(ContactNumber, blank=True)

    def __unicode__(self):
        return self.title

    # TODO depreciate this and refactor accordingly
    @classmethod
    def to_TAGS(klass):
        from opal.utils import Tag
        teams = klass.objects.filter(active=True).order_by('order')
        tags = collections.OrderedDict()
        for team in teams:
            if not team.parent:
                tags[team.name] = Tag(team.name, team.title, [])
        for team in teams:
            if team.parent:
                tags[team.parent.name].subtags.append(Tag(team.name, team.title, None))
        return tags.values()

class Subrecord(UpdatesFromDictMixin, models.Model):
    consistency_token = models.CharField(max_length=8)

    _is_singleton = False

    class Meta:
        abstract = True

    def __unicode__(self):
        return u'{0}: {1}'.format(self.get_api_name(), self.id)

    @classmethod
    def get_api_name(cls):
        return camelcase_to_underscore(cls._meta.object_name)

    @classmethod
    def get_display_name(cls):
        if hasattr(cls, '_title'):
            return cls._title
        else:
            return cls._meta.object_name

    @classmethod
    def build_field_schema(cls):
        field_schema = []
        for fieldname in cls._get_fieldnames_to_serialize():
            if fieldname in ['id', 'patient_id', 'episode_id']:
                continue
            elif fieldname.endswith('_fk_id'):
                continue
            elif fieldname.endswith('_ft'):
                continue

            getter = getattr(cls, 'get_field_type_for_' + fieldname, None)
            if getter is None:
                field = cls._get_field_type(fieldname)
                if field in [models.CharField, ForeignKeyOrFreeText]:
                    field_type = 'string'
                else:
                    field_type = camelcase_to_underscore(field.__name__[:-5])
            else:
                field_type = getter()
            lookup_list = None
            if cls._get_field_type(fieldname) == ForeignKeyOrFreeText:
                fld = getattr(cls, fieldname)
                lookup_list = camelcase_to_underscore(fld.foreign_model.__name__)
            field_schema.append({'name': fieldname,
                                 'type': field_type,
                                 'lookup_list': lookup_list})
        return field_schema

    def _to_dict(self, user, fieldnames):
        """
        Allow a subset of FIELDNAMES
        """
        d = {}
        for name in fieldnames:
            getter = getattr(self, 'get_' + name, None)
            if getter is not None:
                value = getter(user)
            else:
                value = getattr(self, name)
            d[name] = value

        return d

    def to_dict(self, user):
        return self._to_dict(user, self._get_fieldnames_to_serialize())

class PatientSubrecord(Subrecord):
    patient = models.ForeignKey(Patient)

    class Meta:
        abstract = True

class EpisodeSubrecord(Subrecord):
    episode = models.ForeignKey(Episode, null=True)  # TODO make null=False

    class Meta:
        abstract = True

class Tagging(models.Model):
    _is_singleton = True
    _title = 'Teams'

    team = models.ForeignKey(Team, blank=True, null=True)
    user = models.ForeignKey(auth.models.User, null=True)
    episode = models.ForeignKey(Episode, null=True) # TODO make null=False

    def __unicode__(self):
        if self.user is not None:
            return 'User: %s - %s' % (self.user.username, self.team.name)
        else:
            return self.team.name

    @staticmethod
    def get_api_name():
        return 'tagging'

    @staticmethod
    def get_display_name():
        return 'Teams'

    @staticmethod
    def build_field_schema():
        teams = [{'name': t.name, 'type':'boolean'} for t in Team.objects.filter(active=True)]
        return teams
#        return [dict(name='team__name', type='string')]

    @classmethod
    def historic_tags_for_episodes(cls, episodes):
        """
        Given a list of episodes, return a dict indexed by episode id
        that contains historic tags for those episodes.
        """
        teams = {t.id: t.name for t in Team.objects.all()}
        deleted = reversion.get_deleted(cls)
        historic = collections.defaultdict(dict)
        for d in deleted:
            data = json.loads(d.serialized_data)[0]['fields']
            if data['episode'] in episodes:
                try:
                    tag_name = teams[data['team']]
                except KeyError:
                    tag_name = data['tag_name']
                historic[data['episode']][tag_name] = True
        return historic

    @classmethod
    def historic_episodes_for_tag(cls, tag):
        """
        Given a TAG return a list of episodes that have historically been
        tagged with it.
        """
        teams = {t.id: t.name for t in Team.objects.all()}
        deleted = reversion.get_deleted(cls)
        eids = set()
        for d in deleted:
            data = json.loads(d.serialized_data)[0]['fields']
            try:
                tag_name = teams[data['team']]
            except KeyError:
                tag_name = data['tag_name']
            if tag_name == tag:
                eids.add(data['episode'])

        historic = Episode.objects.filter(id__in=eids)
        return historic

class Synonym(models.Model):
    name = models.CharField(max_length=255)
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    content_object = generic.GenericForeignKey('content_type', 'object_id')

    class Meta:
        unique_together = (('name', 'content_type'))

    def __unicode__(self):
        return self.name

class LocatedModel(models.Model):
    address_line1 = models.CharField("Address line 1", max_length = 45,
                                     blank=True, null=True)
    address_line2 = models.CharField("Address line 2", max_length = 45,
                                     blank=True, null=True)
    city = models.CharField(max_length = 50, blank = True)
    county = models.CharField("County", max_length = 40,
                              blank=True, null=True)
    post_code = models.CharField("Post Code", max_length = 10,
                                 blank=True, null=True)

    class Meta:
        abstract = True

class GP(LocatedModel):
    name = models.CharField(blank=True, null=True, max_length=255)
    tel1 = models.CharField(blank=True, null=True, max_length=50)
    tel2 = models.CharField(blank=True, null=True, max_length=50)

class CommunityNurse(LocatedModel):
    name = models.CharField(blank=True, null=True, max_length=255)
    tel1 = models.CharField(blank=True, null=True, max_length=50)
    tel2 = models.CharField(blank=True, null=True, max_length=50)

option_models = {}

model_names = options.model_names

for name in model_names:
    class_name = name.capitalize() # TODO handle camelcase properly
    bases = (models.Model,)
    attrs = {
        'name': models.CharField(max_length=255, unique=True),
        'synonyms': generic.GenericRelation('Synonym'),
        'Meta': type('Meta', (object,), {'ordering': ['name']}),
        '__unicode__': lambda self: self.name,
        '__module__': __name__,
    }
    option_models[name] = type(class_name, bases, attrs)

# TODO
@receiver(models.signals.post_save, sender=Patient)
def create_patient_singletons(sender, **kwargs):
    if kwargs['created']:
        patient = kwargs['instance']
        for subclass in PatientSubrecord.__subclasses__():
            if subclass._is_singleton:
                obj = subclass.objects.create(patient=patient)

@receiver(models.signals.post_save, sender=Episode)
def create_episode_singletons(sender, **kwargs):
    if kwargs['created']:
        episode = kwargs['instance']
        for subclass in EpisodeSubrecord.__subclasses__():
            if subclass._is_singleton:
                subclass.objects.create(episode=episode)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright (c) Juniper Networks, Inc., 2015-2016.
# All rights reserved.
#
# coding: utf-8
# Authors: psagrera@juniper.net
# Version: 1.0 20150122

import IxNetwork
import os
from time import sleep
import sys
from pprint import pprint

class ContinuableError(RuntimeError):
    ROBOT_CONTINUE_ON_FAILURE = True

class FatalError(RuntimeError):
    ROBOT_EXIT_ON_FAILURE = True

class pybot_jixia (object):

    # -----------------------------------------------------------------------
    # CONSTRUCTOR
    # -----------------------------------------------------------------------

        def __init__(self,**kvargs):

            self.server = kvargs['ixia_server'] # ixia server currently available 172.30.153.72
            self.ixia_port = kvargs['ixia_port']
            self.ixNet = IxNetwork.IxNet ()
            self.ixNet.setTimeout = 20

    # -----------------------------------------------------------------------
    # METHODS
    # -----------------------------------------------------------------------

        # Function to open a connection to ixia server

        def ixia_connect (self):

            """
                Opens a connection to the device (ixia) using existing server/port
            """
            getVersion = self.ixNet.getVersion()

            try:

               #print "Establishing connection to ixia.......,please be patient"
               #print "Connection established properly"
               return  self.ixNet.connect(self.server,'-port',self.ixia_port,'-version',getVersion, '-clientusername')

            except IxNetwork.IxNetError as ix_conn_err:
                #print " Problems establishing connection to IXIA: %s" % (ix_conn_err)
                raise FatalError("Expected result is True but was False,test will be stopped at this point")
                return False

        # Function to close a connection to the ixia server

        def ixia_disconnect (self):

            """
                Close a connection to the device
            """
            try:

                #print "Closing connection to ixia.......,bye"
                #print self.ixNet.execute('ClearOwnershipForAllPorts')
                return self.ixNet.disconnect ()

            except IxNetwork.IxNetError as ix_conn_err:

                #print "Problems found closing connection to IXIA: %s" % (ix_conn_err)
                raise FatalError("Expected result is True but was False,test will go on")

        # Function to initialize IxNetwork

        def ixia_initialize (self):

            try:
                if self.ixia_return(self.ixNet.execute('newConfig')) is False:
                    print '---- Could not initialize IxNetwork: [NOK]'
                    raise FatalError("Expected result is True but was False,test will stopped")
            except:
                raise FatalError("Expected result is True but was False,test will stopped")

            print '---- IxNetwork initialization: [OK]'

        # Specific function to convert socket return values in boolean

        def ixia_return (self,tclreturn):

            if tclreturn == '::ixNet::OK':
                return True
            else:
                return False

        # Function to start all protocols from ixia

        def start_protocols (self):

            print "Starting protocols......."

            try:
                self.ixNet.execute('startAllProtocols')

            except IxNetwork.IxNetError as ix_start_err:
                print '---- Could not start protocols %s' % (ix_start_err)
                raise FatalError("Expected result is True but was False,test will stopped")

        # Function to running traffic from ixia

        def running_traffic (self):

            traffic = self.ixNet.getList(self.ixNet.getRoot(),'traffic')

            try:
                if self.ixia_return(self.ixNet.execute('apply', traffic[0])) is False:
                    print '---- Could not create stats gathering at IxNetwork: [NOK] %s ' % (ix_run_err)
                    raise FatalError("Expected result is True but was False,test will stopped")

            except IxNetwork.IxNetError as ix_run_err:
                  # Terminate protocols and abort existing connections
                print '---- Could not create stats gathering at IxNetwork: [NOK] %s ' % (ix_run_err)
                raise FatalError("Expected result is True but was False,test will stopped")

            print '---- IxNetwork traffic data gathering infra created: [OK]'

            # Sending stateless traffic
            try:
                if self.ixia_return(self.ixNet.execute('startStatelessTraffic', traffic[0])) is False:
                    print '---- Could not start stateless traffic at IxNetwork: [NOK]'
                    raise FatalError("Expected result is True but was False,test will stopped")

            except IxNetwork.IxNetError as ix_run_err:
                # Terminate protocols and abort existing connections
                print '---- Could not start stateless traffic at IxNetwork: [NOK] %s' % (ix_run_err)
                raise FatalError("Expected result is True but was False,test will stopped")

            print '---- IxNetwork sending stateless traffic started: [OK]'

        # Function to stopping traffic from ixia

        def stopping_traffic (self):

            # Stop sending traffic at IxNetwork

            traffic = self.ixNet.getList(self.ixNet.getRoot(),'traffic')
            print 'Stopping sending 60s traffic at IxNetwork for 2nd test'

            # Stop sending stateless traffic

            try:
                if self.ixia_return(self.ixNet.execute('stopStatelessTraffic', traffic[0])) is False:
                    print '---- Could not stop stateless traffic at IxNetwork: [NOK]'
                    raise FatalError("Expected result is True but was False,test will stopped")
            except IxNetwork.IxNetError as ix_stop_err:
                # Terminate protocols and abort existing connections
                print '---- Could not stop stateless traffic at IxNetwork: [NOK] %s' %(ix_stop_err)
                raise FatalError("Expected result is True but was False,test will stopped")
            print '---- IxNetwork stopping 60s stateless traffic: [OK]'

        # Function to load ixia config file

        def load_config (self,**kvargs):

            print "Loading configuration......."

            file_path = kvargs['ixia_file']

            try:

                if self.ixia_return (self.ixNet.execute('loadConfig',self.ixNet.readFrom(file_path))) is False:
                    sys.exit('---- Could not load IxNetwork config: [NOK]')

            except IxNetwork.IxNetError as ix_conn_err:

                print "Problems loading configuration into IXIA: %s" % (ix_conn_err)
                raise FatalError("Expected result is True but was False,test will go on")

            print '---- IxNetwork config loading: [OK]'

        def stopping_protocols (self):

            print '!!! Stopping protocols'

            try:
                self.ixNet.execute('stopAllProtocols')
            except IxNetwork.IxNetError as ix_stoppro_err:
                print '!!! Could not stop protocols and close existing connections'
                raise FatalError("Expected result is True but was False,test will stopped")

        # Function to gather stats from ixia

        def gather_stats (self, **kwargs):
            view = kwargs['view']
            verbose = True
            if 'verbose' in kwargs.keys():
                verbose = kwargs['verbose']

            try:

                Stats = self.ixNet.getFilteredList(self.ixNet.getRoot()+'statistics', 'view', '-caption', view)
                #print '---- Retrieved IxNetwork loss percentage stats: [OK]'

            except IxNetwork.IxNetError as ix_gather_err:
                #print "Problems gathering stats from IXIA: %s" % (ix_gather_err)
                raise FatalError("Expected result is True but was False,test will go on")

            result = []
            for Stat in Stats:
                tmp = {}
                tmp['columnCaption'] = self.ixNet.getAttribute(Stat+'/page', '-columnCaptions')
                tmp['RowVal'] = self.ixNet.getAttribute(Stat+'/page', '-rowValues')
                result.append(tmp)

            if verbose:
                for tmp in result:
                    print tmp['columnCaption']
                    for tmp2 in tmp['RowVal']:
                        print tmp2
                return True
            else:
                return result

        def clear_stats (self):

            print "Clearing statistics......."

            try:
                self.ixNet.execute('clearStats')

            except IxNetwork.IxNetError as ix_clear_stats_err:
                print '---- Could not clear statistics %s' % (ix_clear_stats_err)
                raise FatalError("Expected result is True but was False,test will stopped")

#kk = pybot_jixia (ixia_server="172.30.153.72",ixia_port=8009)
#kk.ixia_connect ()
#pprint(kk.gather_stats(view='Traffic Item Statistics',verbose=False))
#pprint(kk.gather_stats(view='BGP Aggregated Statistics',verbose=False))
#kk.ixia_disconnect ()#
<commit_msg>Removing unused libraries<commit_after><|endoftext|>