<commit_before>"""
This module contains the main class that describes an isotherm.
"""

import hashlib

import pandas

import pygaps

from ..utilities.exceptions import ParameterError
from ..utilities.unit_converter import _MASS_UNITS
from ..utilities.unit_converter import _MATERIAL_MODE
from ..utilities.unit_converter import _MOLAR_UNITS
from ..utilities.unit_converter import _PRESSURE_MODE
from ..utilities.unit_converter import _PRESSURE_UNITS
from ..utilities.unit_converter import _VOLUME_UNITS

class Isotherm(object):
    """
    Class which contains the general data for an isotherm, real or model.

    The isotherm class is the parent class that both PointIsotherm and
    ModelIsotherm inherit. It is designed to contain the information about
    an isotherm, such as material, adsorbate, data units etc., but without
    any of the data itself.

    Think of this class as a extended python dictionary.

    Parameters
    ----------

    loading_key : str
        The title of the pressure data in the DataFrame provided.
    pressure_key : str
        The title of the loading data in the DataFrame provided.
    sample_name : str
        Name of the sample on which the isotherm is measured.
    sample_batch : str
        Batch (or identifier) of the sample on which the isotherm is measured.
    adsorbate : str
        The adsorbate used in the experiment.
    t_exp : float
        Experiment temperature.

    Other Parameters
    ----------------
    adsorbent_basis : str, optional
        Whether the adsorption is read in terms of either 'per volume'
        'per molar amount' or 'per mass' of material.
    adsorbent_unit : str, optional
        Unit in which the adsorbent basis is expressed.
    loading_basis : str, optional
        Whether the adsorbed material is read in terms of either 'volume'
        'molar' or 'mass'.
    loading_unit : str, optional
        Unit in which the loading basis is expressed.
    pressure_mode : str, optional
        The pressure mode, either 'absolute' pressures or 'relative' in
        the form of p/p0.
    pressure_unit : str, optional
        Unit of pressure.

    Notes
    -----

    The class is also used to prevent duplication of code within the child
    classes, by calling the common inherited function before any other specific
    implementation additions.

    The minimum arguments required to instantiate the class are
    ``sample_name``, ``sample_batch``, ``t_exp', ``adsorbate``.
    """

    _named_params = [
        # required
        'sample_name',
        'sample_batch',
        't_exp',
        'adsorbate',

        # others
        'user',
        'machine',
        'exp_type',
        'date',
        'is_real',
        't_act',
        'lab',
        'project',
        'comment',
    ]

    _unit_params = [
        'pressure_unit',
        'pressure_mode',
        'adsorbent_unit',
        'adsorbent_basis',
        'loading_unit',
        'loading_basis',
    ]

    _db_columns = ['id'] + _named_params
    _id_params = _named_params + _unit_params + ['other_properties']

    def __init__(self,
                 loading_key=None,
                 pressure_key=None,

                 adsorbent_basis="mass",
                 adsorbent_unit="g",
                 loading_basis="molar",
                 loading_unit="mmol",
                 pressure_mode="absolute",
                 pressure_unit="bar",

                 **isotherm_parameters):
        """
        Instantiation is done by passing a dictionary with the parameters,
        as well as the info about units, modes and data columns.
        """

        # Checks
        if any(k not in isotherm_parameters
               for k in ('sample_name', 'sample_batch', 't_exp', 'adsorbate')):
            raise ParameterError(
                "Isotherm MUST have the following properties:"
                "'sample_name', 'sample_batch', 't_exp', 'adsorbate'")

        # Basis and mode
        if adsorbent_basis is None or pressure_mode is None or loading_basis is None:
            raise ParameterError(
                "One of the modes or bases is not specified.")

        if adsorbent_basis not in _MATERIAL_MODE:
            raise ParameterError(
                "Basis selected for adsorbent is not an option. See viable "
                "values: {0}".format(_MATERIAL_MODE))

        if loading_basis not in _MATERIAL_MODE:
            raise ParameterError(
                "Basis selected for loading is not an option. See viable "
                "values: {0}".format(_MATERIAL_MODE))

        if pressure_mode not in _PRESSURE_MODE:
            raise ParameterError(
                "Mode selected for pressure is not an option. See viable "
                "values: {0}".format(_PRESSURE_MODE))

        # Units
        if loading_unit is None or adsorbent_unit is None:
            raise ParameterError(
                "One of the units is not specified.")

        if loading_unit not in _MATERIAL_MODE[loading_basis]:
            raise ParameterError(
                "Unit selected for loading is not an option. See viable "
                "values: {0}".format(_MOLAR_UNITS))

        if pressure_mode == 'absolute' and pressure_unit not in _PRESSURE_UNITS:
            raise ParameterError(
                "Unit selected for pressure is not an option. See viable "
                "values: {0}".format(_PRESSURE_UNITS))

        if adsorbent_unit not in _MATERIAL_MODE[adsorbent_basis]:
            raise ParameterError(
                "Unit selected for adsorbent is not an option. See viable "
                "values: {0} {1}".format(_VOLUME_UNITS, _MASS_UNITS))

        # Column titles
        if None in [loading_key, pressure_key]:
            raise ParameterError(
                "Pass loading_key and pressure_key, the names of the loading and"
                " pressure columns in the DataFrame, to the constructor.")

        #: Basis for the adsorbent.
        self.adsorbent_basis = str(adsorbent_basis)
        #: Unit for the adsorbent.
        self.adsorbent_unit = str(adsorbent_unit)
        #: Basis for the loading.
        self.loading_basis = str(loading_basis)
        #: Units for loading.
        self.loading_unit = str(loading_unit)
        #: Mode for the pressure.
        self.pressure_mode = str(pressure_mode)
        if pressure_mode == 'relative':
            #: Units for pressure.
            self.pressure_unit = None
        else:
            #: Units for pressure.
            self.pressure_unit = str(pressure_unit)

        # Save column names
        #: Name of column in the dataframe that contains adsorbed amount.
        self.loading_key = loading_key

        #: Name of column in the dataframe that contains pressure.
        self.pressure_key = pressure_key

        # Must-have properties of the isotherm
        #

        # ID
        self.id = isotherm_parameters.pop('id', None)
        #: Isotherm material name.
        self.sample_name = str(isotherm_parameters.pop('sample_name', None))
        #: Isotherm material batch.
        self.sample_batch = str(isotherm_parameters.pop('sample_batch', None))
        #: Isotherm experimental temperature.
        self.t_exp = float(isotherm_parameters.pop('t_exp', None))
        #: Isotherm adsorbate used.
        self.adsorbate = str(isotherm_parameters.pop('adsorbate', None))

        # Good-to-have properties of the isotherm
        #: Isotherm experiment date.
        self.date = None
        date = isotherm_parameters.pop('date', None)
        if date:
            self.date = str(date)

        #: Isotherm sample activation temperature.
        self.t_act = None
        t_act = isotherm_parameters.pop('t_act', None)
        if t_act:
            self.t_act = float(t_act)

        #: Isotherm lab.
        self.lab = None
        lab = isotherm_parameters.pop('lab', None)
        if lab:
            self.lab = str(lab)

        #: Isotherm comments.
        self.comment = None
        comment = isotherm_parameters.pop('comment', None)
        if comment:
            self.comment = str(comment)

        #
        # Other properties
        #: Isotherm user.
        self.user = None
        user = isotherm_parameters.pop('user', None)
        if user:
            self.user = str(user)

        #: Isotherm project.
        self.project = None
        project = isotherm_parameters.pop('project', None)
        if project:
            self.project = str(project)

        #: Isotherm machine used.
        self.machine = None
        machine = isotherm_parameters.pop('machine', None)
        if machine:
            self.machine = str(machine)

        #: Isotherm physicality (real or simulation).
        self.is_real = None
        is_real = isotherm_parameters.pop('is_real', None)
        if is_real:
            self.is_real = bool(is_real)

        #: Isotherm type (calorimetry/isotherm).
        self.exp_type = None
        exp_type = isotherm_parameters.pop('exp_type', None)
        if exp_type:
            self.exp_type = str(exp_type)

        # Save the rest of the properties as an extra dict
        # now that the named properties were taken out of
        #: Other properties of the isotherm.
        self.other_properties = isotherm_parameters

        # Finish instantiation process
        # (check if none in case its part of a Point/Model Isotherm instantiation)
        if not hasattr(self, '_instantiated'):
            self._instantiated = True
            if self.id is None:
                self._check_if_hash(True, [True])

    ##########################################################
    #   Overloaded and private functions

    def __setattr__(self, name, value):
        """
        We overload the usual class setter to make sure that the id is always
        representative of the data inside the isotherm.

        The '_instantiated' attribute gets set to true after isotherm __init__
        From then afterwards, each call to modify the isotherm properties
        recalculates the md5 hash.
        This is done to ensure uniqueness and also to allow isotherm objects to
        be easily compared to each other.
        """
        object.__setattr__(self, name, value)
        self._check_if_hash(name)

    def _check_if_hash(self, name, extra_params=[]):
        """Checks if the hash needs to be generated"""
        if getattr(self, '_instantiated', False) and name in self._id_params + extra_params:
            # Generate the unique id using md5
            self.id = None        # set the id to none for repeatability of json
            md_hasher = hashlib.md5(
                pygaps.isotherm_to_json(self).encode('utf-8'))
            self.id = md_hasher.hexdigest()

    def __eq__(self, other_isotherm):
        """
        We overload the equality operator of the isotherm. Since id's should be unique and
        representative of the data inside the isotherm, all we need to ensure equality
        is to compare the two hashes of the isotherms.
        """

        return self.id == other_isotherm.id

    ###########################################################
    #   Info functions

    def __str__(self):
        '''
        Prints a short summary of all the isotherm parameters.
        '''
        string = ""

        if self.is_real:
            string += ("Experimental isotherm" + '\n')
        else:
            string += ("Simulated isotherm" + '\n')

        # Required
        string += ("Material: " + str(self.sample_name) + '\n')
        string += ("Batch: " + str(self.sample_batch) + '\n')
        string += ("Adsorbate used: " + str(self.adsorbate) + '\n')
        string += ("Isotherm temperature: " + str(self.t_exp) + "K" + '\n')

        if self.exp_type:
            string += ("Isotherm type: " + str(self.exp_type) + '\n')
        if self.date:
            string += ("Isotherm date: " + str(self.date) + '\n')
        if self.machine:
            string += ("Machine: " + str(self.machine) + '\n')
        if self.user:
            string += ("User: " + str(self.user) + '\n')
        if self.t_act:
            string += ("Activation temperature: " +
                       str(self.t_act) + "°C" + '\n')
        if self.comment:
            string += ("Isotherm comments: " + str(self.comment) + '\n')

        # Units/basis
        string += ("Units: \n")
        string += ("Unit for loading: " + str(self.loading_unit) +
                   "/" + str(self.adsorbent_unit) + '\n')
        if self.pressure_mode == 'relative':
            string += ("Relative pressure \n")
        else:
            string += ("Unit for pressure: " + str(self.pressure_unit) + '\n')

        string += ("Other properties: \n")
        for prop in self.other_properties:
            string += (prop + ": " + str(self.other_properties[prop]) + '\n')

        return string

    def to_dict(self):
        """
        Returns a dictionary of the isotherm class
        Is the same dictionary that was used to create it.

        Returns
        -------
        dict
            Dictionary of all parameters.
        """
        parameter_dict = {}

        # Add named parameters
        for param in self._named_params + self._unit_params + ['id']:
            if hasattr(self, param):
                parameter_dict.update({param: getattr(self, param)})

        # Now add the rest
        parameter_dict.update(self.other_properties)

        return parameter_dict

    # Figure out the adsorption and desorption branches
    def _splitdata(self, _data):
        """
        Splits isotherm data into an adsorption and desorption part and
        adds a column to mark the transition between the two.
        """
        # Get a column where all increasing are False and all decreasing are True
        increasing = _data.loc[:, self.pressure_key].diff().fillna(0) < 0
        # Get the first inflexion point (assume where des starts)
        inflexion = increasing.idxmax()

        # If there is an inflexion point
        if inflexion != 0:
            # If the first point is where the isotherm starts decreasing
            # Then it is a complete desorption curve
            if inflexion == 1:
                inflexion = 0

            # Set all instances after the inflexion point to True
            increasing[inflexion:] = True

        # Return the new array with the branch column
        return pandas.concat([_data, increasing.rename('branch')], axis=1)
<commit_msg>fixed lab special field not recognised<commit_after>"""
This module contains the main class that describes an isotherm.
"""

import hashlib

import pandas

import pygaps

from ..utilities.exceptions import ParameterError
from ..utilities.unit_converter import _MASS_UNITS
from ..utilities.unit_converter import _MATERIAL_MODE
from ..utilities.unit_converter import _MOLAR_UNITS
from ..utilities.unit_converter import _PRESSURE_MODE
from ..utilities.unit_converter import _PRESSURE_UNITS
from ..utilities.unit_converter import _VOLUME_UNITS

class Isotherm(object):
    """
    Class which contains the general data for an isotherm, real or model.

    The isotherm class is the parent class that both PointIsotherm and
    ModelIsotherm inherit. It is designed to contain the information about
    an isotherm, such as material, adsorbate, data units etc., but without
    any of the data itself.

    Think of this class as a extended python dictionary.

    Parameters
    ----------

    loading_key : str
        The title of the pressure data in the DataFrame provided.
    pressure_key : str
        The title of the loading data in the DataFrame provided.
    sample_name : str
        Name of the sample on which the isotherm is measured.
    sample_batch : str
        Batch (or identifier) of the sample on which the isotherm is measured.
    adsorbate : str
        The adsorbate used in the experiment.
    t_exp : float
        Experiment temperature.

    Other Parameters
    ----------------
    adsorbent_basis : str, optional
        Whether the adsorption is read in terms of either 'per volume'
        'per molar amount' or 'per mass' of material.
    adsorbent_unit : str, optional
        Unit in which the adsorbent basis is expressed.
    loading_basis : str, optional
        Whether the adsorbed material is read in terms of either 'volume'
        'molar' or 'mass'.
    loading_unit : str, optional
        Unit in which the loading basis is expressed.
    pressure_mode : str, optional
        The pressure mode, either 'absolute' pressures or 'relative' in
        the form of p/p0.
    pressure_unit : str, optional
        Unit of pressure.

    Notes
    -----

    The class is also used to prevent duplication of code within the child
    classes, by calling the common inherited function before any other specific
    implementation additions.

    The minimum arguments required to instantiate the class are
    ``sample_name``, ``sample_batch``, ``t_exp', ``adsorbate``.
    """

    _named_params = [
        # required
        'sample_name',
        'sample_batch',
        't_exp',
        'adsorbate',

        # others
        'user',
        'machine',
        'exp_type',
        'date',
        'is_real',
        't_act',
        'lab',
        'project',
        'comment',
    ]

    _unit_params = [
        'pressure_unit',
        'pressure_mode',
        'adsorbent_unit',
        'adsorbent_basis',
        'loading_unit',
        'loading_basis',
    ]

    _db_columns = ['id'] + _named_params
    _id_params = _named_params + _unit_params + ['other_properties']

    def __init__(self,
                 loading_key=None,
                 pressure_key=None,

                 adsorbent_basis="mass",
                 adsorbent_unit="g",
                 loading_basis="molar",
                 loading_unit="mmol",
                 pressure_mode="absolute",
                 pressure_unit="bar",

                 **isotherm_parameters):
        """
        Instantiation is done by passing a dictionary with the parameters,
        as well as the info about units, modes and data columns.
        """

        # Checks
        if any(k not in isotherm_parameters
               for k in ('sample_name', 'sample_batch', 't_exp', 'adsorbate')):
            raise ParameterError(
                "Isotherm MUST have the following properties:"
                "'sample_name', 'sample_batch', 't_exp', 'adsorbate'")

        # Basis and mode
        if adsorbent_basis is None or pressure_mode is None or loading_basis is None:
            raise ParameterError(
                "One of the modes or bases is not specified.")

        if adsorbent_basis not in _MATERIAL_MODE:
            raise ParameterError(
                "Basis selected for adsorbent is not an option. See viable "
                "values: {0}".format(_MATERIAL_MODE))

        if loading_basis not in _MATERIAL_MODE:
            raise ParameterError(
                "Basis selected for loading is not an option. See viable "
                "values: {0}".format(_MATERIAL_MODE))

        if pressure_mode not in _PRESSURE_MODE:
            raise ParameterError(
                "Mode selected for pressure is not an option. See viable "
                "values: {0}".format(_PRESSURE_MODE))

        # Units
        if loading_unit is None or adsorbent_unit is None:
            raise ParameterError(
                "One of the units is not specified.")

        if loading_unit not in _MATERIAL_MODE[loading_basis]:
            raise ParameterError(
                "Unit selected for loading is not an option. See viable "
                "values: {0}".format(_MOLAR_UNITS))

        if pressure_mode == 'absolute' and pressure_unit not in _PRESSURE_UNITS:
            raise ParameterError(
                "Unit selected for pressure is not an option. See viable "
                "values: {0}".format(_PRESSURE_UNITS))

        if adsorbent_unit not in _MATERIAL_MODE[adsorbent_basis]:
            raise ParameterError(
                "Unit selected for adsorbent is not an option. See viable "
                "values: {0} {1}".format(_VOLUME_UNITS, _MASS_UNITS))

        # Column titles
        if None in [loading_key, pressure_key]:
            raise ParameterError(
                "Pass loading_key and pressure_key, the names of the loading and"
                " pressure columns in the DataFrame, to the constructor.")

        #: Basis for the adsorbent.
        self.adsorbent_basis = str(adsorbent_basis)
        #: Unit for the adsorbent.
        self.adsorbent_unit = str(adsorbent_unit)
        #: Basis for the loading.
        self.loading_basis = str(loading_basis)
        #: Units for loading.
        self.loading_unit = str(loading_unit)
        #: Mode for the pressure.
        self.pressure_mode = str(pressure_mode)
        if pressure_mode == 'relative':
            #: Units for pressure.
            self.pressure_unit = None
        else:
            #: Units for pressure.
            self.pressure_unit = str(pressure_unit)

        # Save column names
        #: Name of column in the dataframe that contains adsorbed amount.
        self.loading_key = loading_key

        #: Name of column in the dataframe that contains pressure.
        self.pressure_key = pressure_key

        # Must-have properties of the isotherm
        #

        # ID
        self.id = isotherm_parameters.pop('id', None)
        #: Isotherm material name.
        self.sample_name = str(isotherm_parameters.pop('sample_name', None))
        #: Isotherm material batch.
        self.sample_batch = str(isotherm_parameters.pop('sample_batch', None))
        #: Isotherm experimental temperature.
        self.t_exp = float(isotherm_parameters.pop('t_exp', None))
        #: Isotherm adsorbate used.
        self.adsorbate = str(isotherm_parameters.pop('adsorbate', None))

        # Good-to-have properties of the isotherm
        #: Isotherm experiment date.
        self.date = None
        date = isotherm_parameters.pop('date', None)
        if date:
            self.date = str(date)

        #: Isotherm sample activation temperature.
        self.t_act = None
        t_act = isotherm_parameters.pop('t_act', None)
        if t_act:
            self.t_act = float(t_act)

        #: Isotherm lab.
        self.lab = None
        lab = isotherm_parameters.pop('lab', None)
        if lab:
            self.lab = str(lab)

        #: Isotherm comments.
        self.comment = None
        comment = isotherm_parameters.pop('comment', None)
        if comment:
            self.comment = str(comment)

        #
        # Other properties
        #: Isotherm user.
        self.user = None
        user = isotherm_parameters.pop('user', None)
        if user:
            self.user = str(user)

        #: Isotherm project.
        self.project = None
        project = isotherm_parameters.pop('project', None)
        if project:
            self.project = str(project)

        #: Isotherm machine used.
        self.machine = None
        machine = isotherm_parameters.pop('machine', None)
        if machine:
            self.machine = str(machine)

        #: Isotherm physicality (real or simulation).
        self.is_real = None
        is_real = isotherm_parameters.pop('is_real', None)
        if is_real:
            self.is_real = bool(is_real)

        #: Isotherm type (calorimetry/isotherm).
        self.exp_type = None
        exp_type = isotherm_parameters.pop('exp_type', None)
        if exp_type:
            self.exp_type = str(exp_type)

        # Save the rest of the properties as an extra dict
        # now that the named properties were taken out of
        #: Other properties of the isotherm.
        self.other_properties = isotherm_parameters

        # Finish instantiation process
        # (check if none in case its part of a Point/Model Isotherm instantiation)
        if not hasattr(self, '_instantiated'):
            self._instantiated = True
            if self.id is None:
                self._check_if_hash(True, [True])

    ##########################################################
    #   Overloaded and private functions

    def __setattr__(self, name, value):
        """
        We overload the usual class setter to make sure that the id is always
        representative of the data inside the isotherm.

        The '_instantiated' attribute gets set to true after isotherm __init__
        From then afterwards, each call to modify the isotherm properties
        recalculates the md5 hash.
        This is done to ensure uniqueness and also to allow isotherm objects to
        be easily compared to each other.
        """
        object.__setattr__(self, name, value)
        self._check_if_hash(name)

    def _check_if_hash(self, name, extra_params=[]):
        """Checks if the hash needs to be generated"""
        if getattr(self, '_instantiated', False) and name in self._id_params + extra_params:
            # Generate the unique id using md5
            self.id = None        # set the id to none for repeatability of json
            md_hasher = hashlib.md5(
                pygaps.isotherm_to_json(self).encode('utf-8'))
            self.id = md_hasher.hexdigest()

    def __eq__(self, other_isotherm):
        """
        We overload the equality operator of the isotherm. Since id's should be unique and
        representative of the data inside the isotherm, all we need to ensure equality
        is to compare the two hashes of the isotherms.
        """

        return self.id == other_isotherm.id

    ###########################################################
    #   Info functions

    def __str__(self):
        '''
        Prints a short summary of all the isotherm parameters.
        '''
        string = ""

        if self.is_real:
            string += ("Experimental isotherm" + '\n')
        else:
            string += ("Simulated isotherm" + '\n')

        # Required
        string += ("Material: " + str(self.sample_name) + '\n')
        string += ("Batch: " + str(self.sample_batch) + '\n')
        string += ("Adsorbate used: " + str(self.adsorbate) + '\n')
        string += ("Isotherm temperature: " + str(self.t_exp) + "K" + '\n')

        if self.exp_type:
            string += ("Isotherm type: " + str(self.exp_type) + '\n')
        if self.date:
            string += ("Isotherm date: " + str(self.date) + '\n')
        if self.machine:
            string += ("Machine: " + str(self.machine) + '\n')
        if self.user:
            string += ("User: " + str(self.user) + '\n')
        if self.t_act:
            string += ("Activation temperature: " +
                       str(self.t_act) + "°C" + '\n')
        if self.lab:
            string += ("Source: " +
                       str(self.lab) + '\n')
        if self.comment:
            string += ("Isotherm comments: " + str(self.comment) + '\n')

        # Units/basis
        string += ("Units: \n")
        string += ("Unit for loading: " + str(self.loading_unit) +
                   "/" + str(self.adsorbent_unit) + '\n')
        if self.pressure_mode == 'relative':
            string += ("Relative pressure \n")
        else:
            string += ("Unit for pressure: " + str(self.pressure_unit) + '\n')

        string += ("Other properties: \n")
        for prop in self.other_properties:
            string += (prop + ": " + str(self.other_properties[prop]) + '\n')

        return string

    def to_dict(self):
        """
        Returns a dictionary of the isotherm class
        Is the same dictionary that was used to create it.

        Returns
        -------
        dict
            Dictionary of all parameters.
        """
        parameter_dict = {}

        # Add named parameters
        for param in self._named_params + self._unit_params + ['id']:
            if hasattr(self, param):
                parameter_dict.update({param: getattr(self, param)})

        # Now add the rest
        parameter_dict.update(self.other_properties)

        return parameter_dict

    # Figure out the adsorption and desorption branches
    def _splitdata(self, _data):
        """
        Splits isotherm data into an adsorption and desorption part and
        adds a column to mark the transition between the two.
        """
        # Get a column where all increasing are False and all decreasing are True
        increasing = _data.loc[:, self.pressure_key].diff().fillna(0) < 0
        # Get the first inflexion point (assume where des starts)
        inflexion = increasing.idxmax()

        # If there is an inflexion point
        if inflexion != 0:
            # If the first point is where the isotherm starts decreasing
            # Then it is a complete desorption curve
            if inflexion == 1:
                inflexion = 0

            # Set all instances after the inflexion point to True
            increasing[inflexion:] = True

        # Return the new array with the branch column
        return pandas.concat([_data, increasing.rename('branch')], axis=1)
<|endoftext|>


<commit_before>from __future__ import division

import argparse
import gc
import sys
import math
from collections import defaultdict
from decimal import Decimal

import pytest

from .stats import RunningStats
from .timers import compute_timer_precision
from .timers import default_timer
from .compat import XRANGE, PY3

def loadtimer(string):
    if "." not in string:
        raise argparse.ArgumentTypeError("Value for --benchmark-timer must be in dotted form. Eg: 'module.attr'.")
    mod, attr = string.rsplit(".", 1)
    if mod == 'pep418':
        if PY3:
            import time
            return getattr(time, attr)
        else:
            from . import pep418
            return getattr(pep418, attr)
    else:
        __import__(mod)
        mod = sys.modules[mod]
        return getattr(mod, attr)

def loadsort(string):
    if string not in ("min", "max", "mean", "stddev"):
        raise argparse.ArgumentTypeError("Value for --benchmark-sort must be one of: 'min', 'max', 'mean' or 'stddev'.")
    return string

def loadrounds(string):
    try:
        value = int(string)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(exc)
    else:
        if value < 1:
            raise argparse.ArgumentTypeError("Value for --benchmark-rounds must be at least 1.")
        return value

def pytest_addoption(parser):
    group = parser.getgroup("benchmark")
    group.addoption(
        "--benchmark-min-time",
        action="store", type=Decimal, default=Decimal("0.1"),
        help="Minimum time per round."
    )
    group.addoption(
        "--benchmark-max-time",
        action="store", type=Decimal, default=Decimal("1.0"),
        help="Maximum time to spend in a benchmark."
    )
    group.addoption(
        "--benchmark-min-rounds",
        action="store", type=loadrounds, default=5,
        help="Minimum rounds, even if total time would exceed `--max-time`."
    )
    group.addoption(
        "--benchmark-sort",
        action="store", type=loadsort, default="min",
        help="Column to sort on. Can be one of: 'min', 'max', 'mean' or 'stddev' ."
    )
    group.addoption(
        "--benchmark-timer",
        action="store", type=loadtimer, default=default_timer,
        help="Timer to use when measuring time."
    )
    group.addoption(
        "--benchmark-warmup",
        action="store_true", default=False,
        help="Runs the benchmarks two times. Discards data from the first run."
    )
    group.addoption(
        "--benchmark-disable-gc",
        action="store_true", default=False,
        help="Disable GC during benchmarks."
    )
    group.addoption(
        "--benchmark-skip",
        action="store_true", default=False,
        help="Skip running any benchmarks."
    )
    group.addoption(
        "--benchmark-only",
        action="store_true", default=False,
        help="Only run benchmarks."
    )

class BenchmarkStats(RunningStats):
    def __init__(self, name, group, scale):
        self.name = name
        self.group = group
        self.scale = scale
        super(BenchmarkStats, self).__init__()

    def __getitem__(self, key):
        return getattr(self, key)

    def update(self, duration):
        super(BenchmarkStats, self).update(duration / self.scale)

class BenchmarkFixture(object):
    _precisions = {}

    @classmethod
    def _get_precision(cls, timer):
        if timer in cls._precisions:
            return cls._precisions[timer]
        else:
            return cls._precisions.setdefault(timer, compute_timer_precision(timer))

    def __init__(self, name, disable_gc, timer, min_rounds, min_time, max_time, warmup, add_stats, group=None):
        self._disable_gc = disable_gc
        self._name = name
        self._group = group
        self._timer = timer
        self._min_rounds = min_rounds
        self._max_time = float(max_time)
        self._min_time = float(min_time)
        self._add_stats = add_stats
        self._warmup = warmup

    def __call__(self, function_to_benchmark):
        def runner(loops, timer=self._timer):
            loops_range = XRANGE(loops)
            gcenabled = gc.isenabled()
            if self._disable_gc:
                gc.disable()
            start = timer()
            for index in loops_range:
                function_to_benchmark()
            end = timer()
            if gcenabled:
                gc.enable()
            return end - start

        duration, scale = self._calibrate_timer(runner)

        # Choose how many time we must repeat the test
        rounds = int(math.ceil(self._max_time / duration))
        rounds = max(rounds, self._min_rounds)

        stats = BenchmarkStats(self._name, group=self._group, scale=scale)
        self._add_stats(stats)

        if self._warmup:
            for _ in XRANGE(rounds):
                runner(scale)

        for _ in XRANGE(rounds):
            stats.update(runner(scale))

        return function_to_benchmark()

    def _calibrate_timer(self, runner):
        # Calibrate the number of loops
        if self._min_time < 1.0:
            timer_precision = self._get_precision(self._timer)
            min_time = max(self._min_time, timer_precision * 100)
            min_time_estimate = timer_precision * 10
        else:
            min_time = self._min_time
            min_time_estimate = self._min_time / 100
        # print("")
        # print("Calibrate to %.12f seconds:" % min_time)

        loops = 1
        min_progress = 1.0
        while True:
            duration = runner(loops)
            # print("Calibrate scale: %.24f seconds for %s loops" % (duration, loops))

            if duration / min_time >= min_progress:
                break

            if duration >= min_time_estimate:
                # coarse estimation of the number of loops
                loops = max(int(min_time * loops / duration), loops * 2)
                # print("Calibrate scale: estimate %s loops" % loops)
                min_progress = 0.75
            else:
                loops *= 10
        return duration, loops

class BenchmarkSession(object):
    def __init__(self, config):
        timer = config.getoption("benchmark_timer")
        self._options = dict(
            min_time=config.getoption("benchmark_min_time"),
            min_rounds=config.getoption("benchmark_min_rounds"),
            max_time=config.getoption("benchmark_max_time"),
            timer=timer,
            disable_gc=config.getoption("benchmark_disable_gc"),
            warmup=config.getoption("benchmark_warmup")
        )
        self._skip = config.getoption("benchmark_skip")
        self._only = config.getoption("benchmark_only")
        self._sort = config.getoption("benchmark_sort")
        if self._skip and self._only:
            raise pytest.UsageError("Can't have both --benchmark-only and --benchmark-skip options.")
        self._benchmarks = []

    @pytest.fixture(scope="function")
    def benchmark(self, request):
        if self._skip:
            pytest.skip("Benchmarks are disabled.")
        else:
            node = request.node
            marker = node.get_marker("benchmark")
            options = marker.kwargs if marker else {}
            benchmark = BenchmarkFixture(node.name, add_stats=self._benchmarks.append, **dict(self._options, **options))
            return benchmark

    def pytest_runtest_call(self, item, __multicall__):
        benchmark = hasattr(item, "funcargs") and item.funcargs.get("benchmark")
        if isinstance(benchmark, BenchmarkFixture):
            if self._skip:
                pytest.skip("Skipping benchmark (--benchmark-skip active).")
            else:
                __multicall__.execute()
        else:
            if self._only:
                pytest.skip("Skipping non-benchmark (--benchmark-only active).")
            else:
                __multicall__.execute()

    def pytest_terminal_summary(self, terminalreporter):
        tr = terminalreporter

        if not self._benchmarks:
            return

        timer = self._options.get('timer') or 'default'
        timer_name = timer.__module__ + "." if hasattr(timer, '__module__') else ""
        timer_name += timer.__name__ if hasattr(timer, '__name__') else repr(timer)

        groups = defaultdict(list)
        for bench in self._benchmarks:
            groups[bench.group].append(bench)
        for group, benchmarks in sorted(groups.items(), key=lambda pair: pair[0] or ""):
            worst = {}
            best = {}
            for prop in "min", "max", "mean", "stddev", "runs", "scale":
                worst[prop] = max(benchmark[prop] for benchmark in benchmarks)
            for prop in "min", "max", "mean", "stddev":
                best[prop] = min(benchmark[prop] for benchmark in benchmarks)

            overall_min = best[self._sort]
            if overall_min < 1e-6:
                unit, adjustment = "n", 1e9
            elif overall_min < 1e-3:
                unit, adjustment = "u", 1e6
            elif overall_min < 1:
                unit, adjustment = "m", 1e3
            else:
                unit, adjustment = "", 1.
            labels = {
                "name": "Name (time in %ss)" % unit,
                "min": "Min",
                "max": "Max",
                "mean": "Mean",
                "stddev": "StdDev",
                "runs": "Rounds",
                "scale": "Iterations",
            }
            widths = {
                "name": 3 + max(len(labels["name"]), max(len(benchmark.name) for benchmark in benchmarks)),
                "runs": 2 + max(len(labels["runs"]), len(str(worst["runs"]))),
                "scale": 2 + max(len(labels["scale"]), len(str(worst["scale"]))),
            }
            for prop in "min", "max", "mean", "stddev":
                widths[prop] = 2 + max(len(labels[prop]), max(
                    len("{0:.4f}".format(benchmark[prop] * adjustment))
                    for benchmark in benchmarks
                ))
            tr.write_line(
                " benchmark{2}: {0} tests, min {1.min_rounds} rounds (of min {1.min_time:f}s),"
                " {1.max_time:f}s max time, timer: {3} ".format(
                    len(benchmarks),
                    type("", (), self._options),
                    "" if group is None else " %r" % group,
                    timer_name,
                ).center(sum(widths.values()), '-'),
                yellow=True,
            )
            tr.write_line(labels["name"].ljust(widths["name"]) + "".join(
                labels[prop].rjust(widths[prop])
                for prop in ("min", "max", "mean", "stddev", "runs", "scale")
            ))
            tr.write_line("-" * sum(widths.values()), yellow=True)

            for benchmark in benchmarks:
                tr.write(benchmark.name.ljust(widths["name"]))
                for prop in "min", "max", "mean", "stddev":
                    tr.write(
                        "{0:>{1}.4f}".format(benchmark[prop] * adjustment, widths[prop]),
                        green=benchmark[prop] == best[prop],
                        red=benchmark[prop] == worst[prop],
                        bold=True,
                    )
                for prop in "runs", "scale":
                    tr.write("{0:>{1}}".format(benchmark[prop], widths[prop]))
                tr.write("\n")

            tr.write_line("-" * sum(widths.values()), yellow=True)
            tr.write_line("")

def pytest_runtest_setup(item):
    benchmark = item.get_marker("benchmark")
    if benchmark:
        if benchmark.args:
            raise ValueError("benchmark mark can't have positional arguments.")
        for name in benchmark.kwargs:
            if name not in ("max_time", "min_rounds", "min_time", "timer", "group", "disable_gc", "warmup"):
                raise ValueError("benchmark mark can't have %r keyword argument." % name)

def pytest_configure(config):
    config.addinivalue_line("markers", "benchmark: mark a test with custom benchmark settings.")
    config.pluginmanager.register(BenchmarkSession(config), "pytest-benchmark")
<commit_msg>Add diagnostics support and simplify the calibration code to have smaller rounds.<commit_after>from __future__ import division

import argparse
import gc
import sys
import math
from collections import defaultdict
from decimal import Decimal

import pytest
import py

from .stats import RunningStats
from .timers import compute_timer_precision
from .timers import default_timer
from .compat import XRANGE, PY3

def loadtimer(string):
    if "." not in string:
        raise argparse.ArgumentTypeError("Value for --benchmark-timer must be in dotted form. Eg: 'module.attr'.")
    mod, attr = string.rsplit(".", 1)
    if mod == 'pep418':
        if PY3:
            import time
            return getattr(time, attr)
        else:
            from . import pep418
            return getattr(pep418, attr)
    else:
        __import__(mod)
        mod = sys.modules[mod]
        return getattr(mod, attr)

def loadsort(string):
    if string not in ("min", "max", "mean", "stddev"):
        raise argparse.ArgumentTypeError("Value for --benchmark-sort must be one of: 'min', 'max', 'mean' or 'stddev'.")
    return string

def loadrounds(string):
    try:
        value = int(string)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(exc)
    else:
        if value < 1:
            raise argparse.ArgumentTypeError("Value for --benchmark-rounds must be at least 1.")
        return value

def pytest_addoption(parser):
    group = parser.getgroup("benchmark")
    group.addoption(
        "--benchmark-min-time",
        action="store", type=Decimal, default=Decimal("0.1"),
        help="Minimum time per round."
    )
    group.addoption(
        "--benchmark-max-time",
        action="store", type=Decimal, default=Decimal("1.0"),
        help="Maximum time to spend in a benchmark."
    )
    group.addoption(
        "--benchmark-min-rounds",
        action="store", type=loadrounds, default=5,
        help="Minimum rounds, even if total time would exceed `--max-time`."
    )
    group.addoption(
        "--benchmark-sort",
        action="store", type=loadsort, default="min",
        help="Column to sort on. Can be one of: 'min', 'max', 'mean' or 'stddev' ."
    )
    group.addoption(
        "--benchmark-timer",
        action="store", type=loadtimer, default=default_timer,
        help="Timer to use when measuring time."
    )
    group.addoption(
        "--benchmark-warmup",
        action="store_true", default=False,
        help="Runs the benchmarks two times. Discards data from the first run."
    )
    group.addoption(
        "--benchmark-verbose",
        action="store_true", default=False,
        help="Dump diagnostic and progress information."
    )
    group.addoption(
        "--benchmark-disable-gc",
        action="store_true", default=False,
        help="Disable GC during benchmarks."
    )
    group.addoption(
        "--benchmark-skip",
        action="store_true", default=False,
        help="Skip running any benchmarks."
    )
    group.addoption(
        "--benchmark-only",
        action="store_true", default=False,
        help="Only run benchmarks."
    )

class BenchmarkStats(RunningStats):
    def __init__(self, name, group, scale):
        self.name = name
        self.group = group
        self.scale = scale
        super(BenchmarkStats, self).__init__()

    def __getitem__(self, key):
        return getattr(self, key)

    def update(self, duration):
        super(BenchmarkStats, self).update(duration / self.scale)

class BenchmarkFixture(object):
    _precisions = {}

    @classmethod
    def _get_precision(cls, timer):
        if timer in cls._precisions:
            return cls._precisions[timer]
        else:
            return cls._precisions.setdefault(timer, compute_timer_precision(timer))

    def __init__(self, name, disable_gc, timer, min_rounds, min_time, max_time, warmup, add_stats, logger, group=None):
        self._disable_gc = disable_gc
        self._name = name
        self._group = group
        self._timer = timer
        self._min_rounds = min_rounds
        self._max_time = float(max_time)
        self._min_time = float(min_time)
        self._add_stats = add_stats
        self._warmup = warmup
        self._logger = logger

    def __call__(self, function_to_benchmark):
        def runner(loops, timer=self._timer):
            loops_range = XRANGE(loops)
            gcenabled = gc.isenabled()
            if self._disable_gc:
                gc.disable()
            start = timer()
            for index in loops_range:
                function_to_benchmark()
            end = timer()
            if gcenabled:
                gc.enable()
            return end - start

        duration, scale = self._calibrate_timer(runner)

        # Choose how many time we must repeat the test
        rounds = int(math.ceil(self._max_time / duration))
        rounds = max(rounds, self._min_rounds)

        stats = BenchmarkStats(self._name, group=self._group, scale=scale)
        self._add_stats(stats)

        if self._warmup:
            for _ in XRANGE(rounds):
                runner(scale)

        for _ in XRANGE(rounds):
            stats.update(runner(scale))

        return function_to_benchmark()

    def _calibrate_timer(self, runner):
        timer_precision = self._get_precision(self._timer)
        min_time = max(self._min_time, timer_precision * 100)
        min_time_estimate = min_time / 10
        self._logger.write("")
        self._logger.write("  Timer precision: %ss" % time_format(timer_precision))
        self._logger.write("  Calibrating to target round %ss; will estimate when reaching %ss." % (
            time_format(min_time), time_format(min_time_estimate)))

        loops = 1
        while True:
            duration = runner(loops)
            self._logger.write("  Calibrate: %ss for %s iterations." % (time_format(duration), loops))

            if duration / min_time >= 0.75:
                break

            if duration >= min_time_estimate:
                # coarse estimation of the number of loops
                loops = int(min_time * loops / duration)
                self._logger.write("  Calibrate estimate: %s iterations." % loops)
            else:
                loops *= 10
        return duration, loops

def time_unit(value):
    if value < 1e-6:
        return "n", 1e9
    elif value < 1e-3:
        return "u", 1e6
    elif value < 1:
        return "m", 1e3
    else:
        return "", 1.

def time_format(value):
    unit, adjustment = time_unit(value)
    return "{0:.2f}{1:s}".format(value * adjustment, unit)

class DiagnosticLogger(object):
    def __init__(self, verbose
                 # , capman
                 ):
        self.term = verbose and py.io.TerminalWriter()
        # self.capman = capman

    def write(self, text):
        if self.term:
            # if self.capman:
            #     self.capman.suspendcapture(in_=True)
            self.term.line(text, yellow=True)
            # if self.capman:
            #     self.capman.resumecapture()

class BenchmarkSession(object):
    def __init__(self, config):
        timer = config.getoption("benchmark_timer")
        self._options = dict(
            min_time=config.getoption("benchmark_min_time"),
            min_rounds=config.getoption("benchmark_min_rounds"),
            max_time=config.getoption("benchmark_max_time"),
            timer=timer,
            disable_gc=config.getoption("benchmark_disable_gc"),
            warmup=config.getoption("benchmark_warmup")
        )
        self._skip = config.getoption("benchmark_skip")
        self._only = config.getoption("benchmark_only")
        self._sort = config.getoption("benchmark_sort")
        self._verbose = config.getoption("benchmark_verbose")
        if self._skip and self._only:
            raise pytest.UsageError("Can't have both --benchmark-only and --benchmark-skip options.")
        self._benchmarks = []

    @pytest.fixture(scope="function")
    def benchmark(self, request):
        if self._skip:
            pytest.skip("Benchmarks are disabled.")
        else:
            node = request.node
            marker = node.get_marker("benchmark")
            options = marker.kwargs if marker else {}
            benchmark = BenchmarkFixture(
                node.name,
                add_stats=self._benchmarks.append,
                logger=DiagnosticLogger(
                    self._verbose,
                    # node.config.pluginmanager.getplugin("capturemanager")
                ),
                **dict(self._options, **options)
            )
            return benchmark

    def pytest_runtest_call(self, item, __multicall__):
        benchmark = hasattr(item, "funcargs") and item.funcargs.get("benchmark")
        if isinstance(benchmark, BenchmarkFixture):
            if self._skip:
                pytest.skip("Skipping benchmark (--benchmark-skip active).")
            else:
                __multicall__.execute()
        else:
            if self._only:
                pytest.skip("Skipping non-benchmark (--benchmark-only active).")
            else:
                __multicall__.execute()

    def pytest_terminal_summary(self, terminalreporter):
        tr = terminalreporter

        if not self._benchmarks:
            return

        timer = self._options.get('timer') or 'default'
        timer_name = timer.__module__ + "." if hasattr(timer, '__module__') else ""
        timer_name += timer.__name__ if hasattr(timer, '__name__') else repr(timer)

        groups = defaultdict(list)
        for bench in self._benchmarks:
            groups[bench.group].append(bench)
        for group, benchmarks in sorted(groups.items(), key=lambda pair: pair[0] or ""):
            worst = {}
            best = {}
            for prop in "min", "max", "mean", "stddev", "runs", "scale":
                worst[prop] = max(benchmark[prop] for benchmark in benchmarks)
            for prop in "min", "max", "mean", "stddev":
                best[prop] = min(benchmark[prop] for benchmark in benchmarks)

            unit, adjustment = time_unit(best[self._sort])
            labels = {
                "name": "Name (time in %ss)" % unit,
                "min": "Min",
                "max": "Max",
                "mean": "Mean",
                "stddev": "StdDev",
                "runs": "Rounds",
                "scale": "Iterations",
            }
            widths = {
                "name": 3 + max(len(labels["name"]), max(len(benchmark.name) for benchmark in benchmarks)),
                "runs": 2 + max(len(labels["runs"]), len(str(worst["runs"]))),
                "scale": 2 + max(len(labels["scale"]), len(str(worst["scale"]))),
            }
            for prop in "min", "max", "mean", "stddev":
                widths[prop] = 2 + max(len(labels[prop]), max(
                    len("{0:.4f}".format(benchmark[prop] * adjustment))
                    for benchmark in benchmarks
                ))
            tr.write_line(
                " benchmark{2}: {0} tests, min {1.min_rounds} rounds (of min {1.min_time:f}s),"
                " {1.max_time:f}s max time, timer: {3} ".format(
                    len(benchmarks),
                    type("", (), self._options),
                    "" if group is None else " %r" % group,
                    timer_name,
                ).center(sum(widths.values()), '-'),
                yellow=True,
            )
            tr.write_line(labels["name"].ljust(widths["name"]) + "".join(
                labels[prop].rjust(widths[prop])
                for prop in ("min", "max", "mean", "stddev", "runs", "scale")
            ))
            tr.write_line("-" * sum(widths.values()), yellow=True)

            for benchmark in benchmarks:
                tr.write(benchmark.name.ljust(widths["name"]))
                for prop in "min", "max", "mean", "stddev":
                    tr.write(
                        "{0:>{1}.4f}".format(benchmark[prop] * adjustment, widths[prop]),
                        green=benchmark[prop] == best[prop],
                        red=benchmark[prop] == worst[prop],
                        bold=True,
                    )
                for prop in "runs", "scale":
                    tr.write("{0:>{1}}".format(benchmark[prop], widths[prop]))
                tr.write("\n")

            tr.write_line("-" * sum(widths.values()), yellow=True)
            tr.write_line("")

def pytest_runtest_setup(item):
    benchmark = item.get_marker("benchmark")
    if benchmark:
        if benchmark.args:
            raise ValueError("benchmark mark can't have positional arguments.")
        for name in benchmark.kwargs:
            if name not in ("max_time", "min_rounds", "min_time", "timer", "group", "disable_gc", "warmup"):
                raise ValueError("benchmark mark can't have %r keyword argument." % name)

def pytest_configure(config):
    config.addinivalue_line("markers", "benchmark: mark a test with custom benchmark settings.")
    config.pluginmanager.register(BenchmarkSession(config), "pytest-benchmark")
<|endoftext|>


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

import os
from collections import OrderedDict

from ansible.module_utils.basic import AnsibleModule

import yaml

DOCUMENTATION = '''
---
module: load_input_model
short_description: Load an Ardana input model
description: |
  Load an Ardana input model from a given path
author: SUSE Linux GmbH
options:
  path:
    description: Root path where the input model files are located
'''

EXAMPLES = '''
- load_input_model:
    path: path/to/input/model
  register: _result
- debug: msg="{{ _result.input_model }}"
'''

def merge_input_model(data, input_model):
    for key, value in data.iteritems():
        if key in input_model and isinstance(input_model[key], list) and value:
            input_model[key] += value
        else:
            input_model[key] = value

def load_input_model_file(file_name, input_model):
    if file_name.endswith('.yml') or file_name.endswith('.yaml'):
        with open(file_name, 'r') as data_file:
            data = yaml.load(data_file.read())
            merge_input_model(
                data,
                input_model)
    return input_model

def load_input_model(input_model_path):
    input_model = OrderedDict()

    if os.path.exists(input_model_path):
        if os.path.isdir(input_model_path):
            for root, dirs, files in os.walk(input_model_path):
                for f in files:
                    file_name = os.path.join(root, f)
                    input_model = load_input_model_file(
                        file_name,
                        input_model)
        else:
            input_model = load_input_model_file(
                input_model_path,
                input_model)

    return input_model

def main():

    argument_spec = dict(
        path=dict(type='str', required=True)
    )
    module = AnsibleModule(argument_spec=argument_spec,
                           supports_check_mode=False)
    input_model_path = module.params['path']

    try:
        input_model = load_input_model(input_model_path)
    except Exception as e:
        module.fail_json(msg=e.message)
    module.exit_json(rc=0, changed=False, input_model=input_model)

if __name__ == '__main__':
    main()
<commit_msg>Use `items()` instead of `iteritems()`<commit_after>#!/usr/bin/python
#
# (c) Copyright 2018 SUSE LLC
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
#

import os
from collections import OrderedDict

from ansible.module_utils.basic import AnsibleModule

import yaml

DOCUMENTATION = '''
---
module: load_input_model
short_description: Load an Ardana input model
description: |
  Load an Ardana input model from a given path
author: SUSE Linux GmbH
options:
  path:
    description: Root path where the input model files are located
'''

EXAMPLES = '''
- load_input_model:
    path: path/to/input/model
  register: _result
- debug: msg="{{ _result.input_model }}"
'''

def merge_input_model(data, input_model):
    for key, value in data.items():
        if key in input_model and isinstance(input_model[key], list) and value:
            input_model[key] += value
        else:
            input_model[key] = value

def load_input_model_file(file_name, input_model):
    if file_name.endswith('.yml') or file_name.endswith('.yaml'):
        with open(file_name, 'r') as data_file:
            data = yaml.load(data_file.read())
            merge_input_model(
                data,
                input_model)
    return input_model

def load_input_model(input_model_path):
    input_model = OrderedDict()

    if os.path.exists(input_model_path):
        if os.path.isdir(input_model_path):
            for root, dirs, files in os.walk(input_model_path):
                for f in files:
                    file_name = os.path.join(root, f)
                    input_model = load_input_model_file(
                        file_name,
                        input_model)
        else:
            input_model = load_input_model_file(
                input_model_path,
                input_model)

    return input_model

def main():

    argument_spec = dict(
        path=dict(type='str', required=True)
    )
    module = AnsibleModule(argument_spec=argument_spec,
                           supports_check_mode=False)
    input_model_path = module.params['path']

    try:
        input_model = load_input_model(input_model_path)
    except Exception as e:
        module.fail_json(msg=e.message)
    module.exit_json(rc=0, changed=False, input_model=input_model)

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


<commit_before>
import sys
import types
import inspect

def isstring(s):
    # if we use Python 3
    if (sys.version_info[0] >= 3):
        return isinstance(s, str)
    # we use Python 2
    return isinstance(s, basestring)

def normalize_func(func):
    # convert class to __init__ method if we can
    if inspect.isclass(func):
      if (inspect.ismethod(func.__init__)):
        return func.__init__
      else:
        return None

    # convert func to __call__
    if not inspect.isfunction(func) and hasattr(func, '__call__'):
      return func.__call__

    # return None for builtins
    if (inspect.isbuiltin(func)):
        return None

    return func

def get_doc(func):
  # if this func is actually a class, grab the doc from its __init__ method
  if isinstance(func, type):
    func = func.__init__
  doc = inspect.getdoc(func)
  if doc is None:
    func = normalize_func(func)
    if func is None:
      return None
    else:
      return inspect.getdoc(func)
  else:
    return doc

def get_property_doc(target, prop):
  for name, obj in inspect.getmembers(type(target), inspect.isdatadescriptor):
    if (isinstance(obj, property) and name == prop):
      return inspect.getdoc(obj.fget)
  return None

def get_argspec(func):
  try:
    if sys.version_info[0] >= 3:
      argspec = inspect.getfullargspec(func)
    else:
      argspec = inspect.getargspec(func)
  except TypeError:
    return None

def get_arguments(func):
    func = normalize_func(func)
    if func is None:
      return None
    argspec = get_argspec(func)
    if argspec is None:
      return None
    args = argspec.args
    if 'self' in args:
      args.remove('self')
    return args

def get_r_representation(default):
  if callable(default) and hasattr(default, '__name__'):
    arg_value = default.__name__
  else:
    if default is None:
      arg_value = "NULL"
    elif type(default) == type(True):
      if default == True:
        arg_value = "TRUE"
      else:
        arg_value = "FALSE"
    elif isstring(default):
      arg_value = "\"%s\"" % default
    elif isinstance(default, int):
      arg_value = "%rL" % default
    elif isinstance(default, float):
      arg_value = "%r" % default
    elif isinstance(default, list):
      arg_value = "c("
      for i, item in enumerate(default):
        if i is (len(default) - 1):
          arg_value += "%s)" % get_r_representation(item)
        else:
          arg_value += "%s, " % get_r_representation(item)
    elif isinstance(default, (tuple, set)):
      arg_value = "list("
      for i, item in enumerate(default):
        if i is (len(default) - 1):
          arg_value += "%s)" % get_r_representation(item)
        else:
          arg_value += "%s, " % get_r_representation(item)
    elif isinstance(default, dict):
      arg_value = "list("
      for i in range(len(default)):
        i_arg_value = "%s = %s" % \
          (default.keys()[i], get_r_representation(default.values()[i]))
        if i is (len(default) - 1):
          arg_value += "%s)" % i_arg_value
        else:
          arg_value += "%s, " % i_arg_value
    else:
      arg_value = "%r" % default

  # if the value starts with "tf." then convert to $ usage
  if (arg_value.startswith("tf.")):
    arg_value = arg_value.replace(".", "$")

  return(arg_value)

def generate_signature_for_function(func):
    """Given a function, returns a string representing its args."""

    func = normalize_func(func)
    if func is None:
      return None

    args_list = []

    argspec = get_argspec(func)
    if argspec is None:
      return None

    first_arg_with_default = (
        len(argspec.args or []) - len(argspec.defaults or []))
    for arg in argspec.args[:first_arg_with_default]:
      if arg == "self":
        # Python documentation typically skips `self` when printing method
        # signatures.
        continue
      args_list.append(arg)

    # TODO(mrry): This is a workaround for documenting signature of
    # functions that have the @contextlib.contextmanager decorator.
    # We should do something better.
    if argspec.varargs == "args" and argspec.keywords == "kwds":
      original_func = func.__closure__[0].cell_contents
      return generate_signature_for_function(original_func)

    if argspec.defaults:
      for arg, default in zip(
          argspec.args[first_arg_with_default:], argspec.defaults):
        arg_value = get_r_representation(default)
        args_list.append("%s = %s" % (arg, arg_value))
    if argspec.varargs:
      args_list.append("...")
    if argspec.keywords:
      args_list.append("...")
    return "(" + ", ".join(args_list) + ")"

<commit_msg>return from getargspec function<commit_after>
import sys
import types
import inspect

def isstring(s):
    # if we use Python 3
    if (sys.version_info[0] >= 3):
        return isinstance(s, str)
    # we use Python 2
    return isinstance(s, basestring)

def normalize_func(func):
    # convert class to __init__ method if we can
    if inspect.isclass(func):
      if (inspect.ismethod(func.__init__)):
        return func.__init__
      else:
        return None

    # convert func to __call__
    if not inspect.isfunction(func) and hasattr(func, '__call__'):
      return func.__call__

    # return None for builtins
    if (inspect.isbuiltin(func)):
        return None

    return func

def get_doc(func):
  # if this func is actually a class, grab the doc from its __init__ method
  if isinstance(func, type):
    func = func.__init__
  doc = inspect.getdoc(func)
  if doc is None:
    func = normalize_func(func)
    if func is None:
      return None
    else:
      return inspect.getdoc(func)
  else:
    return doc

def get_property_doc(target, prop):
  for name, obj in inspect.getmembers(type(target), inspect.isdatadescriptor):
    if (isinstance(obj, property) and name == prop):
      return inspect.getdoc(obj.fget)
  return None

def get_argspec(func):
  try:
    if sys.version_info[0] >= 3:
      return inspect.getfullargspec(func)
    else:
      return inspect.getargspec(func)
  except TypeError:
    return None

def get_arguments(func):
    func = normalize_func(func)
    if func is None:
      return None
    argspec = get_argspec(func)
    if argspec is None:
      return None
    args = argspec.args
    if 'self' in args:
      args.remove('self')
    return args

def get_r_representation(default):
  if callable(default) and hasattr(default, '__name__'):
    arg_value = default.__name__
  else:
    if default is None:
      arg_value = "NULL"
    elif type(default) == type(True):
      if default == True:
        arg_value = "TRUE"
      else:
        arg_value = "FALSE"
    elif isstring(default):
      arg_value = "\"%s\"" % default
    elif isinstance(default, int):
      arg_value = "%rL" % default
    elif isinstance(default, float):
      arg_value = "%r" % default
    elif isinstance(default, list):
      arg_value = "c("
      for i, item in enumerate(default):
        if i is (len(default) - 1):
          arg_value += "%s)" % get_r_representation(item)
        else:
          arg_value += "%s, " % get_r_representation(item)
    elif isinstance(default, (tuple, set)):
      arg_value = "list("
      for i, item in enumerate(default):
        if i is (len(default) - 1):
          arg_value += "%s)" % get_r_representation(item)
        else:
          arg_value += "%s, " % get_r_representation(item)
    elif isinstance(default, dict):
      arg_value = "list("
      for i in range(len(default)):
        i_arg_value = "%s = %s" % \
          (default.keys()[i], get_r_representation(default.values()[i]))
        if i is (len(default) - 1):
          arg_value += "%s)" % i_arg_value
        else:
          arg_value += "%s, " % i_arg_value
    else:
      arg_value = "%r" % default

  # if the value starts with "tf." then convert to $ usage
  if (arg_value.startswith("tf.")):
    arg_value = arg_value.replace(".", "$")

  return(arg_value)

def generate_signature_for_function(func):
    """Given a function, returns a string representing its args."""

    func = normalize_func(func)
    if func is None:
      return None

    args_list = []

    argspec = get_argspec(func)
    if argspec is None:
      return None

    first_arg_with_default = (
        len(argspec.args or []) - len(argspec.defaults or []))
    for arg in argspec.args[:first_arg_with_default]:
      if arg == "self":
        # Python documentation typically skips `self` when printing method
        # signatures.
        continue
      args_list.append(arg)

    # TODO(mrry): This is a workaround for documenting signature of
    # functions that have the @contextlib.contextmanager decorator.
    # We should do something better.
    if argspec.varargs == "args" and argspec.keywords == "kwds":
      original_func = func.__closure__[0].cell_contents
      return generate_signature_for_function(original_func)

    if argspec.defaults:
      for arg, default in zip(
          argspec.args[first_arg_with_default:], argspec.defaults):
        arg_value = get_r_representation(default)
        args_list.append("%s = %s" % (arg, arg_value))
    if argspec.varargs:
      args_list.append("...")
    if argspec.keywords:
      args_list.append("...")
    return "(" + ", ".join(args_list) + ")"

<|endoftext|>


<commit_before>import ddapp.applogic as app

from ddapp import lcmUtils
from ddapp import transformUtils
from ddapp import visualization as vis
from ddapp.shallowCopy import shallowCopy
from ddapp.timercallback import TimerCallback
from ddapp import vtkNumpy
import ddapp.vtkAll as vtk

import PythonQt
from PythonQt import QtCore, QtGui
import bot_core as lcmbotcore
import numpy as np
from ddapp.simpletimer import SimpleTimer
from ddapp import ioUtils
import sys
import drc as lcmdrc

def clipRange(dataObj, arrayName, thresholdRange):
    if not dataObj.GetPointData().GetArray(arrayName):
        raise Exception('clipRange: could not locate array: %s' % arrayName)

    dataObj.GetPointData().SetScalars(dataObj.GetPointData().GetArray(arrayName))

    f = vtk.vtkClipPolyData()
    f.SetInput(dataObj)
    f.SetValue(thresholdRange[0])
    f.SetInputArrayToProcess(0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, arrayName)

    f2 = vtk.vtkClipPolyData()
    f2.AddInputConnection(f.GetOutputPort())
    f2.SetValue(thresholdRange[1])
    f2.InsideOutOn()
    f2.SetInputArrayToProcess(0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, arrayName)

    f2.Update()
    return shallowCopy(f2.GetOutput())

def makeSphere(radius, resolution):

    s = vtk.vtkSphereSource()
    s.SetThetaResolution(resolution)
    s.SetPhiResolution(resolution)
    s.SetRadius(radius)
    s.SetEndPhi(85)
    s.Update()
    return shallowCopy(s.GetOutput())

def colorizePoints(polyData, cameraName='CAMERA_LEFT'):
    cameraView.queue.colorizePoints(cameraName, polyData)

def sendFOVRequest(channel, imagePoints):

    channelToImageType = {
        'CAMERA_LEFT' : lcmdrc.data_request_t.CAMERA_IMAGE_HEAD_LEFT,
        'CAMERACHEST_LEFT' : lcmdrc.data_request_t.CAMERA_IMAGE_LCHEST,
        'CAMERACHEST_RIGHT' : lcmdrc.data_request_t.CAMERA_IMAGE_RCHEST,
                         }

    dataRequest = lcmdrc.data_request_t()
    dataRequest.type = channelToImageType[channel]

    message = lcmdrc.subimage_request_t()
    message.data_request = dataRequest

    imagePoints = np.array([[pt[0], pt[1]] for pt in imagePoints])
    minX, maxX = imagePoints[:,0].min(), imagePoints[:,0].max()
    minY, maxY = imagePoints[:,1].min(), imagePoints[:,1].max()

    message.x = minX
    message.y = minY
    message.w = maxX - minX
    message.h = maxY - minY

    print message.x, message.y, message.w, message.h

    requestChannel = 'SUBIMAGE_REQUEST'
    lcmUtils.publish(requestChannel, message)

def testColorize():

    radius = 10
    resolution = 400
    s = makeSphere(radius, resolution)
    cameraView.queue.colorizePoints(s)
    showPolyData(p, 'sphere', colorByName='rgb')

def rayDebug(position, ray):
    d = vis.DebugData()
    d.addLine(position, position+ray*5.0)
    drcView = app.getViewManager().findView('DRC View')
    vis.updatePolyData(d.getPolyData(), 'camera ray', view=drcView)

class CameraView(object):

    def __init__(self):

        self.initImageQueue()

        if app.getMainWindow():
            self.initView()
            self.initEventFilter()
            self.rayCallback = rayDebug
        else:
            self.view = None

        self.timerCallback = TimerCallback()
        self.timerCallback.targetFps = 60
        self.timerCallback.callback = self.updateView
        self.timerCallback.start()

    def initImageQueue(self):

        imageNames = [
            'CAMERA_LEFT',
            'CAMERACHEST_LEFT',
            'CAMERACHEST_RIGHT',
            ]

        self.images = {}
        self.transforms = {}
        self.imageUtimes = {}
        self.textures = {}

        for name in imageNames:
            self.addImage(name)

        self.queue = PythonQt.dd.ddBotImageQueue(lcmUtils.getGlobalLCMThread())
        self.queue.init(lcmUtils.getGlobalLCMThread())

    def addImage(self, name):
        image = vtk.vtkImageData()
        tex = vtk.vtkTexture()
        tex.SetInput(image)
        tex.EdgeClampOn()
        tex.RepeatOff()

        self.imageUtimes[name] = 0
        self.images[name] = image
        self.textures[name] = tex
        self.transforms[name] = vtk.vtkTransform()

    def onViewDoubleClicked(self, displayPoint):

        obj, pickedPoint = vis.findPickedObject(displayPoint, self.view)

        if pickedPoint is None or not obj:
            return

        imageName = obj.getProperty('Name')
        imageUtime = self.imageUtimes[imageName]

        cameraToLocal = vtk.vtkTransform()
        self.queue.getTransform(imageName, 'local', imageUtime, cameraToLocal)

        utorsoToLocal = vtk.vtkTransform()
        self.queue.getTransform('utorso', 'local', imageUtime, utorsoToLocal)

        p = range(3)
        utorsoToLocal.TransformPoint(pickedPoint, p)

        ray = np.array(p) - np.array(cameraToLocal.GetPosition())
        ray /= np.linalg.norm(ray)

        if self.rayCallback:
            self.rayCallback(np.array(cameraToLocal.GetPosition()), ray)

    def filterEvent(self, obj, event):
        if event.type() == QtCore.QEvent.MouseButtonDblClick:
            self.eventFilter.setEventHandlerResult(True)
            self.onViewDoubleClicked(vis.mapMousePosition(obj, event))

    def initEventFilter(self):
        self.eventFilter = PythonQt.dd.ddPythonEventFilter()
        qvtkwidget = self.view.vtkWidget()
        qvtkwidget.installEventFilter(self.eventFilter)
        self.eventFilter.addFilteredEventType(QtCore.QEvent.MouseButtonDblClick)
        self.eventFilter.connect('handleEvent(QObject*, QEvent*)', self.filterEvent)

    def initView(self):

        self.view = app.getViewManager().createView('Camera View', 'VTK View')

        self.view.camera().SetViewAngle(90)
        self.view.camera().SetPosition(-7.5, 0.0, 5.0)
        self.view.camera().SetFocalPoint(0.0, 0.0, 0.0)
        self.view.camera().SetViewUp(0.0, 0.0, 1.0)

        self.sphereObjects = {}

        app.toggleCameraTerrainMode(self.view)

    def getSphereGeometry(self, imageName):

        sphereObj = self.sphereObjects.get(imageName)
        if sphereObj:
            return sphereObj

        if not self.images[imageName].GetDimensions()[0]:
            return None

        sphereResolution = 50
        sphereRadii = {
                'CAMERA_LEFT' : 9.85,
                'CAMERACHEST_LEFT' : 9.95,
                'CAMERACHEST_RIGHT' : 10
                }

        geometry = makeSphere(sphereRadii[imageName], sphereResolution)
        self.queue.computeTextureCoords(imageName, geometry)

        tcoordsArrayName = 'tcoords_%s' % imageName
        vtkNumpy.addNumpyToVtk(geometry, vtkNumpy.getNumpyFromVtk(geometry, tcoordsArrayName)[:,0].copy(), 'tcoords_U')
        vtkNumpy.addNumpyToVtk(geometry, vtkNumpy.getNumpyFromVtk(geometry, tcoordsArrayName)[:,1].copy(), 'tcoords_V')
        geometry = clipRange(geometry, 'tcoords_U', [0.0, 1.0])
        geometry = clipRange(geometry, 'tcoords_V', [0.0, 1.0])
        geometry.GetPointData().SetTCoords(geometry.GetPointData().GetArray(tcoordsArrayName))

        sphereObj = vis.showPolyData(geometry, imageName, view=self.view, parent='cameras')
        sphereObj.actor.SetTexture(self.textures[imageName])
        sphereObj.actor.GetProperty().LightingOff()

        self.sphereObjects[imageName] = sphereObj
        return sphereObj

    def updateSphereGeometry(self):

        for imageName, transform in self.transforms.iteritems():
            sphereObj = self.getSphereGeometry(imageName)
            if not sphereObj:
                continue

            self.queue.getBodyToCameraTransform(imageName, transform)
            sphereObj.actor.SetUserTransform(transform.GetLinearInverse())

        return True

    def writeImage(self, imageName, outFile):
        writer = vtk.vtkPNGWriter()
        writer.SetInput(self.images[imageName])
        writer.SetFileName(outFile)
        writer.Write()

    def updateImages(self):
        updated = False
        for imageName, image in self.images.iteritems():
            imageUtime = self.queue.getCurrentImageTime(imageName)
            if imageUtime != self.imageUtimes.get(imageName, 0):
                self.imageUtimes[imageName] = self.queue.getImage(imageName, image)
                updated = True
        return updated

    def updateView(self):

        if not self.updateImages():
            return

        if not self.view:
            return

        self.updateSphereGeometry()

        if not self.view.isVisible():
            return
        self.view.render()

class CameraImageView(object):

    def __init__(self, imageManager, imageName, viewName=None, view=None):

        self.imageManager = imageManager
        self.viewName = viewName or imageName
        self.imageName = imageName
        self.imageInitialized = False
        self.updateUtime = 0
        self.initView(view)
        self.initEventFilter()

    def onViewDoubleClicked(self, displayPoint):

        obj, pickedPoint = vis.findPickedObject(displayPoint, self.view)

        if pickedPoint is None or not obj:
            return

        imageName = obj.getProperty('Name')
        imageUtime = self.imageUtimes[imageName]

        cameraToLocal = vtk.vtkTransform()
        self.queue.getTransform(imageName, 'local', imageUtime, cameraToLocal)

        utorsoToLocal = vtk.vtkTransform()
        self.queue.getTransform('utorso', 'local', imageUtime, utorsoToLocal)

        drcView = app.getViewManager().findView('DRC View')

        p = range(3)
        utorsoToLocal.TransformPoint(pickedPoint, p)
        d = vis.DebugData()
        d.addLine(cameraToLocal.GetPosition(), p)
        vis.updatePolyData(d.getPolyData(), 'camera ray', view=drcView)

    def filterEvent(self, obj, event):
        if self.eventFilterEnabled and event.type() == QtCore.QEvent.MouseButtonDblClick:
            self.eventFilter.setEventHandlerResult(True)
            self.onViewDoubleClicked(vis.mapMousePosition(obj, event))

    def onRubberBandPick(self, obj, event):
        displayPoints = self.interactorStyle.GetStartPosition(), self.interactorStyle.GetEndPosition()
        imagePoints = [vis.pickImage(point, self.view)[1] for point in displayPoints]
        sendFOVRequest(self.imageName, imagePoints)

    def initView(self, view):
        self.view = view or app.getViewManager().createView(self.viewName, 'VTK View')
        self.view.installImageInteractor()
        self.interactorStyle = self.view.renderWindow().GetInteractor().GetInteractorStyle()
        self.interactorStyle.AddObserver('SelectionChangedEvent', self.onRubberBandPick)

        self.view.renderWindow().GetInteractor().AddObserver('KeyPressEvent', self.onKeyPress)

        self.imageActor = vtk.vtkImageActor()
        self.imageActor.SetInput(self.imageManager.images[self.imageName])
        self.imageActor.SetVisibility(False)
        self.view.renderer().AddActor(self.imageActor)

        self.timerCallback = TimerCallback()
        self.timerCallback.targetFps = 60
        self.timerCallback.callback = self.updateView
        self.timerCallback.start()

    def initEventFilter(self):
        self.eventFilter = PythonQt.dd.ddPythonEventFilter()
        qvtkwidget = self.view.vtkWidget()
        qvtkwidget.installEventFilter(self.eventFilter)
        self.eventFilter.addFilteredEventType(QtCore.QEvent.MouseButtonDblClick)
        self.eventFilter.connect('handleEvent(QObject*, QEvent*)', self.filterEvent)
        self.eventFilterEnabled = True

    def resetCamera(self):
        camera = self.view.camera()
        camera.ParallelProjectionOn()
        camera.SetFocalPoint(0,0,0)
        camera.SetPosition(0,0,-1)
        camera.SetViewUp(0,-1, 0)
        self.view.resetCamera()
        self.fitImageToView()

    def onKeyPress(self, obj, event):
        if obj.GetKeyCode() == 'p':
            return
        if obj.GetKeyCode() == 'r':
            self.fitImageToView()

    def fitImageToView(self):

        camera = self.view.camera()
        image = self.imageManager.images[self.imageName]
        imageWidth, imageHeight, _ = image.GetDimensions()

        viewWidth, viewHeight = self.view.renderWindow().GetSize()
        aspectRatio = float(viewWidth)/viewHeight
        parallelScale = max(imageWidth/aspectRatio, imageHeight) / 2.0
        camera.SetParallelScale(parallelScale)

    def setImageName(self, imageName):
        if imageName == self.imageName:
            return

        assert imageName in self.imageManager.images

        self.imageName = imageName
        self.imageInitialized = False
        self.updateUtime = 0
        self.imageActor.SetInput(self.imageManager.images[self.imageName])
        self.imageActor.SetVisibility(False)
        self.view.render()

    def updateView(self):

        if not self.view.isVisible():
            return

        currentUtime = self.imageManager.imageUtimes[self.imageName]
        if currentUtime != self.updateUtime:
            self.view.render()
            self.updateUtime = currentUtime

            if not self.imageInitialized and self.imageActor.GetInput().GetDimensions()[0]:
                self.imageActor.SetVisibility(True)
                self.resetCamera()
                self.imageInitialized = True

views = {}

def addCameraView(channel, name):

    cameraView.queue.addCameraStream(channel)
    cameraView.addImage(channel)
    view = CameraImageView(cameraView, channel, name)
    global views
    views[channel] = view
    return view

def init():
    global cameraView
    cameraView = CameraView()

    global headView, chestLeft, chestRight
    headView = CameraImageView(cameraView, 'CAMERA_LEFT', 'Head camera')
    chestLeft = CameraImageView(cameraView, 'CAMERACHEST_LEFT', 'Chest left')
    chestRight = CameraImageView(cameraView, 'CAMERACHEST_RIGHT', 'Chest right')

    #addCameraView('AFFORDANCE_OVERLAY', 'Affordance overlay')

    useAnaglyph = False
    if useAnaglyph:
        addCameraView('CAMERA_ANAGLYPH', 'Head Anaglyph')

<commit_msg>fix key press event filter for camera views<commit_after>import ddapp.applogic as app

from ddapp import lcmUtils
from ddapp import transformUtils
from ddapp import visualization as vis
from ddapp.shallowCopy import shallowCopy
from ddapp.timercallback import TimerCallback
from ddapp import vtkNumpy
import ddapp.vtkAll as vtk

import PythonQt
from PythonQt import QtCore, QtGui
import bot_core as lcmbotcore
import numpy as np
from ddapp.simpletimer import SimpleTimer
from ddapp import ioUtils
import sys
import drc as lcmdrc

def clipRange(dataObj, arrayName, thresholdRange):
    if not dataObj.GetPointData().GetArray(arrayName):
        raise Exception('clipRange: could not locate array: %s' % arrayName)

    dataObj.GetPointData().SetScalars(dataObj.GetPointData().GetArray(arrayName))

    f = vtk.vtkClipPolyData()
    f.SetInput(dataObj)
    f.SetValue(thresholdRange[0])
    f.SetInputArrayToProcess(0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, arrayName)

    f2 = vtk.vtkClipPolyData()
    f2.AddInputConnection(f.GetOutputPort())
    f2.SetValue(thresholdRange[1])
    f2.InsideOutOn()
    f2.SetInputArrayToProcess(0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, arrayName)

    f2.Update()
    return shallowCopy(f2.GetOutput())

def makeSphere(radius, resolution):

    s = vtk.vtkSphereSource()
    s.SetThetaResolution(resolution)
    s.SetPhiResolution(resolution)
    s.SetRadius(radius)
    s.SetEndPhi(85)
    s.Update()
    return shallowCopy(s.GetOutput())

def colorizePoints(polyData, cameraName='CAMERA_LEFT'):
    cameraView.queue.colorizePoints(cameraName, polyData)

def sendFOVRequest(channel, imagePoints):

    channelToImageType = {
        'CAMERA_LEFT' : lcmdrc.data_request_t.CAMERA_IMAGE_HEAD_LEFT,
        'CAMERACHEST_LEFT' : lcmdrc.data_request_t.CAMERA_IMAGE_LCHEST,
        'CAMERACHEST_RIGHT' : lcmdrc.data_request_t.CAMERA_IMAGE_RCHEST,
                         }

    dataRequest = lcmdrc.data_request_t()
    dataRequest.type = channelToImageType[channel]

    message = lcmdrc.subimage_request_t()
    message.data_request = dataRequest

    imagePoints = np.array([[pt[0], pt[1]] for pt in imagePoints])
    minX, maxX = imagePoints[:,0].min(), imagePoints[:,0].max()
    minY, maxY = imagePoints[:,1].min(), imagePoints[:,1].max()

    message.x = minX
    message.y = minY
    message.w = maxX - minX
    message.h = maxY - minY

    print message.x, message.y, message.w, message.h

    requestChannel = 'SUBIMAGE_REQUEST'
    lcmUtils.publish(requestChannel, message)

def testColorize():

    radius = 10
    resolution = 400
    s = makeSphere(radius, resolution)
    cameraView.queue.colorizePoints(s)
    showPolyData(p, 'sphere', colorByName='rgb')

def rayDebug(position, ray):
    d = vis.DebugData()
    d.addLine(position, position+ray*5.0)
    drcView = app.getViewManager().findView('DRC View')
    vis.updatePolyData(d.getPolyData(), 'camera ray', view=drcView)

class CameraView(object):

    def __init__(self):

        self.initImageQueue()

        if app.getMainWindow():
            self.initView()
            self.initEventFilter()
            self.rayCallback = rayDebug
        else:
            self.view = None

        self.timerCallback = TimerCallback()
        self.timerCallback.targetFps = 60
        self.timerCallback.callback = self.updateView
        self.timerCallback.start()

    def initImageQueue(self):

        imageNames = [
            'CAMERA_LEFT',
            'CAMERACHEST_LEFT',
            'CAMERACHEST_RIGHT',
            ]

        self.images = {}
        self.transforms = {}
        self.imageUtimes = {}
        self.textures = {}

        for name in imageNames:
            self.addImage(name)

        self.queue = PythonQt.dd.ddBotImageQueue(lcmUtils.getGlobalLCMThread())
        self.queue.init(lcmUtils.getGlobalLCMThread())

    def addImage(self, name):
        image = vtk.vtkImageData()
        tex = vtk.vtkTexture()
        tex.SetInput(image)
        tex.EdgeClampOn()
        tex.RepeatOff()

        self.imageUtimes[name] = 0
        self.images[name] = image
        self.textures[name] = tex
        self.transforms[name] = vtk.vtkTransform()

    def onViewDoubleClicked(self, displayPoint):

        obj, pickedPoint = vis.findPickedObject(displayPoint, self.view)

        if pickedPoint is None or not obj:
            return

        imageName = obj.getProperty('Name')
        imageUtime = self.imageUtimes[imageName]

        cameraToLocal = vtk.vtkTransform()
        self.queue.getTransform(imageName, 'local', imageUtime, cameraToLocal)

        utorsoToLocal = vtk.vtkTransform()
        self.queue.getTransform('utorso', 'local', imageUtime, utorsoToLocal)

        p = range(3)
        utorsoToLocal.TransformPoint(pickedPoint, p)

        ray = np.array(p) - np.array(cameraToLocal.GetPosition())
        ray /= np.linalg.norm(ray)

        if self.rayCallback:
            self.rayCallback(np.array(cameraToLocal.GetPosition()), ray)

    def filterEvent(self, obj, event):
        if event.type() == QtCore.QEvent.MouseButtonDblClick:
            self.eventFilter.setEventHandlerResult(True)
            self.onViewDoubleClicked(vis.mapMousePosition(obj, event))

    def initEventFilter(self):
        self.eventFilter = PythonQt.dd.ddPythonEventFilter()
        qvtkwidget = self.view.vtkWidget()
        qvtkwidget.installEventFilter(self.eventFilter)
        self.eventFilter.addFilteredEventType(QtCore.QEvent.MouseButtonDblClick)
        self.eventFilter.connect('handleEvent(QObject*, QEvent*)', self.filterEvent)

    def initView(self):

        self.view = app.getViewManager().createView('Camera View', 'VTK View')

        self.view.camera().SetViewAngle(90)
        self.view.camera().SetPosition(-7.5, 0.0, 5.0)
        self.view.camera().SetFocalPoint(0.0, 0.0, 0.0)
        self.view.camera().SetViewUp(0.0, 0.0, 1.0)

        self.sphereObjects = {}

        app.toggleCameraTerrainMode(self.view)

    def getSphereGeometry(self, imageName):

        sphereObj = self.sphereObjects.get(imageName)
        if sphereObj:
            return sphereObj

        if not self.images[imageName].GetDimensions()[0]:
            return None

        sphereResolution = 50
        sphereRadii = {
                'CAMERA_LEFT' : 9.85,
                'CAMERACHEST_LEFT' : 9.95,
                'CAMERACHEST_RIGHT' : 10
                }

        geometry = makeSphere(sphereRadii[imageName], sphereResolution)
        self.queue.computeTextureCoords(imageName, geometry)

        tcoordsArrayName = 'tcoords_%s' % imageName
        vtkNumpy.addNumpyToVtk(geometry, vtkNumpy.getNumpyFromVtk(geometry, tcoordsArrayName)[:,0].copy(), 'tcoords_U')
        vtkNumpy.addNumpyToVtk(geometry, vtkNumpy.getNumpyFromVtk(geometry, tcoordsArrayName)[:,1].copy(), 'tcoords_V')
        geometry = clipRange(geometry, 'tcoords_U', [0.0, 1.0])
        geometry = clipRange(geometry, 'tcoords_V', [0.0, 1.0])
        geometry.GetPointData().SetTCoords(geometry.GetPointData().GetArray(tcoordsArrayName))

        sphereObj = vis.showPolyData(geometry, imageName, view=self.view, parent='cameras')
        sphereObj.actor.SetTexture(self.textures[imageName])
        sphereObj.actor.GetProperty().LightingOff()

        self.sphereObjects[imageName] = sphereObj
        return sphereObj

    def updateSphereGeometry(self):

        for imageName, transform in self.transforms.iteritems():
            sphereObj = self.getSphereGeometry(imageName)
            if not sphereObj:
                continue

            self.queue.getBodyToCameraTransform(imageName, transform)
            sphereObj.actor.SetUserTransform(transform.GetLinearInverse())

        return True

    def writeImage(self, imageName, outFile):
        writer = vtk.vtkPNGWriter()
        writer.SetInput(self.images[imageName])
        writer.SetFileName(outFile)
        writer.Write()

    def updateImages(self):
        updated = False
        for imageName, image in self.images.iteritems():
            imageUtime = self.queue.getCurrentImageTime(imageName)
            if imageUtime != self.imageUtimes.get(imageName, 0):
                self.imageUtimes[imageName] = self.queue.getImage(imageName, image)
                updated = True
        return updated

    def updateView(self):

        if not self.updateImages():
            return

        if not self.view:
            return

        self.updateSphereGeometry()

        if not self.view.isVisible():
            return
        self.view.render()

class CameraImageView(object):

    def __init__(self, imageManager, imageName, viewName=None, view=None):

        self.imageManager = imageManager
        self.viewName = viewName or imageName
        self.imageName = imageName
        self.imageInitialized = False
        self.updateUtime = 0
        self.initView(view)
        self.initEventFilter()

    def onViewDoubleClicked(self, displayPoint):

        obj, pickedPoint = vis.findPickedObject(displayPoint, self.view)

        if pickedPoint is None or not obj:
            return

        imageName = obj.getProperty('Name')
        imageUtime = self.imageUtimes[imageName]

        cameraToLocal = vtk.vtkTransform()
        self.queue.getTransform(imageName, 'local', imageUtime, cameraToLocal)

        utorsoToLocal = vtk.vtkTransform()
        self.queue.getTransform('utorso', 'local', imageUtime, utorsoToLocal)

        drcView = app.getViewManager().findView('DRC View')

        p = range(3)
        utorsoToLocal.TransformPoint(pickedPoint, p)
        d = vis.DebugData()
        d.addLine(cameraToLocal.GetPosition(), p)
        vis.updatePolyData(d.getPolyData(), 'camera ray', view=drcView)

    def filterEvent(self, obj, event):
        if self.eventFilterEnabled and event.type() == QtCore.QEvent.MouseButtonDblClick:
            self.eventFilter.setEventHandlerResult(True)
            self.onViewDoubleClicked(vis.mapMousePosition(obj, event))

        elif event.type() == QtCore.QEvent.KeyPress:
            if str(event.text()).lower() == 'p':
                self.eventFilter.setEventHandlerResult(True)
            elif str(event.text()).lower() == 'r':
                self.eventFilter.setEventHandlerResult(True)
                self.resetCamera()

    def onRubberBandPick(self, obj, event):
        displayPoints = self.interactorStyle.GetStartPosition(), self.interactorStyle.GetEndPosition()
        imagePoints = [vis.pickImage(point, self.view)[1] for point in displayPoints]
        sendFOVRequest(self.imageName, imagePoints)

    def initView(self, view):
        self.view = view or app.getViewManager().createView(self.viewName, 'VTK View')
        self.view.installImageInteractor()
        self.interactorStyle = self.view.renderWindow().GetInteractor().GetInteractorStyle()
        self.interactorStyle.AddObserver('SelectionChangedEvent', self.onRubberBandPick)

        self.imageActor = vtk.vtkImageActor()
        self.imageActor.SetInput(self.imageManager.images[self.imageName])
        self.imageActor.SetVisibility(False)
        self.view.renderer().AddActor(self.imageActor)

        self.timerCallback = TimerCallback()
        self.timerCallback.targetFps = 60
        self.timerCallback.callback = self.updateView
        self.timerCallback.start()

    def initEventFilter(self):
        self.eventFilter = PythonQt.dd.ddPythonEventFilter()
        qvtkwidget = self.view.vtkWidget()
        qvtkwidget.installEventFilter(self.eventFilter)
        self.eventFilter.addFilteredEventType(QtCore.QEvent.MouseButtonDblClick)
        self.eventFilter.addFilteredEventType(QtCore.QEvent.KeyPress)
        self.eventFilter.connect('handleEvent(QObject*, QEvent*)', self.filterEvent)
        self.eventFilterEnabled = True

    def resetCamera(self):
        camera = self.view.camera()
        camera.ParallelProjectionOn()
        camera.SetFocalPoint(0,0,0)
        camera.SetPosition(0,0,-1)
        camera.SetViewUp(0,-1, 0)
        self.view.resetCamera()
        self.fitImageToView()
        self.view.render()

    def fitImageToView(self):

        camera = self.view.camera()
        image = self.imageManager.images[self.imageName]
        imageWidth, imageHeight, _ = image.GetDimensions()

        viewWidth, viewHeight = self.view.renderWindow().GetSize()
        aspectRatio = float(viewWidth)/viewHeight
        parallelScale = max(imageWidth/aspectRatio, imageHeight) / 2.0
        camera.SetParallelScale(parallelScale)

    def setImageName(self, imageName):
        if imageName == self.imageName:
            return

        assert imageName in self.imageManager.images

        self.imageName = imageName
        self.imageInitialized = False
        self.updateUtime = 0
        self.imageActor.SetInput(self.imageManager.images[self.imageName])
        self.imageActor.SetVisibility(False)
        self.view.render()

    def updateView(self):

        if not self.view.isVisible():
            return

        currentUtime = self.imageManager.imageUtimes[self.imageName]
        if currentUtime != self.updateUtime:
            self.view.render()
            self.updateUtime = currentUtime

            if not self.imageInitialized and self.imageActor.GetInput().GetDimensions()[0]:
                self.imageActor.SetVisibility(True)
                self.resetCamera()
                self.imageInitialized = True

views = {}

def addCameraView(channel, name):

    cameraView.queue.addCameraStream(channel)
    cameraView.addImage(channel)
    view = CameraImageView(cameraView, channel, name)
    global views
    views[channel] = view
    return view

def init():
    global cameraView
    cameraView = CameraView()

    global headView, chestLeft, chestRight
    headView = CameraImageView(cameraView, 'CAMERA_LEFT', 'Head camera')
    chestLeft = CameraImageView(cameraView, 'CAMERACHEST_LEFT', 'Chest left')
    chestRight = CameraImageView(cameraView, 'CAMERACHEST_RIGHT', 'Chest right')

    #addCameraView('AFFORDANCE_OVERLAY', 'Affordance overlay')

    useAnaglyph = False
    if useAnaglyph:
        addCameraView('CAMERA_ANAGLYPH', 'Head Anaglyph')

<|endoftext|>