<commit_before>import re
import copy
# https://stackoverflow.com/questions/5121931/in-python-how-can-you-load-yaml-mappings-as-ordereddicts

# stdlib
import os
from inspect import isclass
from lxml import etree
import json

# local
from jnpr.junos.exception import RpcError
from jnpr.junos.utils.start_shell import StartShell
from jnpr.junos.factory.state_machine import StateMachine
from jnpr.junos.factory.to_json import TableJSONEncoder

from jinja2 import Template
from ntc_templates import parse as ntc_parse

import logging
logger = logging.getLogger("jnpr.junos.factory.cmdtable")

class CMDTable(object):

    def __init__(self, dev=None, data=None, path=None, template_dir=None):
        """
        :dev: Device instance
        :data: string blob of the command output
        :path: file path to XML, to be used rather than :dev:
        :template_dir: To look for textfsm templates in this folder first
        """
        self._dev = dev
        self.xml = None
        self.view = None
        self.ITEM_FILTER = 'name'
        self._key_list = []
        self._path = path
        self._parser = None
        self.output = None
        self.data = data
        self.template_dir = template_dir

    # -------------------------------------------------------------------------
    # PUBLIC METHODS
    # -------------------------------------------------------------------------

    def get(self, *vargs, **kvargs):
        """
        Retrieve the XML (string blob under <output> tag of table data from the
        Device instance and returns back the Table instance - for call-chaining
        purposes.

        If the Table was created with a :path: rather than a Device,
        then this method will load the string blob from that file.  In this
        case, the \*vargs, and \**kvargs are not used.

        ALIAS: __call__

        :vargs:
          [0] is the table :arg_key: value.  This is used so that
          the caller can retrieve just one item from the table without
          having to know the Junos RPC argument.

        :kvargs:
          these are the name/value pairs relating to the specific Junos
          command attached to the table.  For example, if the command
          is 'show ithrottle id {{ id }}', there is a parameter 'id'.
          Any valid command argument can be passed to :kvargs: to further filter
          the results of the :get(): operation.  neato!

        NOTES:
          If you need to create a 'stub' for unit-testing
          purposes, you want to create a subclass of your table and
          overload this methods.
        """
        self._clearkeys()

        if self._path and self.data:
            raise AttributeError('path and data are mutually exclusive')
        if self._path is not None:
            # for loading from local file-path
            with open(self._path, 'r') as fp:
                self.data = fp.read().strip()
            if self.data.startswith('<output>') and self.data.endswith(
                    '</output>'):
                self.data = etree.fromstring(self.data).text

        if 'target' in kvargs:
            self.TARGET = kvargs['target']

        if 'key' in kvargs:
            self.KEY = kvargs['key']

        if 'key_items' in kvargs:
            self.KEY_ITEMS = kvargs['key_items']

        if 'filters' in kvargs:
            self.VIEW.FILTERS = [kvargs['filters']] if isinstance(kvargs[
                                                                   'filters'],
                                                                str) else \
                kvargs['filters']

        if 'args' in kvargs and isinstance(kvargs['args'], dict):
            self.CMD_ARGS.update(kvargs['args'])

        if len(self.CMD_ARGS) > 0:
            self.GET_CMD = Template(self.GET_CMD).render(**self.CMD_ARGS)

        if self.data is None:
            # execute the Junos RPC to retrieve the table
            if hasattr(self, 'TARGET'):
                if self.TARGET is None:
                    raise ValueError('"target" value not provided')
                rpc_args = {'target': self.TARGET, 'command': self.GET_CMD,
                            'timeout': '0'}
                try:
                    self.xml = getattr(self.RPC, 'request_pfe_execute')(**rpc_args)
                    self.data = self.xml.text
                    ver_info = self._dev.facts.get('version_info')
                    if ver_info and ver_info.major[0] <= 15:
                        # Junos <=15.x output has output something like below
                        #
                        # <rpc-reply>
                        # <output>
                        # SENT: Ukern command: show memory
                        # GOT:
                        # GOT: ID      Base      Total(b)       Free(b)     Used(b)
                        # GOT: --  --------     ---------     ---------   ---------
                        # GOT:  0  44e72078    1882774284    1689527364   193246920
                        # GOT:  1  b51ffb88      67108860      57651900     9456960
                        # GOT:  2  bcdfffe0      52428784      52428784           0
                        # GOT:  3  b91ffb88      62914556      62914556           0
                        # LOCAL: End of file
                        # </output>
                        # </rpc-reply>
                        # hence need to do cleanup
                        self.data = self.data.replace("GOT: ", "")
                except RpcError:
                    with StartShell(self.D) as ss:
                        ret = ss.run('cprod -A %s -c "%s"' % (self.TARGET,
                                                              self.GET_CMD))
                        if ret[0]:
                            self.data = ret[1]
            else:
                self.xml = self.RPC.cli(self.GET_CMD)
                self.data = self.xml.text

        if self.USE_TEXTFSM:
            self.output = self._parse_textfsm(platform=self.PLATFORM,
                                              command=self.GET_CMD,
                                              data=self.data)
        else:
            # state machine
            sm = StateMachine(self)
            self.output = sm.parse(self.data.splitlines())

        # returning self for call-chaining purposes, yo!
        return self

    # -------------------------------------------------------------------------
    # PROPERTIES
    # -------------------------------------------------------------------------

    @property
    def D(self):
        """ the Device instance """
        return self._dev

    @property
    def CLI(self):
        """ the Device.cli instance """
        return self.D.cli

    @property
    def RPC(self):
        """ the Device.rpc instance """
        return self.D.rpc

    @property
    def view(self):
        """ returns the current view assigned to this table """
        return self._view

    @view.setter
    def view(self, cls):
        """ assigns a new view to the table """
        if cls is None:
            self._view = None
            return

        if not isclass(cls):
            raise ValueError("Must be given RunstatView class")

        self._view = cls

    @property
    def hostname(self):
        return self.D.hostname

    @property
    def key_list(self):
        """ the list of keys, as property for caching """
        return self._key_list

    # -------------------------------------------------------------------------
    # PRIVATE METHODS
    # -------------------------------------------------------------------------

    def _assert_data(self):
        if self.data is None:
            raise RuntimeError("Table is empty, use get()")

    def _clearkeys(self):
        self._key_list = []

    # -------------------------------------------------------------------------
    # PUBLIC METHODS
    # -------------------------------------------------------------------------

    # ------------------------------------------------------------------------
    # keys
    # ------------------------------------------------------------------------

    def _keys(self):
        """ return a list of data item keys from the data string """

        return self.output.keys()

    def keys(self):
        # if the key_list has been cached, then use it
        if len(self.key_list):
            return self.key_list

        # otherwise, build the list of keys into the cache
        self._key_list = self._keys()
        return self._key_list

    # ------------------------------------------------------------------------
    # values
    # ------------------------------------------------------------------------

    def values(self):
        """ returns list of table entry items() """

        self._assert_data()
        return self.output.values()

    # ------------------------------------------------------------------------
    # items
    # ------------------------------------------------------------------------

    def items(self):
        """ returns list of tuple(name,values) for each table entry """
        return self.output.items()

    def to_json(self):
        """
        :returns: JSON encoded string of entire Table contents
        """
        return json.dumps(self, cls=TableJSONEncoder)

    # -------------------------------------------------------------------------
    # OVERLOADS
    # -------------------------------------------------------------------------

    __call__ = get

    def __repr__(self):
        cls_name = self.__class__.__name__
        source = self.D.hostname if self.D is not None else self._path

        if self.data is None:
            return "%s:%s - Table empty" % (cls_name, source)
        else:
            n_items = len(self.keys())
            return "%s:%s: %s items" % (cls_name, source, n_items)

    def __len__(self):
        self._assert_data()
        return len(self.keys())

    def __iter__(self):
        """ iterate over each time in the table """
        self._assert_data()

        for key, value in self.output.items():
            yield key, value

    def __getitem__(self, value):
        """
        returns a table item. If a table view is set (should be by default)
        then the item will be converted to the view upon return.  if there is
        no table view, then the XML object will be returned.

        :value:
          for <string>, this will perform a select based on key-name
          for <tuple>, this will perform a select based on compsite key-name
          for <int>, this will perform a select based by position, like <list>
            [0] is the first item
            [-1] is the last item
          when it is a <slice> then this will return a <list> of View widgets
        """
        self._assert_data()
        return self.output[value]

    def __contains__(self, key):
        """ membership for use with 'in' """
        return bool(key in self.keys())

    # ------------------------------------------------------------------------
    # textfsm
    # ------------------------------------------------------------------------

    def _parse_textfsm(self, platform=None, command=None, data=None):
        """ textfsm returns list of dict, make it JSON/dict """
        attrs = dict(
            Command=command,
            Platform=platform
        )

        template = None
        template_dir = None
        if self.template_dir is not None:
            # we dont need index file foor lookup
            index = None
            template_path = os.path.join(self.template_dir, '{}_{}.textfsm'.format(
                platform, '_'.join(command.split())))
            if not os.path.exists(template_path):
                logger.debug('Template %s file missing, '
                             'looking default ntc-templates location' % template_path)
            else:
                template = template_path
                template_dir = self.template_dir
        if template_dir is None:
            index = 'index'
            template_dir = ntc_parse._get_template_dir()

        cli_table = ntc_parse.clitable.CliTable(index, template_dir)
        try:
            cli_table.ParseCmd(data, attrs, template)
        except ntc_parse.clitable.CliTableError as ex:
            logger.error('Unable to parse command "%s" on platform %s' % (command, platform))
            raise ex
        return self._filter_output(cli_table)

    def _filter_output(self, cli_table):
        """
        textfsm return list of list, covert it into more consumable format

        :param cli_table: CLiTable object from textfsm
        :return: dict of key, fields and its values
        """
        self._set_key(cli_table)
        output = {}
        fields = self.VIEW.FIELDS if self.VIEW is not None else {}
        reverse_fields = {val: key for key, val in fields.items()}
        for row in cli_table:
            temp_dict = {}
            for index, element in enumerate(row):
                key = cli_table.header[index]
                if key in self.KEY:
                    temp_dict[key] = element
                if reverse_fields:
                    if key in reverse_fields:
                        temp_dict[reverse_fields[key]] = element
                else:
                    temp_dict[key] = element
            logger.debug("data at index {} is {}".format(row.row, temp_dict))
            if self.KEY in temp_dict:
                if self.KEY not in fields:
                    output[temp_dict.pop(self.KEY)] = temp_dict
                else:
                    output[temp_dict[self.KEY]] = temp_dict
            else:
                logger.debug("Key {} not present in {}".format(self.KEY,
                                                               temp_dict))
        return output

    def _set_key(self, cli_table):
        """
        Preference to user provided key, then template and at last default
        Checks and update if we need KEY from template file

        :param cli_table: CLiTable object from textfsm
        :return:
        """
        if self.KEY == 'name' and len(cli_table._keys) > 0:
            template_keys = list(cli_table._keys)
            self.KEY = template_keys[0] if len(template_keys) == 1 else \
                template_keys
        logger.debug("KEY being used: {}".format(self.KEY))
<commit_msg>handle single row with null as key<commit_after>import re
import copy
# https://stackoverflow.com/questions/5121931/in-python-how-can-you-load-yaml-mappings-as-ordereddicts

# stdlib
import os
from inspect import isclass
from lxml import etree
import json

# local
from jnpr.junos.exception import RpcError
from jnpr.junos.utils.start_shell import StartShell
from jnpr.junos.factory.state_machine import StateMachine
from jnpr.junos.factory.to_json import TableJSONEncoder

from jinja2 import Template
from ntc_templates import parse as ntc_parse

import logging
logger = logging.getLogger("jnpr.junos.factory.cmdtable")

class CMDTable(object):

    def __init__(self, dev=None, data=None, path=None, template_dir=None):
        """
        :dev: Device instance
        :data: string blob of the command output
        :path: file path to XML, to be used rather than :dev:
        :template_dir: To look for textfsm templates in this folder first
        """
        self._dev = dev
        self.xml = None
        self.view = None
        self.ITEM_FILTER = 'name'
        self._key_list = []
        self._path = path
        self._parser = None
        self.output = None
        self.data = data
        self.template_dir = template_dir

    # -------------------------------------------------------------------------
    # PUBLIC METHODS
    # -------------------------------------------------------------------------

    def get(self, *vargs, **kvargs):
        """
        Retrieve the XML (string blob under <output> tag of table data from the
        Device instance and returns back the Table instance - for call-chaining
        purposes.

        If the Table was created with a :path: rather than a Device,
        then this method will load the string blob from that file.  In this
        case, the \*vargs, and \**kvargs are not used.

        ALIAS: __call__

        :vargs:
          [0] is the table :arg_key: value.  This is used so that
          the caller can retrieve just one item from the table without
          having to know the Junos RPC argument.

        :kvargs:
          these are the name/value pairs relating to the specific Junos
          command attached to the table.  For example, if the command
          is 'show ithrottle id {{ id }}', there is a parameter 'id'.
          Any valid command argument can be passed to :kvargs: to further filter
          the results of the :get(): operation.  neato!

        NOTES:
          If you need to create a 'stub' for unit-testing
          purposes, you want to create a subclass of your table and
          overload this methods.
        """
        self._clearkeys()

        if self._path and self.data:
            raise AttributeError('path and data are mutually exclusive')
        if self._path is not None:
            # for loading from local file-path
            with open(self._path, 'r') as fp:
                self.data = fp.read().strip()
            if self.data.startswith('<output>') and self.data.endswith(
                    '</output>'):
                self.data = etree.fromstring(self.data).text

        if 'target' in kvargs:
            self.TARGET = kvargs['target']

        if 'key' in kvargs:
            self.KEY = kvargs['key']

        if 'key_items' in kvargs:
            self.KEY_ITEMS = kvargs['key_items']

        if 'filters' in kvargs:
            self.VIEW.FILTERS = [kvargs['filters']] if isinstance(kvargs[
                                                                   'filters'],
                                                                str) else \
                kvargs['filters']

        if 'args' in kvargs and isinstance(kvargs['args'], dict):
            self.CMD_ARGS.update(kvargs['args'])

        if len(self.CMD_ARGS) > 0:
            self.GET_CMD = Template(self.GET_CMD).render(**self.CMD_ARGS)

        if self.data is None:
            # execute the Junos RPC to retrieve the table
            if hasattr(self, 'TARGET'):
                if self.TARGET is None:
                    raise ValueError('"target" value not provided')
                rpc_args = {'target': self.TARGET, 'command': self.GET_CMD,
                            'timeout': '0'}
                try:
                    self.xml = getattr(self.RPC, 'request_pfe_execute')(**rpc_args)
                    self.data = self.xml.text
                    ver_info = self._dev.facts.get('version_info')
                    if ver_info and ver_info.major[0] <= 15:
                        # Junos <=15.x output has output something like below
                        #
                        # <rpc-reply>
                        # <output>
                        # SENT: Ukern command: show memory
                        # GOT:
                        # GOT: ID      Base      Total(b)       Free(b)     Used(b)
                        # GOT: --  --------     ---------     ---------   ---------
                        # GOT:  0  44e72078    1882774284    1689527364   193246920
                        # GOT:  1  b51ffb88      67108860      57651900     9456960
                        # GOT:  2  bcdfffe0      52428784      52428784           0
                        # GOT:  3  b91ffb88      62914556      62914556           0
                        # LOCAL: End of file
                        # </output>
                        # </rpc-reply>
                        # hence need to do cleanup
                        self.data = self.data.replace("GOT: ", "")
                except RpcError:
                    with StartShell(self.D) as ss:
                        ret = ss.run('cprod -A %s -c "%s"' % (self.TARGET,
                                                              self.GET_CMD))
                        if ret[0]:
                            self.data = ret[1]
            else:
                self.xml = self.RPC.cli(self.GET_CMD)
                self.data = self.xml.text

        if self.USE_TEXTFSM:
            self.output = self._parse_textfsm(platform=self.PLATFORM,
                                              command=self.GET_CMD,
                                              data=self.data)
        else:
            # state machine
            sm = StateMachine(self)
            self.output = sm.parse(self.data.splitlines())

        # returning self for call-chaining purposes, yo!
        return self

    # -------------------------------------------------------------------------
    # PROPERTIES
    # -------------------------------------------------------------------------

    @property
    def D(self):
        """ the Device instance """
        return self._dev

    @property
    def CLI(self):
        """ the Device.cli instance """
        return self.D.cli

    @property
    def RPC(self):
        """ the Device.rpc instance """
        return self.D.rpc

    @property
    def view(self):
        """ returns the current view assigned to this table """
        return self._view

    @view.setter
    def view(self, cls):
        """ assigns a new view to the table """
        if cls is None:
            self._view = None
            return

        if not isclass(cls):
            raise ValueError("Must be given RunstatView class")

        self._view = cls

    @property
    def hostname(self):
        return self.D.hostname

    @property
    def key_list(self):
        """ the list of keys, as property for caching """
        return self._key_list

    # -------------------------------------------------------------------------
    # PRIVATE METHODS
    # -------------------------------------------------------------------------

    def _assert_data(self):
        if self.data is None:
            raise RuntimeError("Table is empty, use get()")

    def _clearkeys(self):
        self._key_list = []

    # -------------------------------------------------------------------------
    # PUBLIC METHODS
    # -------------------------------------------------------------------------

    # ------------------------------------------------------------------------
    # keys
    # ------------------------------------------------------------------------

    def _keys(self):
        """ return a list of data item keys from the data string """

        return self.output.keys()

    def keys(self):
        # if the key_list has been cached, then use it
        if len(self.key_list):
            return self.key_list

        # otherwise, build the list of keys into the cache
        self._key_list = self._keys()
        return self._key_list

    # ------------------------------------------------------------------------
    # values
    # ------------------------------------------------------------------------

    def values(self):
        """ returns list of table entry items() """

        self._assert_data()
        return self.output.values()

    # ------------------------------------------------------------------------
    # items
    # ------------------------------------------------------------------------

    def items(self):
        """ returns list of tuple(name,values) for each table entry """
        return self.output.items()

    def to_json(self):
        """
        :returns: JSON encoded string of entire Table contents
        """
        return json.dumps(self, cls=TableJSONEncoder)

    # -------------------------------------------------------------------------
    # OVERLOADS
    # -------------------------------------------------------------------------

    __call__ = get

    def __repr__(self):
        cls_name = self.__class__.__name__
        source = self.D.hostname if self.D is not None else self._path

        if self.data is None:
            return "%s:%s - Table empty" % (cls_name, source)
        else:
            n_items = len(self.keys())
            return "%s:%s: %s items" % (cls_name, source, n_items)

    def __len__(self):
        self._assert_data()
        return len(self.keys())

    def __iter__(self):
        """ iterate over each time in the table """
        self._assert_data()

        for key, value in self.output.items():
            yield key, value

    def __getitem__(self, value):
        """
        returns a table item. If a table view is set (should be by default)
        then the item will be converted to the view upon return.  if there is
        no table view, then the XML object will be returned.

        :value:
          for <string>, this will perform a select based on key-name
          for <tuple>, this will perform a select based on compsite key-name
          for <int>, this will perform a select based by position, like <list>
            [0] is the first item
            [-1] is the last item
          when it is a <slice> then this will return a <list> of View widgets
        """
        self._assert_data()
        return self.output[value]

    def __contains__(self, key):
        """ membership for use with 'in' """
        return bool(key in self.keys())

    # ------------------------------------------------------------------------
    # textfsm
    # ------------------------------------------------------------------------

    def _parse_textfsm(self, platform=None, command=None, data=None):
        """
        textfsm returns list of dict, make it JSON/dict

        :param platform: vendor platform, for ex cisco_xr
        :param command: cli command to be parsed
        :param data: string blob output from the cli command execution
        :return: dict of parsed data.
        """
        attrs = dict(
            Command=command,
            Platform=platform
        )

        template = None
        template_dir = None
        if self.template_dir is not None:
            # we dont need index file for lookup
            index = None
            template_path = os.path.join(self.template_dir,
                                         '{}_{}.textfsm'.format(
                                             platform,
                                             '_'.join(command.split())))
            if not os.path.exists(template_path):
                msg = 'Template file %s missing' % template_path
                logger.error(msg)
                raise FileNotFoundError(msg)
            else:
                template = template_path
                template_dir = self.template_dir
        if template_dir is None:
            index = 'index'
            template_dir = ntc_parse._get_template_dir()

        cli_table = ntc_parse.clitable.CliTable(index, template_dir)
        try:
            cli_table.ParseCmd(data, attrs, template)
        except ntc_parse.clitable.CliTableError as ex:
            logger.error('Unable to parse command "%s" on platform %s' % (
                command, platform))
            raise ex
        return self._filter_output(cli_table)

    def _filter_output(self, cli_table):
        """
        textfsm return list of list, covert it into more consumable format

        :param cli_table: CLiTable object from textfsm
        :return: dict of key, fields and its values, list of dict when key is None
        """
        self._set_key(cli_table)

        fields = self.VIEW.FIELDS if self.VIEW is not None else {}
        reverse_fields = {val: key for key, val in fields.items()}
        if self.KEY is None:
            cli_table_size = cli_table.size
            if cli_table_size > 1:
                raise KeyError("Key is Mandatory for parsed o/p of %s "
                               "length" % cli_table_size)
            elif cli_table_size == 1:
                temp_dict = self._parse_row(cli_table[1], cli_table,
                                            reverse_fields)
                logger.debug("For Null Key, data returned: {}".format(temp_dict))
                return temp_dict
        output = {}
        for row in cli_table:
            temp_dict = self._parse_row(row, cli_table, reverse_fields)
            logger.debug("data at index {} is {}".format(row.row, temp_dict))
            if self.KEY in temp_dict:
                if self.KEY not in fields:
                    output[temp_dict.pop(self.KEY)] = temp_dict
                else:
                    output[temp_dict[self.KEY]] = temp_dict
            else:
                logger.debug("Key {} not present in {}".format(self.KEY,
                                                               temp_dict))
        return output

    def _parse_row(self, row, cli_table, reverse_fields):
        temp_dict = {}
        for index, element in enumerate(row):
            key = cli_table.header[index]
            if self.KEY and key in self.KEY:
                temp_dict[key] = element
            if reverse_fields:
                if key in reverse_fields:
                    temp_dict[reverse_fields[key]] = element
            else:
                temp_dict[key] = element
        return temp_dict

    def _set_key(self, cli_table):
        """
        Preference to user provided key, then template and at last default
        Checks and update if we need KEY from template file

        :param cli_table: CLiTable object from textfsm
        :return:
        """
        if self.KEY == 'name' and len(cli_table._keys) > 0:
            template_keys = list(cli_table._keys)
            self.KEY = template_keys[0] if len(template_keys) == 1 else \
                template_keys
        logger.debug("KEY being used: {}".format(self.KEY))
<|endoftext|>


<commit_before>"""
Emissivity models and atomic calculations for simulating radiation from coronal loops
"""

from .chianti import ChIon
from .emiss_models import EquilibriumModel
<commit_msg>pull in emission model<commit_after>"""
Emissivity models and atomic calculations for simulating radiation from coronal loops
"""

from .chianti import ChIon
from .emission_models import EquilibriumEmissionModel
<|endoftext|>


<commit_before>import json
import requests

from .device import Device
from .channel import Channel
from .contact import Contact
from .errors import PushBulletError, InvalidKeyError, PushError
from .filetype import get_file_type

class PushBullet(object):

    DEVICES_URL = "https://api.pushbullet.com/v2/devices"
    CONTACTS_URL = "https://api.pushbullet.com/v2/contacts"
    CHANNELS_URL = "https://api.pushbullet.com/v2/channels"
    ME_URL = "https://api.pushbullet.com/v2/users/me"
    PUSH_URL = "https://api.pushbullet.com/v2/pushes"
    UPLOAD_REQUEST_URL = "https://api.pushbullet.com/v2/upload-request"

    def __init__(self, api_key):
        self.api_key = api_key
        self._json_header = {'Content-Type': 'application/json'}

        self._session = requests.Session()
        self._session.auth = (self.api_key, "")
        self._session.headers.update(self._json_header)

        self.refresh()

    def _get_data(self, url):
        resp = self._session.get(url)

        if resp.status_code == 401:
            raise InvalidKeyError()

        return resp.json()

    def _load_devices(self):
        self.devices = []

        resp_dict = self._get_data(self.DEVICES_URL)
        device_list = resp_dict.get("devices", [])

        for device_info in device_list:
            if device_info.get("active"):
                d = Device(self, device_info)
                self.devices.append(d)

    def _load_contacts(self):
        self.contacts = []

        resp_dict = self._get_data(self.CONTACTS_URL)
        contacts_list = resp_dict.get("contacts", [])

        for contact_info in contacts_list:
            if contact_info.get("active"):
                c = Contact(self, contact_info)
                self.contacts.append(c)

    def _load_user_info(self):
        self.user_info = self._get_data(self.ME_URL)

    def _load_channels(self):
        self.channels = []

        resp_dict = self._get_data(self.CHANNELS_URL)
        channel_list = resp_dict.get("channels", [])

        for channel_info in channel_list:
            if channel_info.get("active"):
                c = Channel(self, channel_info)
                self.channels.append(c)

    @staticmethod
    def _recipient(device=None, contact=None, email=None, channel=None):
        data = dict()

        if device:
            data["device_iden"] = device.device_iden
        elif contact:
            data["email"] = contact.email
        elif email:
            data["email"] = email
        elif channel:
            data["channel_tag"] = channel.channel_tag

        return data

    def new_device(self, nickname):
        data = {"nickname": nickname, "type": "stream"}
        r = self._session.post(self.DEVICES_URL, data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_device = Device(self, r.json())
            self.devices.append(new_device)
            return new_device
        else:
            raise PushBulletError(r.text)

    def new_contact(self, name, email):
        data = {"name": name, "email": email}
        r = self._session.post(self.CONTACTS_URL, data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_contact = Contact(self, r.json())
            self.contacts.append(new_contact)
            return new_contact
        else:
            raise PushBulletError(r.text)

    def edit_device(self, device, nickname=None, model=None, manufacturer=None):
        data = {"nickname": nickname}
        iden = device.device_iden
        r = self._session.post("{}/{}".format(self.DEVICES_URL, iden), data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_device = Device(self, r.json())
            self.devices[self.devices.index(device)] = new_device
            return new_device
        else:
            raise PushBulletError(r.text)

    def edit_contact(self, contact, name):
        data = {"name": name}
        iden = contact.iden
        r = self._session.post("{}/{}".format(self.CONTACTS_URL, iden), data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_contact = Contact(self, r.json())
            self.contacts[self.contacts.index(contact)] = new_contact
            return new_contact
        else:
            raise PushBulletError(r.text)

    def remove_device(self, device):
        iden = device.device_iden
        r = self._session.delete("{}/{}".format(self.DEVICES_URL, iden))
        if r.status_code == requests.codes.ok:
            self.devices.remove(device)
        else:
            raise PushBulletError(r.text)

    def remove_contact(self, contact):
        iden = contact.iden
        r = self._session.delete("{}/{}".format(self.CONTACTS_URL, iden))
        if r.status_code == requests.codes.ok:
            self.contacts.remove(contact)
            return True
        else:
            raise PushBulletError(r.text)

    def get_pushes(self, modified_after=None, limit=None):
        data = {"modified_after": modified_after, "limit": limit}

        pushes_list = []
        get_more_pushes = True
        while get_more_pushes:
            r = self._session.get(self.PUSH_URL, params=data)
            if r.status_code != requests.codes.ok:
                raise PushBulletError(r.text)

            pushes_list += r.json().get("pushes")
            if 'cursor' in r.json() and (not limit or len(pushes_list) < limit):
                data['cursor'] = r.json()['cursor']
            else:
                get_more_pushes = False

        return True, pushes_list

    def dismiss_push(self, iden):
        data = {"dismissed": True}
        r = self._session.post("{}/{}".format(self.PUSH_URL, iden), data=json.dumps(data))

        if r.status_code =! requests.codes.ok:
            raise PushBulletError(r.text)

    def delete_push(self, iden):
        r = self._session.delete("{}/{}".format(self.PUSH_URL, iden))

        if r.status_code != requests.codes.ok:
            raise PushBulletError(r.text)

    def upload_file(self, f, file_name, file_type=None):
        if not file_type:
            file_type = get_file_type(f, file_name)

        data = {"file_name": file_name, "file_type": file_type}

        # Request url for file upload
        r = self._session.post(self.UPLOAD_REQUEST_URL, data=json.dumps(data))

        if r.status_code != requests.codes.ok:
            raise PushBulletError(r.text)

        upload_data = r.json().get("data")
        file_url = r.json().get("file_url")
        upload_url = r.json().get("upload_url")

        upload = requests.post(upload_url, data=upload_data, files={"file": f})

        return {"file_type": file_type, "file_url": file_url, "file_name": file_name}

    def push_file(self, file_name, file_url, file_type, body=None, device=None, contact=None, email=None, channel=None):
        data = {"type": "file", "file_type": file_type, "file_url": file_url, "file_name": file_name}
        if body:
            data["body"] = body

        data.update(PushBullet._recipient(device, contact, email, channel))

        return self._push(data)

    def push_note(self, title, body, device=None, contact=None, email=None):
        data = {"type": "note", "title": title, "body": body}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def push_address(self, name, address, device=None, contact=None, email=None):
        data = {"type": "address", "name": name, "address": address}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def push_list(self, title, items, device=None, contact=None, email=None):
        data = {"type": "list", "title": title, "items": items}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def push_link(self, title, url, body=None, device=None, contact=None, email=None):
        data = {"type": "link", "title": title, "url": url, "body": body}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def _push(self, data):
        r = self._session.post(self.PUSH_URL, data=json.dumps(data))

        if r.status_code == requests.codes.ok:
            return r.json()
        else:
            raise PushError(r.text)

    def refresh(self):
        self._load_devices()
        self._load_contacts()
        self._load_user_info()
        self._load_channels()
<commit_msg>Syntax hiccup<commit_after>import json
import requests

from .device import Device
from .channel import Channel
from .contact import Contact
from .errors import PushBulletError, InvalidKeyError, PushError
from .filetype import get_file_type

class PushBullet(object):

    DEVICES_URL = "https://api.pushbullet.com/v2/devices"
    CONTACTS_URL = "https://api.pushbullet.com/v2/contacts"
    CHANNELS_URL = "https://api.pushbullet.com/v2/channels"
    ME_URL = "https://api.pushbullet.com/v2/users/me"
    PUSH_URL = "https://api.pushbullet.com/v2/pushes"
    UPLOAD_REQUEST_URL = "https://api.pushbullet.com/v2/upload-request"

    def __init__(self, api_key):
        self.api_key = api_key
        self._json_header = {'Content-Type': 'application/json'}

        self._session = requests.Session()
        self._session.auth = (self.api_key, "")
        self._session.headers.update(self._json_header)

        self.refresh()

    def _get_data(self, url):
        resp = self._session.get(url)

        if resp.status_code == 401:
            raise InvalidKeyError()

        return resp.json()

    def _load_devices(self):
        self.devices = []

        resp_dict = self._get_data(self.DEVICES_URL)
        device_list = resp_dict.get("devices", [])

        for device_info in device_list:
            if device_info.get("active"):
                d = Device(self, device_info)
                self.devices.append(d)

    def _load_contacts(self):
        self.contacts = []

        resp_dict = self._get_data(self.CONTACTS_URL)
        contacts_list = resp_dict.get("contacts", [])

        for contact_info in contacts_list:
            if contact_info.get("active"):
                c = Contact(self, contact_info)
                self.contacts.append(c)

    def _load_user_info(self):
        self.user_info = self._get_data(self.ME_URL)

    def _load_channels(self):
        self.channels = []

        resp_dict = self._get_data(self.CHANNELS_URL)
        channel_list = resp_dict.get("channels", [])

        for channel_info in channel_list:
            if channel_info.get("active"):
                c = Channel(self, channel_info)
                self.channels.append(c)

    @staticmethod
    def _recipient(device=None, contact=None, email=None, channel=None):
        data = dict()

        if device:
            data["device_iden"] = device.device_iden
        elif contact:
            data["email"] = contact.email
        elif email:
            data["email"] = email
        elif channel:
            data["channel_tag"] = channel.channel_tag

        return data

    def new_device(self, nickname):
        data = {"nickname": nickname, "type": "stream"}
        r = self._session.post(self.DEVICES_URL, data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_device = Device(self, r.json())
            self.devices.append(new_device)
            return new_device
        else:
            raise PushBulletError(r.text)

    def new_contact(self, name, email):
        data = {"name": name, "email": email}
        r = self._session.post(self.CONTACTS_URL, data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_contact = Contact(self, r.json())
            self.contacts.append(new_contact)
            return new_contact
        else:
            raise PushBulletError(r.text)

    def edit_device(self, device, nickname=None, model=None, manufacturer=None):
        data = {"nickname": nickname}
        iden = device.device_iden
        r = self._session.post("{}/{}".format(self.DEVICES_URL, iden), data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_device = Device(self, r.json())
            self.devices[self.devices.index(device)] = new_device
            return new_device
        else:
            raise PushBulletError(r.text)

    def edit_contact(self, contact, name):
        data = {"name": name}
        iden = contact.iden
        r = self._session.post("{}/{}".format(self.CONTACTS_URL, iden), data=json.dumps(data))
        if r.status_code == requests.codes.ok:
            new_contact = Contact(self, r.json())
            self.contacts[self.contacts.index(contact)] = new_contact
            return new_contact
        else:
            raise PushBulletError(r.text)

    def remove_device(self, device):
        iden = device.device_iden
        r = self._session.delete("{}/{}".format(self.DEVICES_URL, iden))
        if r.status_code == requests.codes.ok:
            self.devices.remove(device)
        else:
            raise PushBulletError(r.text)

    def remove_contact(self, contact):
        iden = contact.iden
        r = self._session.delete("{}/{}".format(self.CONTACTS_URL, iden))
        if r.status_code == requests.codes.ok:
            self.contacts.remove(contact)
            return True
        else:
            raise PushBulletError(r.text)

    def get_pushes(self, modified_after=None, limit=None):
        data = {"modified_after": modified_after, "limit": limit}

        pushes_list = []
        get_more_pushes = True
        while get_more_pushes:
            r = self._session.get(self.PUSH_URL, params=data)
            if r.status_code != requests.codes.ok:
                raise PushBulletError(r.text)

            pushes_list += r.json().get("pushes")
            if 'cursor' in r.json() and (not limit or len(pushes_list) < limit):
                data['cursor'] = r.json()['cursor']
            else:
                get_more_pushes = False

        return True, pushes_list

    def dismiss_push(self, iden):
        data = {"dismissed": True}
        r = self._session.post("{}/{}".format(self.PUSH_URL, iden), data=json.dumps(data))

        if r.status_code != requests.codes.ok:
            raise PushBulletError(r.text)

    def delete_push(self, iden):
        r = self._session.delete("{}/{}".format(self.PUSH_URL, iden))

        if r.status_code != requests.codes.ok:
            raise PushBulletError(r.text)

    def upload_file(self, f, file_name, file_type=None):
        if not file_type:
            file_type = get_file_type(f, file_name)

        data = {"file_name": file_name, "file_type": file_type}

        # Request url for file upload
        r = self._session.post(self.UPLOAD_REQUEST_URL, data=json.dumps(data))

        if r.status_code != requests.codes.ok:
            raise PushBulletError(r.text)

        upload_data = r.json().get("data")
        file_url = r.json().get("file_url")
        upload_url = r.json().get("upload_url")

        upload = requests.post(upload_url, data=upload_data, files={"file": f})

        return {"file_type": file_type, "file_url": file_url, "file_name": file_name}

    def push_file(self, file_name, file_url, file_type, body=None, device=None, contact=None, email=None, channel=None):
        data = {"type": "file", "file_type": file_type, "file_url": file_url, "file_name": file_name}
        if body:
            data["body"] = body

        data.update(PushBullet._recipient(device, contact, email, channel))

        return self._push(data)

    def push_note(self, title, body, device=None, contact=None, email=None):
        data = {"type": "note", "title": title, "body": body}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def push_address(self, name, address, device=None, contact=None, email=None):
        data = {"type": "address", "name": name, "address": address}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def push_list(self, title, items, device=None, contact=None, email=None):
        data = {"type": "list", "title": title, "items": items}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def push_link(self, title, url, body=None, device=None, contact=None, email=None):
        data = {"type": "link", "title": title, "url": url, "body": body}

        data.update(PushBullet._recipient(device, contact, email))

        return self._push(data)

    def _push(self, data):
        r = self._session.post(self.PUSH_URL, data=json.dumps(data))

        if r.status_code == requests.codes.ok:
            return r.json()
        else:
            raise PushError(r.text)

    def refresh(self):
        self._load_devices()
        self._load_contacts()
        self._load_user_info()
        self._load_channels()
<|endoftext|>


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

"""
Copyright [2009-2021] EMBL-European Bioinformatics Institute
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.
"""

from pathlib import Path

import click

from rnacentral_pipeline.databases.zwd import parser
from rnacentral_pipeline.writers import entry_writer

@click.group("zwd")
def cli():
    """
    Commands for parsing ENA data.
    """

@cli.command("parse")
@click.argument("context", type=click.Path())
@click.argument("json_file", type=click.File("r"))
@click.argument(
    "output",
    default=".",
    type=click.Path(writable=True, dir_okay=True, file_okay=False),
)
def process_json_schema(context, json_file, output):
    """
    This parses our JSON schema files to produce the importable CSV files.
    """
    entries = parser.parse(Path(context), Path(json_file))
    with entry_writer(Path(output)) as writer:
        writer.write(entries)
<commit_msg>Fix JSON file argument<commit_after># -*- coding: utf-8 -*-

"""
Copyright [2009-2021] EMBL-European Bioinformatics Institute
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.
"""

from pathlib import Path

import click

from rnacentral_pipeline.databases.zwd import parser
from rnacentral_pipeline.writers import entry_writer

@click.group("zwd")
def cli():
    """
    Commands for parsing ENA data.
    """

@cli.command("parse")
@click.argument("context", type=click.Path())
@click.argument("json_file", type=click.Path())
@click.argument(
    "output",
    default=".",
    type=click.Path(writable=True, dir_okay=True, file_okay=False),
)
def process_json_schema(context, json_file, output):
    """
    This parses our JSON schema files to produce the importable CSV files.
    """
    entries = parser.parse(Path(context), Path(json_file))
    with entry_writer(Path(output)) as writer:
        writer.write(entries)
<|endoftext|>


<commit_before>"""
Unit tests for plugins
"""

from django.test import TestCase

import plugin.integration
from plugin.samples.integration.sample import SampleIntegrationPlugin
from plugin.samples.integration.another_sample import WrongIntegrationPlugin, NoIntegrationPlugin
import plugin.templatetags.plugin_extras as plugin_tags
from plugin import registry, InvenTreePluginBase

class InvenTreePluginTests(TestCase):
    """ Tests for InvenTreePlugin """
    def setUp(self):
        self.plugin = InvenTreePluginBase()

        class NamedPlugin(InvenTreePluginBase):
            """a named plugin"""
            PLUGIN_NAME = 'abc123'

        self.named_plugin = NamedPlugin()

    def test_basic_plugin_init(self):
        """check if a basic plugin intis"""
        self.assertEqual(self.plugin.PLUGIN_NAME, '')
        self.assertEqual(self.plugin.plugin_name(), '')

    def test_basic_plugin_name(self):
        """check if the name of a basic plugin can be set"""
        self.assertEqual(self.named_plugin.PLUGIN_NAME, 'abc123')
        self.assertEqual(self.named_plugin.plugin_name(), 'abc123')

class PluginTagTests(TestCase):
    """ Tests for the plugin extras """

    def setUp(self):
        self.sample = SampleIntegrationPlugin()
        self.plugin_no = NoIntegrationPlugin()
        self.plugin_wrong = WrongIntegrationPlugin()

    def test_tag_plugin_list(self):
        """test that all plugins are listed"""
        self.assertEqual(plugin_tags.plugin_list(), registry.plugins)

    def test_tag_incative_plugin_list(self):
        """test that all inactive plugins are listed"""
        self.assertEqual(plugin_tags.inactive_plugin_list(), registry.plugins_inactive)

    def test_tag_plugin_settings(self):
        """check all plugins are listed"""
        self.assertEqual(
            plugin_tags.plugin_settings(self.sample),
            registry.mixins_settings.get(self.sample)
        )

    def test_tag_mixin_enabled(self):
        """check that mixin enabled functions work"""
        key = 'urls'
        # mixin enabled
        self.assertEqual(plugin_tags.mixin_enabled(self.sample, key), True)
        # mixin not enabled
        self.assertEqual(plugin_tags.mixin_enabled(self.plugin_wrong, key), False)
        # mxixn not existing
        self.assertEqual(plugin_tags.mixin_enabled(self.plugin_no, key), False)

    def test_tag_safe_url(self):
        """test that the safe url tag works expected"""
        # right url
        self.assertEqual(plugin_tags.safe_url('api-plugin-install'), '/api/plugin/install/')
        # wrong url
        self.assertEqual(plugin_tags.safe_url('indexas'), None)

    def test_tag_plugin_errors(self):
        """test that all errors are listed"""
        self.assertEqual(plugin_tags.plugin_errors(), registry.errors)
<commit_msg>PEP fix<commit_after>"""
Unit tests for plugins
"""

from django.test import TestCase

from plugin.samples.integration.sample import SampleIntegrationPlugin
from plugin.samples.integration.another_sample import WrongIntegrationPlugin, NoIntegrationPlugin
import plugin.templatetags.plugin_extras as plugin_tags
from plugin import registry, InvenTreePluginBase

class InvenTreePluginTests(TestCase):
    """ Tests for InvenTreePlugin """
    def setUp(self):
        self.plugin = InvenTreePluginBase()

        class NamedPlugin(InvenTreePluginBase):
            """a named plugin"""
            PLUGIN_NAME = 'abc123'

        self.named_plugin = NamedPlugin()

    def test_basic_plugin_init(self):
        """check if a basic plugin intis"""
        self.assertEqual(self.plugin.PLUGIN_NAME, '')
        self.assertEqual(self.plugin.plugin_name(), '')

    def test_basic_plugin_name(self):
        """check if the name of a basic plugin can be set"""
        self.assertEqual(self.named_plugin.PLUGIN_NAME, 'abc123')
        self.assertEqual(self.named_plugin.plugin_name(), 'abc123')

class PluginTagTests(TestCase):
    """ Tests for the plugin extras """

    def setUp(self):
        self.sample = SampleIntegrationPlugin()
        self.plugin_no = NoIntegrationPlugin()
        self.plugin_wrong = WrongIntegrationPlugin()

    def test_tag_plugin_list(self):
        """test that all plugins are listed"""
        self.assertEqual(plugin_tags.plugin_list(), registry.plugins)

    def test_tag_incative_plugin_list(self):
        """test that all inactive plugins are listed"""
        self.assertEqual(plugin_tags.inactive_plugin_list(), registry.plugins_inactive)

    def test_tag_plugin_settings(self):
        """check all plugins are listed"""
        self.assertEqual(
            plugin_tags.plugin_settings(self.sample),
            registry.mixins_settings.get(self.sample)
        )

    def test_tag_mixin_enabled(self):
        """check that mixin enabled functions work"""
        key = 'urls'
        # mixin enabled
        self.assertEqual(plugin_tags.mixin_enabled(self.sample, key), True)
        # mixin not enabled
        self.assertEqual(plugin_tags.mixin_enabled(self.plugin_wrong, key), False)
        # mxixn not existing
        self.assertEqual(plugin_tags.mixin_enabled(self.plugin_no, key), False)

    def test_tag_safe_url(self):
        """test that the safe url tag works expected"""
        # right url
        self.assertEqual(plugin_tags.safe_url('api-plugin-install'), '/api/plugin/install/')
        # wrong url
        self.assertEqual(plugin_tags.safe_url('indexas'), None)

    def test_tag_plugin_errors(self):
        """test that all errors are listed"""
        self.assertEqual(plugin_tags.plugin_errors(), registry.errors)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Watchmaker module."""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals, with_statement)

import codecs
import collections
import datetime
import logging
import os
import platform
import subprocess

import yaml
from six.moves import urllib

from watchmaker import static
from watchmaker.exceptions import WatchmakerException
from watchmaker.managers.workers import (LinuxWorkersManager,
                                         WindowsWorkersManager)

__version__ = '0.1.3.dev'

class Arguments(dict):
    """
    Create an arguments object for the :class:`Client`.

    Args:
        config_path (:obj:`str`):
            (Defaults to ``None``) Path or URL to the Watchmaker configuration
            file. If ``None``, the default config.yaml file is used.
        log_dir (:obj:`str`):
            (Defaults to ``None``) See :func:`logger.prepare_logging`.
        no_reboot (:obj:`bool`):
            (Defaults to ``False``): Switch to control whether to reboot the
            system upon a successfull execution of
            :func:`WatchmakerClient.install`. When this parameter is set,
            Watchmaker will suppress the reboot. Watchmaker automatically
            suppresses the reboot if it encounters an error.
        verbosity (:obj:`int`):
            (Defaults to ``0``) See :func:`prepare_logging`.

    .. important::

        For all **Keyword Arguments**, below, the default value of ``None``
        means Watchmaker will get the value from the configuration file. Be
        aware that ``None`` and ``'None'`` are two different values, with
        different meanings and effects.

    Keyword Args:

        admin_groups (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the domain
            _groups_ that should have root privileges on Linux or admin
            privileges on Windows. Value must be a colon-separated string. On
            Linux, use the ``^`` to denote spaces in the group name.

            .. code-block:: python

                admin_groups = "group1:group2"

                # (Linux only) If there are spaces in the group name, replace
                # the spaces with a '^':
                admin_groups = "space^out"

                # (Windows only) If there are spaces in the group name, no
                # special syntax required.
                admin_groups = "space out"

        admin_users (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the domain
            _users_ that should have root privileges on Linux or admin
            privileges on Windows. Value must be a colon-separated string.

            .. code-block:: python

                admin_users = "user1:user2"

        computer_name (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the
            computername to apply to the system.
        environment (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the
            environment in which the system is being built. For example:
            ``dev``, ``test``, or ``prod``.
        salt_states (:obj:`str`):
            (Defaults to ``None``) Comma-separated string of salt states to
            apply. A value of ``'None'`` (the string) will not apply any salt
            states. A value of ``'Highstate'`` will apply the salt highstate.
        s3_source (:obj:`bool`):
            (Defaults to ``None``) Use S3 utilities to retrieve content instead
            of http/s utilities. For S3 utilities to work, the system must have
            boto credentials configured that allow access to the S3 bucket.
        ou_path (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the full DN
            of the OU where the computer account will be created when joining a
            domain.

            .. code-block:: python

                ou_path="OU=Super Cool App,DC=example,DC=com"

        extra_arguments (:obj:`list`):
            (Defaults to ``[]``) A list of extra arguments to be merged into
            the worker configurations. The list must be formed as pairs of
            named arguments and values. Any leading hypens in the argument name
            are stripped. For example:

            .. code-block:: python

                extra_arguments=['--arg1', 'value1', '--arg2', 'value2']

                # This list would be converted to the following dict and merged
                # into the parameters passed to the worker configurations:
                {'arg1': 'value1', 'arg2': 'value2'}
    """

    def __init__(  # noqa: D102
        self,
        config_path=None,
        log_dir=None,
        no_reboot=False,
        verbosity=0,
        *args,
        **kwargs
    ):
        super(Arguments, self).__init__(*args, **kwargs)
        self.config_path = config_path
        self.log_dir = log_dir
        self.no_reboot = no_reboot
        self.verbosity = verbosity
        self.admin_groups = kwargs.pop('admin_groups', None)
        self.admin_users = kwargs.pop('admin_users', None)
        self.computer_name = kwargs.pop('computer_name', None)
        self.environment = kwargs.pop('environment', None)
        self.salt_states = kwargs.pop('salt_states', None)
        self.s3_source = kwargs.pop('s3_source', None)
        self.ou_path = kwargs.pop('ou_path', None)
        self.extra_arguments = kwargs.pop('extra_arguments', [])

    def __getattr__(self, attr):
        """Support attr-notation for getting dict contents."""
        return super(Arguments, self).__getitem__(attr)

    def __setattr__(self, attr, value):
        """Support attr-notation for setting dict contents."""
        super(Arguments, self).__setitem__(attr, value)

class Client(object):
    """
    Prepare a system for setup and installation.

    Keyword Arguments:
        arguments (:obj:`Arguments`):
            A dictionary of arguments. See :class:`Arguments`.
    """

    def __init__(self, arguments):  # noqa: D102
        self.log = logging.getLogger(
            '{0}.{1}'.format(__name__, self.__class__.__name__)
        )
        # Pop extra_arguments now so we can log it separately
        extra_arguments = arguments.pop('extra_arguments', [])

        header = ' WATCHMAKER RUN '
        header = header.rjust((40 + len(header) // 2), '#').ljust(80, '#')
        self.log.info(header)
        self.log.debug('Parameters: %s', arguments)
        self.log.debug('Extra Parameters: %s', extra_arguments)

        # Pop remaining arguments used by watchmaker.Client itself
        self.default_config = os.path.join(static.__path__[0], 'config.yaml')
        self.no_reboot = arguments.pop('no_reboot', False)
        self.config_path = arguments.pop('config_path')
        self.log_dir = arguments.pop('log_dir')
        self.verbosity = arguments.pop('verbosity')

        # Get the system params
        self.system = platform.system().lower()
        self._set_system_params()

        self.log.debug('System Type: %s', self.system)
        self.log.debug('System Parameters: %s', self.system_params)

        # All remaining arguments are worker_args
        worker_args = arguments

        # Convert extra_arguments to a dict and merge it with worker_args
        worker_args.update(dict(
            (k.lstrip('-'), v) for k, v in zip(*[iter(extra_arguments)] * 2)
        ))
        # Set self.worker_args, removing `None` values from worker_args
        self.worker_args = dict(
            (k, v) for k, v in worker_args.items() if v is not None
        )

        self.config = self._get_config()

    @staticmethod
    def _validate_url(url):
        return urllib.parse.urlparse(url).scheme in ['http', 'https']

    def _get_config(self):
        """
        Read and validate configuration data for installation.

        Returns:
            OrderedDict: Returns the data from the the YAML configuration file,
            scoped to the value of ``self.system`` and merged with the value of
            the ``"All"`` key.
        """
        if not self.config_path:
            self.log.warning(
                'User did not supply a config. Using the default config.'
            )
            self.config_path = self.default_config
        else:
            self.log.info('User supplied config being used.')

        # Get the raw config data
        data = ''
        if self._validate_url(self.config_path):
            try:
                data = urllib.request.urlopen(self.config_path).read()
            except urllib.error.URLError:
                msg = (
                    'The URL used to get the user config.yaml file did not '
                    'work! Please make sure your config is available.'
                )
                self.log.critical(msg)
                raise
        elif self.config_path and not os.path.exists(self.config_path):
            msg = (
                'User supplied config {0} does not exist. Please '
                'double-check your config path or use the default config '
                'path.'.format(self.config_path)
            )
            self.log.critical(msg)
            raise WatchmakerException(msg)
        else:
            with codecs.open(self.config_path, encoding="utf-8") as fh_:
                data = fh_.read()

        config_full = yaml.safe_load(data)
        try:
            config_all = config_full.get('all', [])
            config_system = config_full.get(self.system, [])
        except AttributeError:
            msg = 'Malformed config file. Must be a dictionary.'
            self.log.critical(msg)
            raise

        # If both config and config_system are empty, raise
        if not config_system and not config_all:
            msg = 'Malformed config file. No workers for this system.'
            self.log.critical(msg)
            raise WatchmakerException(msg)

        # Merge the config data, preserving the listed order of workers
        config = collections.OrderedDict()
        for worker in config_system + config_all:
            try:
                # worker is a single-key dict, where the key is the name of the
                # worker and the value is the worker parameters. we need to
                # test if the worker is already in the config, but a dict is
                # is not hashable so cannot be tested directly with
                # `if worker not in config`. this bit of ugliness extracts the
                # key and its value so we can use them directly.
                worker_name, worker_config = list(worker.items())[0]
                if worker_name not in config:
                    # Add worker to config
                    config[worker_name] = {'config': worker_config}
                    self.log.debug('%s config: %s', worker_name, worker_config)
                else:
                    # Worker present in both config_system and config_all
                    config[worker_name]['config'].update(worker_config)
                    self.log.debug(
                        '%s extra config: %s',
                        worker_name, worker_config
                    )
                    # Need to (re)merge cli worker args so they override
                    config[worker_name]['__merged'] = False
                if not config[worker_name].get('__merged'):
                    # Merge worker_args into config params
                    config[worker_name]['config'].update(self.worker_args)
                    config[worker_name]['__merged'] = True
            except Exception:
                msg = (
                    'Failed to merge worker config; worker={0}'
                    .format(worker)
                )
                self.log.critical(msg)
                raise

        self.log.debug(
            'Command-line arguments merged into worker configs: %s',
            self.worker_args
        )

        return config

    def _get_linux_system_params(self):
        """Set ``self.system_params`` attribute for Linux systems."""
        params = {}
        params['prepdir'] = os.path.join(
            '{0}'.format(self.system_drive), 'usr', 'tmp', 'watchmaker')
        params['readyfile'] = os.path.join(
            '{0}'.format(self.system_drive), 'var', 'run', 'system-is-ready')
        params['logdir'] = os.path.join(
            '{0}'.format(self.system_drive), 'var', 'log')
        params['workingdir'] = os.path.join(
            '{0}'.format(params['prepdir']), 'workingfiles')
        params['restart'] = 'shutdown -r +1 &'
        return params

    def _get_windows_system_params(self):
        """Set ``self.system_params`` attribute for Windows systems."""
        params = {}
        # os.path.join does not produce path as expected when first string
        # ends in colon; so using a join on the sep character.
        params['prepdir'] = os.path.sep.join([self.system_drive, 'Watchmaker'])
        params['readyfile'] = os.path.join(
            '{0}'.format(params['prepdir']), 'system-is-ready')
        params['logdir'] = os.path.join(
            '{0}'.format(params['prepdir']), 'Logs')
        params['workingdir'] = os.path.join(
            '{0}'.format(params['prepdir']), 'WorkingFiles')
        params['shutdown_path'] = os.path.join(
            '{0}'.format(os.environ['SYSTEMROOT']), 'system32', 'shutdown.exe')
        params['restart'] = params["shutdown_path"] + \
            ' /r /t 30 /d p:2:4 /c ' + \
            '"Watchmaker complete. Rebooting computer."'
        return params

    def _set_system_params(self):
        """Set OS-specific attributes."""
        if 'linux' in self.system:
            self.system_drive = '/'
            self.workers_manager = LinuxWorkersManager
            self.system_params = self._get_linux_system_params()
        elif 'windows' in self.system:
            self.system_drive = os.environ['SYSTEMDRIVE']
            self.workers_manager = WindowsWorkersManager
            self.system_params = self._get_windows_system_params()
        else:
            msg = 'System, {0}, is not recognized?'.format(self.system)
            self.log.critical(msg)
            raise WatchmakerException(msg)
        if self.log_dir:
            self.system_params['logdir'] = self.log_dir

    def install(self):
        """
        Initiate the installation of the prepared system.

        Upon successful execution, the system will be properly provisioned,
        according to the defined configuration and workers.
        """
        self.log.info('Start time: %s', datetime.datetime.now())
        self.log.info('Workers to execute: %s', self.config.keys())

        # Create watchmaker directories
        try:
            os.makedirs(self.system_params['workingdir'])
        except OSError:
            if not os.path.exists(self.system_params['workingdir']):
                msg = (
                    'Unable create directory - {0}'
                    .format(self.system_params['workingdir'])
                )
                self.log.critical(msg)
                raise

        workers_manager = self.workers_manager(
            system_params=self.system_params,
            workers=self.config
        )

        try:
            workers_manager.worker_cadence()
        except Exception:
            msg = 'Execution of the workers cadence has failed.'
            self.log.critical(msg)
            raise

        if self.no_reboot:
            self.log.info(
                'Detected `no-reboot` switch. System will not be '
                'rebooted.'
            )
        else:
            self.log.info(
                'Reboot scheduled. System will reboot after the script '
                'exits.'
            )
            subprocess.call(self.system_params['restart'], shell=True)
        self.log.info('Stop time: %s', datetime.datetime.now())
<commit_msg>Aligns cli syntax for extra_arguments<commit_after># -*- coding: utf-8 -*-
"""Watchmaker module."""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals, with_statement)

import codecs
import collections
import datetime
import logging
import os
import platform
import subprocess

import yaml
from six.moves import urllib

from watchmaker import static
from watchmaker.exceptions import WatchmakerException
from watchmaker.managers.workers import (LinuxWorkersManager,
                                         WindowsWorkersManager)

__version__ = '0.1.3.dev'

class Arguments(dict):
    """
    Create an arguments object for the :class:`Client`.

    Args:
        config_path (:obj:`str`):
            (Defaults to ``None``) Path or URL to the Watchmaker configuration
            file. If ``None``, the default config.yaml file is used.
        log_dir (:obj:`str`):
            (Defaults to ``None``) See :func:`logger.prepare_logging`.
        no_reboot (:obj:`bool`):
            (Defaults to ``False``): Switch to control whether to reboot the
            system upon a successfull execution of
            :func:`WatchmakerClient.install`. When this parameter is set,
            Watchmaker will suppress the reboot. Watchmaker automatically
            suppresses the reboot if it encounters an error.
        verbosity (:obj:`int`):
            (Defaults to ``0``) See :func:`prepare_logging`.

    .. important::

        For all **Keyword Arguments**, below, the default value of ``None``
        means Watchmaker will get the value from the configuration file. Be
        aware that ``None`` and ``'None'`` are two different values, with
        different meanings and effects.

    Keyword Args:

        admin_groups (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the domain
            _groups_ that should have root privileges on Linux or admin
            privileges on Windows. Value must be a colon-separated string. On
            Linux, use the ``^`` to denote spaces in the group name.

            .. code-block:: python

                admin_groups = "group1:group2"

                # (Linux only) If there are spaces in the group name, replace
                # the spaces with a '^':
                admin_groups = "space^out"

                # (Windows only) If there are spaces in the group name, no
                # special syntax required.
                admin_groups = "space out"

        admin_users (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the domain
            _users_ that should have root privileges on Linux or admin
            privileges on Windows. Value must be a colon-separated string.

            .. code-block:: python

                admin_users = "user1:user2"

        computer_name (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the
            computername to apply to the system.
        environment (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the
            environment in which the system is being built. For example:
            ``dev``, ``test``, or ``prod``.
        salt_states (:obj:`str`):
            (Defaults to ``None``) Comma-separated string of salt states to
            apply. A value of ``'None'`` (the string) will not apply any salt
            states. A value of ``'Highstate'`` will apply the salt highstate.
        s3_source (:obj:`bool`):
            (Defaults to ``None``) Use S3 utilities to retrieve content instead
            of http/s utilities. For S3 utilities to work, the system must have
            boto credentials configured that allow access to the S3 bucket.
        ou_path (:obj:`str`):
            (Defaults to ``None``) Set a salt grain that specifies the full DN
            of the OU where the computer account will be created when joining a
            domain.

            .. code-block:: python

                ou_path="OU=Super Cool App,DC=example,DC=com"

        extra_arguments (:obj:`list`):
            (Defaults to ``[]``) A list of extra arguments to be merged into
            the worker configurations. The list must be formed as pairs of
            named arguments and values. Any leading hypens in the argument name
            are stripped. For example:

            .. code-block:: python

                extra_arguments=['--arg1', 'value1', '--arg2', 'value2']

                # This list would be converted to the following dict and merged
                # into the parameters passed to the worker configurations:
                {'arg1': 'value1', 'arg2': 'value2'}
    """

    def __init__(  # noqa: D102
        self,
        config_path=None,
        log_dir=None,
        no_reboot=False,
        verbosity=0,
        *args,
        **kwargs
    ):
        super(Arguments, self).__init__(*args, **kwargs)
        self.config_path = config_path
        self.log_dir = log_dir
        self.no_reboot = no_reboot
        self.verbosity = verbosity
        self.admin_groups = kwargs.pop('admin_groups', None)
        self.admin_users = kwargs.pop('admin_users', None)
        self.computer_name = kwargs.pop('computer_name', None)
        self.environment = kwargs.pop('environment', None)
        self.salt_states = kwargs.pop('salt_states', None)
        self.s3_source = kwargs.pop('s3_source', None)
        self.ou_path = kwargs.pop('ou_path', None)
        self.extra_arguments = kwargs.pop('extra_arguments', [])

    def __getattr__(self, attr):
        """Support attr-notation for getting dict contents."""
        return super(Arguments, self).__getitem__(attr)

    def __setattr__(self, attr, value):
        """Support attr-notation for setting dict contents."""
        super(Arguments, self).__setitem__(attr, value)

class Client(object):
    """
    Prepare a system for setup and installation.

    Keyword Arguments:
        arguments (:obj:`Arguments`):
            A dictionary of arguments. See :class:`Arguments`.
    """

    def __init__(self, arguments):  # noqa: D102
        self.log = logging.getLogger(
            '{0}.{1}'.format(__name__, self.__class__.__name__)
        )
        # Pop extra_arguments now so we can log it separately
        extra_arguments = arguments.pop('extra_arguments', [])

        header = ' WATCHMAKER RUN '
        header = header.rjust((40 + len(header) // 2), '#').ljust(80, '#')
        self.log.info(header)
        self.log.debug('Parameters: %s', arguments)
        self.log.debug('Extra Parameters: %s', extra_arguments)

        # Pop remaining arguments used by watchmaker.Client itself
        self.default_config = os.path.join(static.__path__[0], 'config.yaml')
        self.no_reboot = arguments.pop('no_reboot', False)
        self.config_path = arguments.pop('config_path')
        self.log_dir = arguments.pop('log_dir')
        self.verbosity = arguments.pop('verbosity')

        # Get the system params
        self.system = platform.system().lower()
        self._set_system_params()

        self.log.debug('System Type: %s', self.system)
        self.log.debug('System Parameters: %s', self.system_params)

        # All remaining arguments are worker_args
        worker_args = arguments

        # Convert extra_arguments to a dict and merge it with worker_args.
        # Leading hypens are removed, and other hyphens are converted to
        # underscores
        worker_args.update(dict(
            (k.lstrip('-').replace('-', '_'), v) for k, v in zip(
                *[iter(extra_arguments)] * 2)
        ))
        # Set self.worker_args, removing `None` values from worker_args
        self.worker_args = dict(
            (k, v) for k, v in worker_args.items() if v is not None
        )

        self.config = self._get_config()

    @staticmethod
    def _validate_url(url):
        return urllib.parse.urlparse(url).scheme in ['http', 'https']

    def _get_config(self):
        """
        Read and validate configuration data for installation.

        Returns:
            OrderedDict: Returns the data from the the YAML configuration file,
            scoped to the value of ``self.system`` and merged with the value of
            the ``"All"`` key.
        """
        if not self.config_path:
            self.log.warning(
                'User did not supply a config. Using the default config.'
            )
            self.config_path = self.default_config
        else:
            self.log.info('User supplied config being used.')

        # Get the raw config data
        data = ''
        if self._validate_url(self.config_path):
            try:
                data = urllib.request.urlopen(self.config_path).read()
            except urllib.error.URLError:
                msg = (
                    'The URL used to get the user config.yaml file did not '
                    'work! Please make sure your config is available.'
                )
                self.log.critical(msg)
                raise
        elif self.config_path and not os.path.exists(self.config_path):
            msg = (
                'User supplied config {0} does not exist. Please '
                'double-check your config path or use the default config '
                'path.'.format(self.config_path)
            )
            self.log.critical(msg)
            raise WatchmakerException(msg)
        else:
            with codecs.open(self.config_path, encoding="utf-8") as fh_:
                data = fh_.read()

        config_full = yaml.safe_load(data)
        try:
            config_all = config_full.get('all', [])
            config_system = config_full.get(self.system, [])
        except AttributeError:
            msg = 'Malformed config file. Must be a dictionary.'
            self.log.critical(msg)
            raise

        # If both config and config_system are empty, raise
        if not config_system and not config_all:
            msg = 'Malformed config file. No workers for this system.'
            self.log.critical(msg)
            raise WatchmakerException(msg)

        # Merge the config data, preserving the listed order of workers
        config = collections.OrderedDict()
        for worker in config_system + config_all:
            try:
                # worker is a single-key dict, where the key is the name of the
                # worker and the value is the worker parameters. we need to
                # test if the worker is already in the config, but a dict is
                # is not hashable so cannot be tested directly with
                # `if worker not in config`. this bit of ugliness extracts the
                # key and its value so we can use them directly.
                worker_name, worker_config = list(worker.items())[0]
                if worker_name not in config:
                    # Add worker to config
                    config[worker_name] = {'config': worker_config}
                    self.log.debug('%s config: %s', worker_name, worker_config)
                else:
                    # Worker present in both config_system and config_all
                    config[worker_name]['config'].update(worker_config)
                    self.log.debug(
                        '%s extra config: %s',
                        worker_name, worker_config
                    )
                    # Need to (re)merge cli worker args so they override
                    config[worker_name]['__merged'] = False
                if not config[worker_name].get('__merged'):
                    # Merge worker_args into config params
                    config[worker_name]['config'].update(self.worker_args)
                    config[worker_name]['__merged'] = True
            except Exception:
                msg = (
                    'Failed to merge worker config; worker={0}'
                    .format(worker)
                )
                self.log.critical(msg)
                raise

        self.log.debug(
            'Command-line arguments merged into worker configs: %s',
            self.worker_args
        )

        return config

    def _get_linux_system_params(self):
        """Set ``self.system_params`` attribute for Linux systems."""
        params = {}
        params['prepdir'] = os.path.join(
            '{0}'.format(self.system_drive), 'usr', 'tmp', 'watchmaker')
        params['readyfile'] = os.path.join(
            '{0}'.format(self.system_drive), 'var', 'run', 'system-is-ready')
        params['logdir'] = os.path.join(
            '{0}'.format(self.system_drive), 'var', 'log')
        params['workingdir'] = os.path.join(
            '{0}'.format(params['prepdir']), 'workingfiles')
        params['restart'] = 'shutdown -r +1 &'
        return params

    def _get_windows_system_params(self):
        """Set ``self.system_params`` attribute for Windows systems."""
        params = {}
        # os.path.join does not produce path as expected when first string
        # ends in colon; so using a join on the sep character.
        params['prepdir'] = os.path.sep.join([self.system_drive, 'Watchmaker'])
        params['readyfile'] = os.path.join(
            '{0}'.format(params['prepdir']), 'system-is-ready')
        params['logdir'] = os.path.join(
            '{0}'.format(params['prepdir']), 'Logs')
        params['workingdir'] = os.path.join(
            '{0}'.format(params['prepdir']), 'WorkingFiles')
        params['shutdown_path'] = os.path.join(
            '{0}'.format(os.environ['SYSTEMROOT']), 'system32', 'shutdown.exe')
        params['restart'] = params["shutdown_path"] + \
            ' /r /t 30 /d p:2:4 /c ' + \
            '"Watchmaker complete. Rebooting computer."'
        return params

    def _set_system_params(self):
        """Set OS-specific attributes."""
        if 'linux' in self.system:
            self.system_drive = '/'
            self.workers_manager = LinuxWorkersManager
            self.system_params = self._get_linux_system_params()
        elif 'windows' in self.system:
            self.system_drive = os.environ['SYSTEMDRIVE']
            self.workers_manager = WindowsWorkersManager
            self.system_params = self._get_windows_system_params()
        else:
            msg = 'System, {0}, is not recognized?'.format(self.system)
            self.log.critical(msg)
            raise WatchmakerException(msg)
        if self.log_dir:
            self.system_params['logdir'] = self.log_dir

    def install(self):
        """
        Initiate the installation of the prepared system.

        Upon successful execution, the system will be properly provisioned,
        according to the defined configuration and workers.
        """
        self.log.info('Start time: %s', datetime.datetime.now())
        self.log.info('Workers to execute: %s', self.config.keys())

        # Create watchmaker directories
        try:
            os.makedirs(self.system_params['workingdir'])
        except OSError:
            if not os.path.exists(self.system_params['workingdir']):
                msg = (
                    'Unable create directory - {0}'
                    .format(self.system_params['workingdir'])
                )
                self.log.critical(msg)
                raise

        workers_manager = self.workers_manager(
            system_params=self.system_params,
            workers=self.config
        )

        try:
            workers_manager.worker_cadence()
        except Exception:
            msg = 'Execution of the workers cadence has failed.'
            self.log.critical(msg)
            raise

        if self.no_reboot:
            self.log.info(
                'Detected `no-reboot` switch. System will not be '
                'rebooted.'
            )
        else:
            self.log.info(
                'Reboot scheduled. System will reboot after the script '
                'exits.'
            )
            subprocess.call(self.system_params['restart'], shell=True)
        self.log.info('Stop time: %s', datetime.datetime.now())
<|endoftext|>


<commit_before>"""
Interface between loop object and ebtel++ simulation
"""

import os
import logging
import copy

import numpy as np
from scipy.interpolate import splprep,splev
import astropy.units as u

from synthesizAR.util import InputHandler,OutputHandler

class EbtelInterface(object):
    """
    Interface between field/loop model for the EBTEL model

    Parameters
    ----------
    base_config : `dict`
        Config dictionary with default parameters for all loops.
    heating_model
    """

    def __init__(self,base_config,heating_model,dt=None,ds=None):
        """
        Create EBTEL interface
        """
        self.logger = logging.getLogger(name=type(self).__name__)
        self.base_config = base_config
        self.heating_model = heating_model
        self.heating_model.base_config = base_config
        if dt is None:
            self.global_time = None
            self.logger.warning('Global time not set. Evolution of loops may not be synchronized. Set global time before importing loops to change this.')
        else:
            self.global_time = np.linspace(0.0,self.base_config['total_time'],
                int(np.ceil(self.base_config['total_time']/dt)))*u.s
        self.ds = ds
        if self.ds is None:
            self.logger.warning('Interpolated loop spacing set to None. You will not be able to load loop results until this is set.')

    def configure_input(self,loop,parent_config_dir,parent_results_dir):
        """
        Configure EBTEL input for a given loop object.

        Parameters
        ----------
        loop
        parent_config_dir : `string`
        parent_results_dir : `string`
        """
        oh = OutputHandler(os.path.join(parent_config_dir,loop.name+'.xml'), copy.deepcopy(self.base_config))
        oh.output_dict['output_filename'] = os.path.join(parent_results_dir,loop.name)
        oh.output_dict['loop_length'] = loop.full_length.value/2.0
        event_properties = self.heating_model.calculate_event_properties(loop)
        events = []
        for i in range(self.heating_model.number_events):
            events.append({'event':{
            'magnitude':event_properties['magnitude'][i], 'rise_start':event_properties['rise_start'][i], 'rise_end':event_properties['rise_end'][i], 'decay_start':event_properties['decay_start'][i], 'decay_end':event_properties['decay_end'][i]}})
        oh.output_dict['heating']['events'] = events
        oh.print_to_xml()
        oh.output_dict['config_filename'] = oh.output_filename
        loop.hydro_configuration = oh.output_dict

    def load_results(self,loop):
        """
        Load EBTEL output for a given loop object.

        Parameters
        ----------
        loop
        """
        if self.ds is not None:
            #interpolate loop lengths to higher resolution with a B-spline
            N_interp = int(np.ceil(loop.full_length/self.ds.to(loop.full_length.unit)))
            nots,_ = splprep(loop.coordinates.value.T)
            _tmp = splev(np.linspace(0,1,N_interp),nots)
            loop.coordinates = [(x,y,z) for x,y,z in zip(_tmp[0],_tmp[1],_tmp[2])]*loop.coordinates.unit

        #load in data and interpolate to universal time
        N_s = len(loop.field_aligned_coordinate)
        _tmp = np.loadtxt(loop.hydro_configuration['output_filename'])

        if self.global_time is not None:
            loop.time = self.global_time
        else:
            loop.time = _tmp[:,0]

        loop.temperature = np.outer(np.interp(loop.time, _tmp[:,0], _tmp[:,1]), np.ones(N_s))*u.K
        loop.density = np.outer(np.interp(loop.time, _tmp[:,0], _tmp[:,3]), np.ones(N_s))*(u.cm**(-3))
<commit_msg>PEP8 fixes<commit_after>"""
Interface between loop object and ebtel++ simulation
"""

import os
import logging
import copy

import numpy as np
from scipy.interpolate import splprep,splev
import astropy.units as u

from synthesizAR.util import InputHandler,OutputHandler

class EbtelInterface(object):
    """
    Interface between field/loop model for the EBTEL model

    Parameters
    ----------
    base_config : `dict`
        Config dictionary with default parameters for all loops.
    heating_model
    """

    def __init__(self,base_config,heating_model,dt=None,ds=None):
        """
        Create EBTEL interface
        """
        self.logger = logging.getLogger(name=type(self).__name__)
        self.base_config = base_config
        self.heating_model = heating_model
        self.heating_model.base_config = base_config
        if dt is None:
            self.global_time = None
            self.logger.warning('Global time not set. Evolution of loops may not be synchronized. Set global time before importing loops to change this.')
        else:
            self.global_time = np.linspace(0.0,self.base_config['total_time'],
                int(np.ceil(self.base_config['total_time']/dt)))*u.s
        self.ds = ds
        if self.ds is None:
            self.logger.warning('Interpolated loop spacing set to None. You will not be able to load loop results until this is set.')

    def configure_input(self,loop,parent_config_dir,parent_results_dir):
        """
        Configure EBTEL input for a given loop object.

        Parameters
        ----------
        loop
        parent_config_dir : `string`
        parent_results_dir : `string`
        """
        oh = OutputHandler(os.path.join(parent_config_dir,loop.name+'.xml'), copy.deepcopy(self.base_config))
        oh.output_dict['output_filename'] = os.path.join(parent_results_dir,loop.name)
        oh.output_dict['loop_length'] = loop.full_length.value/2.0
        event_properties = self.heating_model.calculate_event_properties(loop)
        events = []
        for i in range(self.heating_model.number_events):
            events.append({'event':{
            'magnitude':event_properties['magnitude'][i],
            'rise_start':event_properties['rise_start'][i],
            'rise_end':event_properties['rise_end'][i],
            'decay_start':event_properties['decay_start'][i],
            'decay_end':event_properties['decay_end'][i]}})
        oh.output_dict['heating']['events'] = events
        oh.print_to_xml()
        oh.output_dict['config_filename'] = oh.output_filename
        loop.hydro_configuration = oh.output_dict

    def load_results(self,loop):
        """
        Load EBTEL output for a given loop object.

        Parameters
        ----------
        loop
        """
        if self.ds is not None:
            #interpolate loop lengths to higher resolution with a B-spline
            N_interp = int(np.ceil(loop.full_length/self.ds.to(loop.full_length.unit)))
            nots,_ = splprep(loop.coordinates.value.T)
            _tmp = splev(np.linspace(0,1,N_interp),nots)
            loop.coordinates = [(x,y,z) for x,y,z in zip(_tmp[0],_tmp[1],_tmp[2])]*loop.coordinates.unit

        #load in data and interpolate to universal time
        N_s = len(loop.field_aligned_coordinate)
        _tmp = np.loadtxt(loop.hydro_configuration['output_filename'])

        if self.global_time is not None:
            loop.time = self.global_time
        else:
            loop.time = _tmp[:,0]

        loop.temperature = np.outer(np.interp(loop.time, _tmp[:,0], _tmp[:,1]), np.ones(N_s))*u.K
        loop.density = np.outer(np.interp(loop.time, _tmp[:,0], _tmp[:,3]), np.ones(N_s))*(u.cm**(-3))
<|endoftext|>