<commit_before>"""

$Id$
"""
__version__='$Revision$'

### JABHACKALERT: Need to remove Back/Forward buttons, because they do not
### do what the user would expect.

# CEBALERT: taggedsliders refresh the display even when auto-refresh is off.

# CB changing pattern doesn't refresh plots even when autorefresh is on.

import copy
import Pmw

import topo

import topo.base.patterngenerator
import topo.base.sheetview
from topo.base.parameterclasses import BooleanParameter,Number

from topo.sheets.generatorsheet import GeneratorSheet
from topo.commands.basic import pattern_present
from topo.plotting.plotgroup import identity

from topo.misc.keyedlist import KeyedList

import parametersframe
from plotgrouppanel import XPGPanel

from topo.plotting.plot import make_template_plot

from Tkinter import IntVar, StringVar, Checkbutton
from Tkinter import TOP, LEFT, RIGHT, BOTTOM, YES, N, S, E, W, X,NO,NONE

from tkparameterizedobject import ButtonParameter

DEFAULT_PRESENTATION = 1.0

from topo.plotting.plotgroup import XPlotGroup

class TestPatternPlotGroup(XPlotGroup):
    def __init__(self,**params):
        super(TestPatternPlotGroup,self).__init__(**params)

class TestPattern(XPGPanel):

    sheet_type = GeneratorSheet
    plotgroup_type = TestPatternPlotGroup

    learning = BooleanParameter(default=False,doc="""Whether to enable learning during presentation.""")
    duration = Number(default=1.0,doc="""How long to run the simulator when presenting.""")

    present = ButtonParameter(doc="""Present this pattern to the simulation.""")
    reset = ButtonParameter(doc="""Reset the parameters for this pattern back to their defaults.""")

    def __init__(self,console,master,label="Preview",**params):
	super(TestPattern,self).__init__(console,master,label,**params)
        self.auto_refresh=True

        # FIND GENERATOR SHEETS
        ###############################################################################################
        self.plotgroup.params()['sheet'].range=topo.sim.objects(GeneratorSheet).values()

        ### Find the GeneratorSheets in the simulation, set up generator_sheet_patterns dictionary
        # CEBALERT: this has a difficult structure to work with.
        # generator_sheets_patterns =
        # {generator_sheet_name:  { 'generator_sheet': <gs_obj>,
        #                           'editing': True/False,
        #                           'pattern_generator': <pg_obj> }
        #                         }
        self.generator_sheets_patterns = {}
        for (gen_sheet_name,gen_sheet) in topo.sim.objects(GeneratorSheet).items():
            self.generator_sheets_patterns[gen_sheet_name] = {'generator_sheet': gen_sheet,
                                                              'editing': True,
                                                              'pattern_generator': None}
        ###############################################################################################

        self.pack_param('learning')
        self.pack_param('duration')

        self.pack_param('present',on_change=self.present_pattern)
        self.pack_param('reset',on_change=self.reset_to_defaults)

        # LIST OF PATTERNGENERATORS
        ###############################################################################################
        ### Menu of PatternGenerator types
        # Take the list of PatternGenerators from the first
        # GeneratorSheet's input_generator PatternGeneratorParameter.
        # pattern_generators = {'Pattern': <PatternGenerator_obj>}
        for generator_sheet_name in self.generator_sheets_patterns.values():
            generator_sheet_params = generator_sheet_name['generator_sheet'].params()
            pattern_generators = generator_sheet_params['input_generator'].range()
            break

        self.pattern_generators = KeyedList()
        self.pattern_generators.update(pattern_generators)
        self.pattern_generators.sort()  # sorted so the pgs appear alphabetically

        self.__current_pattern_generator = GeneratorSheet.classparams()['input_generator'].default
        self.__current_pattern_generator_name = StringVar()

        for (pg_name,pg) in self.pattern_generators.items():
            if pg==type(self.__current_pattern_generator):
                self.__current_pattern_generator_name.set(pg_name)
                self.__default_pattern_generator_name = pg_name

        # PatternGenerator choice box
        self.pg_choice_box = Pmw.OptionMenu(self,
                                            command = self.__change_pattern_generator,
                                            labelpos = 'w',
                                            label_text = 'Pattern generator:',
                                            menubutton_textvariable = self.__current_pattern_generator_name,
                                            items = self.pattern_generators.keys())
        self.pg_choice_box.pack(side=TOP)
        self.balloon.bind(self.pg_choice_box,
"""Type of pattern to present.
Each type will have various parameters that can be changed.""")
        # (remember to remove audiogen)
        ###############################################################################################

        # PARAMETERSFRAME
        ###############################################################################################
        # Because this window applies changes to an object immediately
        # (unlike ParametersFrame), the button to 'reset' is
        # 'Defaults'. 'Reset' would do nothing.
        self.__params_frame = parametersframe.ParametersFrame(self,buttons_to_remove=['Apply','Close','Reset'])
        self.__params_frame.create_widgets(self.__current_pattern_generator)
        self.__params_frame.pack(side=TOP,expand=YES,fill=X)
        ###############################################################################################

        # SELECT GENERATORSHEET BUTTONS
        ###############################################################################################
        ### 'Edit patterns in' boxes
        #
        # CEBHACKALERT: Buttons seem to have strange behavior!
        # Also, will boxes be in the same order as the plots?
        self.__input_box = Pmw.RadioSelect(self, labelpos = 'w',
                                           command = self._input_change,
                                           label_text = 'Apply to pattern in:',
                                           selectmode = 'multiple')
        self.__input_box.pack(expand=NO,fill=NONE,padx=5)

        keys = copy.copy(self.generator_sheets_patterns.keys())
        keys.reverse()
        for generator_sheet_name in keys:
            self.__input_box.add(generator_sheet_name)
            self.__input_box.invoke(generator_sheet_name)
        ###############################################################################################

        self.__change_pattern_generator(self.__current_pattern_generator_name.get())
	self.refresh()

    @staticmethod
    def valid_context():
        """
        Only open if GeneratorSheets are in the Simulation.
        """
        sim = topo.sim
        if sim.objects(GeneratorSheet).items():
            return True
        else:
            return False

    def refresh(self):
        self.update_plotgroup_variables()
        self.__setup_pattern_generators()
        super(TestPattern,self).refresh()

    # METHODS FOR BUTTONS
    ###############################################################################################
    def _input_change(self,button_name, checked):
        """
        Called by the input box.  The variable self.generator_sheets_patterns records
        all input event processors, and whether they are checked or
        not.
        """
        self.generator_sheets_patterns[button_name]['editing'] = checked

        if not self.generator_sheets_patterns[button_name]['editing']:
            self.generator_sheets_patterns[button_name]['pattern_generator'] = copy.copy(self.__current_pattern_generator)
        else:
            if self.auto_refresh:
                self.__setup_pattern_generators()
                self.refresh()

    def __change_pattern_generator(self,pattern_generator_name):
        """
        Set the current PatternGenerator to the one selected and get the
        ParametersFrame to draw the relevant widgets
        """
        self.__current_pattern_generator = self.pattern_generators[pattern_generator_name]()
        self.__params_frame.create_widgets(self.__current_pattern_generator)
        if self.auto_refresh:
	    self.refresh()

    def present_pattern(self):
        """
        Move the user created patterns into the GeneratorSheets, run for
        the specified length of time, then restore the original
        patterns.
        """
	topo.sim.state_push()
        self.__setup_pattern_generators()

        input_dict = dict([(name,d['pattern_generator'])
                           for (name,d) in self.generator_sheets_patterns.items()])

        pattern_present(input_dict,self.duration,
                        learning=self.learning,overwrite_previous=False)

        self.console.auto_refresh()
	topo.sim.state_pop()

    def reset_to_defaults(self):
        """
        Reset to default presentation length, no learning, and the original PatternGenerator.
        """
        self.duration=DEFAULT_PRESENTATION
        #self.learning_button.deselect()
        self.pg_choice_box.invoke(self.__default_pattern_generator_name)
        if self.auto_refresh: self.refresh()

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

    def update_plotgroup_variables(self):
        plot_list = []
        for each in self.generator_sheets_patterns.keys():
            view_dict = {}
            k = self.generator_sheets_patterns[each]['pattern_generator']
            # xdensity and ydensity are the same for the pattern_generator
            # because they were set from a Sheet xdensity and ydensity,
            # which are guaranteed to be the same
	    density = k.xdensity

	    view_dict[each] = topo.base.sheetview.SheetView((k(),k.bounds),src_name=each)
	    channels = {'Strength':each,'Hue':None,'Confidence':None}
	    ### JCALERT! it is not good to have to pass '' here... maybe a test in plot would be better
	    plot_list.append(make_template_plot(channels,view_dict,density,None,self.normalize,name=''))

        self.plotgroup.plot_list = plot_list
	self.plotgroup.time = topo.sim.time()

    ### JABHACKALERT: This should be replaced with a PatternPresenter
    ### (from topo.commands.analysis), which will allow flexible
    ### support for making objects with different parameters in the
    ### different eyes, e.g. to test ocular dominance or disparity.

    # (there was a) buttonBox.add('Use for future learning',command = self.use_for_learning)
    def __setup_pattern_generators(self):
        """
        Make an instantiation of the current user patterns.

        Also does some other things...
        """
        # CB: remember to replace disparity flip stuff.

        self.__params_frame.set_parameters()

        for (gs_name,o_s_p) in self.generator_sheets_patterns.items():
            if o_s_p['editing']==True:

                newpattern = copy.deepcopy(self.__params_frame.parameterized_object)

                newpattern.bounds   = copy.deepcopy(o_s_p['generator_sheet'].bounds)
                newpattern.xdensity = o_s_p['generator_sheet'].xdensity
                newpattern.ydensity = o_s_p['generator_sheet'].ydensity

                o_s_p['pattern_generator'] = newpattern

<commit_msg>Continued simplifying. 'Apply to' buttons temporarily not functional.<commit_after>"""

$Id$
"""
__version__='$Revision$'

### JABHACKALERT: Need to remove Back/Forward buttons, because they do not
### do what the user would expect.

# CEBALERT: taggedsliders refresh the display even when auto-refresh is off.

# CB changing pattern doesn't refresh plots even when autorefresh is on.

import copy
import Pmw

import topo

import topo.base.patterngenerator
import topo.base.sheetview
from topo.base.parameterclasses import BooleanParameter,Number

from topo.sheets.generatorsheet import GeneratorSheet
from topo.commands.basic import pattern_present
from topo.plotting.plotgroup import identity

from topo.misc.keyedlist import KeyedList

import parametersframe
from plotgrouppanel import XPGPanel

from topo.plotting.plot import make_template_plot

from Tkinter import IntVar, StringVar, Checkbutton
from Tkinter import TOP, LEFT, RIGHT, BOTTOM, YES, N, S, E, W, X,NO,NONE

from tkparameterizedobject import ButtonParameter

DEFAULT_PRESENTATION = 1.0

from topo.plotting.plotgroup import XPlotGroup

class TestPatternPlotGroup(XPlotGroup):
    def __init__(self,**params):
        super(TestPatternPlotGroup,self).__init__(**params)

class TestPattern(XPGPanel):

    sheet_type = GeneratorSheet
    plotgroup_type = TestPatternPlotGroup

    learning = BooleanParameter(default=False,doc="""Whether to enable learning during presentation.""")
    duration = Number(default=1.0,doc="""How long to run the simulator when presenting.""")

    present = ButtonParameter(doc="""Present this pattern to the simulation.""")
    reset = ButtonParameter(doc="""Reset the parameters for this pattern back to their defaults.""")

    def __init__(self,console,master,label="Preview",**params):
	super(TestPattern,self).__init__(console,master,label,**params)
        self.auto_refresh=True

        # FIND GENERATOR SHEETS
        ###############################################################################################
        self.plotgroup.params()['sheet'].range=topo.sim.objects(GeneratorSheet).values()

        ### Find the GeneratorSheets in the simulation, set up generator_sheet_patterns dictionary
        # {generator_sheet_name:pattern_generator}
        self.generator_sheets_patterns = dict.fromkeys(topo.sim.objects(GeneratorSheet),None)

            # editing 'pattern_generator': None}
        ###############################################################################################

        self.pack_param('learning')
        self.pack_param('duration')

        self.pack_param('present',on_change=self.present_pattern)
        self.pack_param('reset',on_change=self.reset_to_defaults)

        # LIST OF PATTERNGENERATORS
        ###############################################################################################
        pattern_generators = GeneratorSheet.classparams()['input_generator'].range()

        self.pattern_generators = KeyedList()
        self.pattern_generators.update(pattern_generators)
        self.pattern_generators.sort()  # sorted so the pgs appear alphabetically

        self.__current_pattern_generator = GeneratorSheet.classparams()['input_generator'].default
        self.__current_pattern_generator_name = StringVar()

        for (pg_name,pg) in self.pattern_generators.items():
            if pg==type(self.__current_pattern_generator):
                self.__current_pattern_generator_name.set(pg_name)
                self.__default_pattern_generator_name = pg_name

        # PatternGenerator choice box
        self.pg_choice_box = Pmw.OptionMenu(self,
                                            command = self.__change_pattern_generator,
                                            labelpos = 'w',
                                            label_text = 'Pattern generator:',
                                            menubutton_textvariable = self.__current_pattern_generator_name,
                                            items = self.pattern_generators.keys())
        self.pg_choice_box.pack(side=TOP)
        self.balloon.bind(self.pg_choice_box,
"""Type of pattern to present.
Each type will have various parameters that can be changed.""")
        # (remember to remove audiogen)
        ###############################################################################################

        # PARAMETERSFRAME
        ###############################################################################################
        # Because this window applies changes to an object immediately
        # (unlike ParametersFrame), the button to 'reset' is
        # 'Defaults'. 'Reset' would do nothing.
        self.__params_frame = parametersframe.ParametersFrame(self,buttons_to_remove=['Apply','Close','Reset'])
        self.__params_frame.create_widgets(self.__current_pattern_generator)
        self.__params_frame.pack(side=TOP,expand=YES,fill=X)
        ###############################################################################################

        # SELECT GENERATORSHEET BUTTONS
        ###############################################################################################
        ### 'Edit patterns in' boxes
        #
        # CEBHACKALERT: Buttons seem to have strange behavior!
        # Also, will boxes be in the same order as the plots?
        self.__input_box = Pmw.RadioSelect(self, labelpos = 'w',
                                           command = self._input_change,
                                           label_text = 'Apply to pattern in:',
                                           selectmode = 'multiple')
        self.__input_box.pack(expand=NO,fill=NONE,padx=5)

        keys = copy.copy(self.generator_sheets_patterns.keys())
        keys.reverse()
        for generator_sheet_name in keys:
            self.__input_box.add(generator_sheet_name)
            self.__input_box.invoke(generator_sheet_name)
        ###############################################################################################

        self.__change_pattern_generator(self.__current_pattern_generator_name.get())
	self.refresh()

    @staticmethod
    def valid_context():
        """
        Only open if GeneratorSheets are in the Simulation.
        """
        sim = topo.sim
        if sim.objects(GeneratorSheet).items():
            return True
        else:
            return False

    def refresh(self):
        self.update_plotgroup_variables()
        self.__setup_pattern_generators()
        super(TestPattern,self).refresh()

    # METHODS FOR BUTTONS
    ###############################################################################################
    def _input_change(self,button_name, checked):
        """
        Called by the input box.  The variable self.generator_sheets_patterns records
        all input event processors, and whether they are checked or
        not.
        """
        #self.generator_sheets_patterns[button_name]['editing'] = checked

        #if not self.generator_sheets_patterns[button_name]['editing']:
        #    self.generator_sheets_patterns[button_name]['pattern_generator'] = copy.copy(self.__current_pattern_generator)
        #else:
        if self.auto_refresh:
            self.__setup_pattern_generators()
            self.refresh()

    def __change_pattern_generator(self,pattern_generator_name):
        """
        Set the current PatternGenerator to the one selected and get the
        ParametersFrame to draw the relevant widgets
        """
        self.__current_pattern_generator = self.pattern_generators[pattern_generator_name]()
        self.__params_frame.create_widgets(self.__current_pattern_generator)
        if self.auto_refresh:
	    self.refresh()

    def present_pattern(self):
        """
        Move the user created patterns into the GeneratorSheets, run for
        the specified length of time, then restore the original
        patterns.
        """
	topo.sim.state_push()
        self.__setup_pattern_generators()

        input_dict = dict([(name,pg)
                           for name,pg in self.generator_sheets_patterns.items()])

        pattern_present(input_dict,self.duration,
                        learning=self.learning,overwrite_previous=False)

        self.console.auto_refresh()
	topo.sim.state_pop()

    def reset_to_defaults(self):
        """
        Reset to default presentation length, no learning, and the original PatternGenerator.
        """
        self.duration=DEFAULT_PRESENTATION
        #self.learning_button.deselect()
        self.pg_choice_box.invoke(self.__default_pattern_generator_name)
        if self.auto_refresh: self.refresh()

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

    def update_plotgroup_variables(self):
        plot_list = []
        for sheetname,pg in self.generator_sheets_patterns.items():
            view_dict = {}
            # xdensity and ydensity are the same for the pattern_generator
            # because they were set from a Sheet xdensity and ydensity,
            # which are guaranteed to be the same
	    density = pg.xdensity

	    view_dict[sheetname] = topo.base.sheetview.SheetView((pg(),pg.bounds),src_name=sheetname)
	    channels = {'Strength':sheetname,'Hue':None,'Confidence':None}
	    ### JCALERT! it is not good to have to pass '' here... maybe a test in plot would be better
	    plot_list.append(make_template_plot(channels,view_dict,density,None,self.normalize,name=''))

        self.plotgroup.plot_list = plot_list
	self.plotgroup.time = topo.sim.time()

    ### JABHACKALERT: This should be replaced with a PatternPresenter
    ### (from topo.commands.analysis), which will allow flexible
    ### support for making objects with different parameters in the
    ### different eyes, e.g. to test ocular dominance or disparity.

    # (there was a) buttonBox.add('Use for future learning',command = self.use_for_learning)
    def __setup_pattern_generators(self):
        """
        Make an instantiation of the current user patterns.

        Also does some other things...
        """
        # CB: remember to replace disparity flip stuff.

        self.__params_frame.set_parameters()

        for sheetname in self.generator_sheets_patterns:
            #if o_s_p['editing']==True:

            newpattern = copy.deepcopy(self.__params_frame.parameterized_object)

            sheet = topo.sim[sheetname]

            newpattern.bounds   = copy.deepcopy(sheet.bounds)
            newpattern.xdensity = sheet.xdensity
            newpattern.ydensity = sheet.ydensity

            self.generator_sheets_patterns[sheetname] = newpattern

<|endoftext|>


<commit_before>from django import template

register = template.Library()

@register.filter
def get_available_modules_for_case_list_configuration(app, module):
    return [
        m for m in app.get_modules()
        if m.id != module.id
           and m.module_type not in ('reports', 'careplan')
           and m.case_type == module.case_type
    ]
<commit_msg>[ci skip] lint<commit_after>from django import template

register = template.Library()

@register.filter
def get_available_modules_for_case_list_configuration(app, module):
    return [
        m for m in app.get_modules()
        if (m.id != module.id
            and m.module_type not in ('reports', 'careplan')
            and m.case_type == module.case_type)
    ]
<|endoftext|>


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

import logging
import os.path
import sys
import unicodedata
from collections import defaultdict

import pkg_resources
import regex as re

import panphon
import unicodecsv as csv
from epitran.exceptions import DatafileError, MappingError
from epitran.ligaturize import ligaturize
from epitran.ppprocessor import PrePostProcessor
from epitran.puncnorm import PuncNorm
from epitran.stripdiacritics import StripDiacritics

if sys.version_info[0] == 3:
    def unicode(x):
        return x

class SimpleEpitran(object):
    def __init__(self, code, preproc=True, postproc=True, ligatures=False,
                 rev=False, rev_preproc=True, rev_postproc=True):
        """Constructs the backend object epitran uses for most languages

        Args:
            code (str): ISO 639-3 code and ISO 15924 code joined with a hyphen
            preproc (bool): if True, apply preprocessor
            postproc (bool): if True, apply postprocessors
            ligatures (bool): if True, use phonetic ligatures for affricates
                              instead of standard IPA
            rev (bool): if True, load reverse transliteration
            rev_preproc (bool): if True, apply preprocessor when reverse transliterating
            rev_postproc (bool): if True, apply postprocessor when reverse transliterating
        """
        self.rev = rev
        self.g2p = self._load_g2p_map(code, False)
        self.regexp = self._construct_regex(self.g2p.keys())
        self.puncnorm = PuncNorm()
        self.ft = panphon.FeatureTable()
        self.num_panphon_fts = len(self.ft.names)
        self.preprocessor = PrePostProcessor(code, 'pre', False)
        self.postprocessor = PrePostProcessor(code, 'post', False)
        self.strip_diacritics = StripDiacritics(code)
        self.preproc = preproc
        self.postproc = postproc
        self.ligatures = ligatures
        self.rev_preproc = rev_preproc
        self.rev_postproc = rev_postproc
        if rev:
            self.rev_g2p = self._load_g2p_map(code, True)
            self.rev_regexp = self._construct_regex(self.rev_g2p.keys())
            self.rev_preprocessor = PrePostProcessor(code, 'pre', True)
            self.rev_postprocessor = PrePostProcessor(code, 'post', True)

        self.nils = defaultdict(int)

    def __enter__(self):
        return self

    def __exit__(self, type_, val, tb):
        for nil, count in self.nils.items():
            sys.stderr.write('Unknown character "{}" occured {} times.\n'.format(nil, count))

    def _one_to_many_gr_by_line_map(self, gr_by_line):
        for g, ls in gr_by_line.items():
            if len(ls) != 1:
                return (g, ls)
        return None

    def _load_g2p_map(self, code, rev):
        """Load the code table for the specified language.

        Args:
            code (str): ISO 639-3 code plus "-" plus ISO 15924 code for the
                        language/script to be loaded
            rev (boolean): True for reversing the table (for reverse transliterating)
        """
        g2p = defaultdict(list)
        gr_by_line = defaultdict(list)
        code += '_rev' if rev else ''
        try:
            path = os.path.join('data', 'map', code + '.csv')
            path = pkg_resources.resource_filename(__name__, path)
        except IndexError:
            raise DatafileError('Add an appropriately-named mapping to the data/maps directory.')
        with open(path, 'rb') as f:
            reader = csv.reader(f, encoding='utf-8')
            orth, phon = next(reader)
            if orth != 'Orth' or phon != 'Phon':
                raise DatafileError('Header is ["{}", "{}"] instead of ["Orth", "Phon"].'.format(orth, phon))
            for (i, fields) in enumerate(reader):
                try:
                    graph, phon = fields
                except ValueError:
                    raise DatafileError('Map file is not well formed at line {}.'.format(i + 2))
                graph = unicodedata.normalize('NFD', graph)
                phon = unicodedata.normalize('NFD', phon)
                g2p[graph].append(phon)
                gr_by_line[graph].append(i)
        if self._one_to_many_gr_by_line_map(g2p):
            graph, lines = self._one_to_many_gr_by_line_map(gr_by_line)
            lines = [l + 2 for l in lines]
            raise MappingError('One-to-many G2P mapping for "{}" on lines {}'.format(graph, ', '.join(map(str, lines))).encode('utf-8'))
        return g2p

    def _load_punc_norm_map(self):
        """Load the map table for normalizing 'down' punctuation."""
        path = os.path.join('data', 'puncnorm.csv')
        path = pkg_resources.resource_filename(__name__, path)
        with open(path, 'rb') as f:
            reader = csv.reader(f, encoding='utf-8',
                                delimiter=str(','),
                                quotechar=str('"'))
            next(reader)
            return {punc: norm for (punc, norm) in reader}

    def _construct_regex(self, g2p_keys):
        """Build a regular expression that will greadily match segments from
           the mapping table.
        """
        graphemes = sorted(g2p_keys, key=len, reverse=True)
        return re.compile(r'({})'.format(r'|'.join(graphemes)), re.I)

    def general_trans(self, text, filter_func,
                      normpunc=False, ligatures=False):
        """Transliaterates a word into IPA, filtering with filter_func

        Args:
            text (str): word to transcribe; unicode strings
            filter_func (function): function for filtering segments; takes
                                    a <segment, is_ipa> tuple and returns a
                                    boolean.
            normpunc (bool): normalize punctuation
            ligatures (bool): use precomposed ligatures instead of
                              standard IPA

        Returns:
            unicode: IPA string, filtered by filter_func.
        """
        text = unicode(text)
        text = unicodedata.normalize('NFD', text.lower())
        logging.debug('(after norm) text=' + repr(list(text)))
        text = self.strip_diacritics.process(text)
        logging.debug('(after strip) text=' + repr(list(text)))
        if self.preproc:
            text = self.preprocessor.process(text)
        logging.debug('(after preproc) text=' + repr(list(text)))
        tr_list = []
        while text:
            logging.debug('text=' + repr(list(text)))
            m = self.regexp.match(text)
            if m:
                source = m.group(0)
                try:
                    target = self.g2p[source][0]
                except KeyError:
                    logging.debug("source = '{}'".format(source))
                    logging.debug("self.g2p[source] = '{}'"
                                  .format(self.g2p[source]))
                    target = source
                except IndexError:
                    logging.debug("self.g2p[source]={}".format(self.g2p[source]))
                    target = source
                tr_list.append((target, True))
                text = text[len(source):]
            else:
                tr_list.append((text[0], False))
                self.nils[text[0]] += 2
                text = text[1:]
        text = ''.join([s for (s, _) in filter(filter_func, tr_list)])
        if self.postproc:
            text = self.postprocessor.process(text)
        if ligatures or self.ligatures:
            text = ligaturize(text)
        if normpunc:
            text = self.puncnorm.norm(text)
        return unicodedata.normalize('NFC', text)

    def transliterate(self, text, normpunc=False, ligatures=False):
        """Transliterates/transcribes a word into IPA

        Passes unmapped characters through to output unchanged.

        Args:
            word (str): word to transcribe; unicode string
            normpunc (bool): normalize punctuation
            ligatures (bool): use precomposed ligatures instead of standard IPA

        Returns:
            unicode: IPA string with unrecognized characters included
        """
        return self.general_trans(text, lambda x: True,
                                  normpunc, ligatures)

    def general_reverse_trans(self, ipa):
        """Reconstructs word from IPA. Does the reverse of transliterate().
        Ignores unmapped characters.

            Args:
                ipa (str): word transcription in ipa; unicode string

            Returns:
                unicode: reconstructed word
        """
        text = unicode(ipa)
        if self.rev_preproc:
            text = self.rev_preprocessor.process(text)
        tr_list = []
        while text:
            m = self.rev_regexp.match(text)
            if m:
                source = m.group(0)
                try:
                    target = self.rev_g2p[source][0]
                except KeyError:
                    logging.debug("source = '{}'".format(source))
                    logging.debug("self.rev_g2p[source] = '{}'"
                                  .format(self.g2p[source]))
                    target = source
                tr_list.append((target, True))
                text = text[len(source):]
            else:
                tr_list.append((text[0], False))
                self.nils[text[0]] += 2
                text = text[1:]
        text = ''.join([s for (s, _) in tr_list])
        if self.rev_postproc:
            text = self.rev_postprocessor.process(text)
        return unicodedata.normalize('NFC', text)

    def reverse_transliterate(self, ipa):
        """Reconstructs word from IPA. Does the reverse of transliterate()

        Args:
            ipa (str): word transcription in ipa; unicode string

        Returns:
            unicode: reconstructed word
        """
        if not self.rev:
            raise ValueError('This Epitran object was initialized with no reverse transliteration loaded')
        return self.general_reverse_trans(ipa)

    def strict_trans(self, text, normpunc=False, ligatures=False):
        """Transliterates/transcribes a word into IPA

        Ignores umapped characters.

        Args:
            word (str): word to transcribe; unicode string
            normpunc (bool): normalize punctuation
            ligatures (bool): use precomposed ligatures instead of standard IPA

        Returns:
            unicode: IPA string
        """
        return self.general_trans(text, lambda x: x[1],
                                  normpunc, ligatures)

    def word_to_tuples(self, word, normpunc=False):
        """Given a word, returns a list of tuples corresponding to IPA segments.

        Args:
            word (unicode): word to transliterate
            normpunc (bool): If True, normalizes punctuation to ASCII inventory

        Returns:
            list: A list of (category, lettercase, orthographic_form,
                  phonetic_form, feature_vectors) tuples.

        The "feature vectors" form a list consisting of (segment, vector)
        pairs. For IPA segments, segment is a substring of phonetic_form such
        that the concatenation of all segments in the list is equal to
        the phonetic_form. The vectors are a sequence of integers drawn from
        the set {-1, 0, 1} where -1 corresponds to '-', 0 corresponds to '0',
        and 1 corresponds to '+'.
        """
        def cat_and_cap(c):
            cat, case = tuple(unicodedata.category(c))
            case = 1 if case == 'u' else 0
            return unicode(cat), case

        def recode_ft(ft):
            try:
                return {'+': 1, '0': 0, '-': -1}[ft]
            except KeyError:
                return None

        def vec2bin(vec):
            return list(map(recode_ft, vec))

        def to_vector(seg):
            return seg, vec2bin(self.ft.segment_to_vector(seg))

        def to_vectors(phon):
            if phon == '':
                return [(-1, [0] * self.num_panphon_fts)]
            else:
                return [to_vector(seg) for seg in self.ft.ipa_segs(phon)]

        tuples = []
        word = unicode(word)
        word = self.strip_diacritics.process(word)
        word = unicodedata.normalize('NFD', word)
        if self.preproc:
            word = self.preprocessor.process(word)
        while word:
            match = self.regexp.match(word)
            if match:
                span = match.group(1)
                cat, case = cat_and_cap(span[0])
                phon = self.g2p[span.lower()][0]
                vecs = to_vectors(phon)
                tuples.append(('L', case, span, phon, vecs))
                word = word[len(span):]
            else:
                span = word[0]
                span = self.puncnorm.norm(span) if normpunc else span
                cat, case = cat_and_cap(span)
                cat = 'P' if normpunc and cat in self.puncnorm else cat
                phon = ''
                vecs = to_vectors(phon)
                tuples.append((cat, case, span, phon, vecs))
                word = word[1:]
        return tuples

    def ipa_segs(self, ipa):
        """Given an IPA string, decompose it into a list of segments

        Args:
            ipa (unicode): a Unicode IPA string

        Returns:
            list: a list of unicode strings corresponding to segments
                  (consonants and vowels) in the input string
        """
        return self.ft.ipa_segs(ipa)
<commit_msg>Enabled tones<commit_after># -*- coding: utf-8 -*-
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import logging
import os.path
import sys
import unicodedata
from collections import defaultdict

import pkg_resources
import regex as re

import panphon
import unicodecsv as csv
from epitran.exceptions import DatafileError, MappingError
from epitran.ligaturize import ligaturize
from epitran.ppprocessor import PrePostProcessor
from epitran.puncnorm import PuncNorm
from epitran.stripdiacritics import StripDiacritics

if sys.version_info[0] == 3:
    def unicode(x):
        return x

class SimpleEpitran(object):
    def __init__(self, code, preproc=True, postproc=True, ligatures=False,
                 rev=False, rev_preproc=True, rev_postproc=True, tones=False):
        """Constructs the backend object epitran uses for most languages

        Args:
            code (str): ISO 639-3 code and ISO 15924 code joined with a hyphen
            preproc (bool): if True, apply preprocessor
            postproc (bool): if True, apply postprocessors
            ligatures (bool): if True, use phonetic ligatures for affricates
                              instead of standard IPA
            rev (bool): if True, load reverse transliteration
            rev_preproc (bool): if True, apply preprocessor when reverse transliterating
            rev_postproc (bool): if True, apply postprocessor when reverse transliterating
        """
        self.rev = rev
        self.tones = tones
        self.g2p = self._load_g2p_map(code, False)
        self.regexp = self._construct_regex(self.g2p.keys())
        self.puncnorm = PuncNorm()
        self.ft = panphon.FeatureTable()
        self.num_panphon_fts = len(self.ft.names)
        self.preprocessor = PrePostProcessor(code, 'pre', False)
        self.postprocessor = PrePostProcessor(code, 'post', False)
        self.strip_diacritics = StripDiacritics(code)
        self.preproc = preproc
        self.postproc = postproc
        self.ligatures = ligatures
        self.rev_preproc = rev_preproc
        self.rev_postproc = rev_postproc
        if rev:
            self.rev_g2p = self._load_g2p_map(code, True)
            self.rev_regexp = self._construct_regex(self.rev_g2p.keys())
            self.rev_preprocessor = PrePostProcessor(code, 'pre', True)
            self.rev_postprocessor = PrePostProcessor(code, 'post', True)

        self.nils = defaultdict(int)

    def get_tones(self):
        return self.tones

    def __enter__(self):
        return self

    def __exit__(self, type_, val, tb):
        for nil, count in self.nils.items():
            sys.stderr.write('Unknown character "{}" occured {} times.\n'.format(nil, count))

    def _one_to_many_gr_by_line_map(self, gr_by_line):
        for g, ls in gr_by_line.items():
            if len(ls) != 1:
                return (g, ls)
        return None

    def _load_g2p_map(self, code, rev):
        """Load the code table for the specified language.

        Args:
            code (str): ISO 639-3 code plus "-" plus ISO 15924 code for the
                        language/script to be loaded
            rev (boolean): True for reversing the table (for reverse transliterating)
        """
        g2p = defaultdict(list)
        gr_by_line = defaultdict(list)
        code += '_rev' if rev else ''
        try:
            path = os.path.join('data', 'map', code + '.csv')
            path = pkg_resources.resource_filename(__name__, path)
        except IndexError:
            raise DatafileError('Add an appropriately-named mapping to the data/maps directory.')
        with open(path, 'rb') as f:
            reader = csv.reader(f, encoding='utf-8')
            orth, phon = next(reader)
            if orth != 'Orth' or phon != 'Phon':
                raise DatafileError('Header is ["{}", "{}"] instead of ["Orth", "Phon"].'.format(orth, phon))
            for (i, fields) in enumerate(reader):
                try:
                    graph, phon = fields
                except ValueError:
                    raise DatafileError('Map file is not well formed at line {}.'.format(i + 2))
                graph = unicodedata.normalize('NFD', graph)
                phon = unicodedata.normalize('NFD', phon)
                if not self.tones:
                    phon = re.sub('[˩˨˧˦˥]', '', phon)
                g2p[graph].append(phon)
                gr_by_line[graph].append(i)
        if self._one_to_many_gr_by_line_map(g2p):
            graph, lines = self._one_to_many_gr_by_line_map(gr_by_line)
            lines = [l + 2 for l in lines]
            raise MappingError('One-to-many G2P mapping for "{}" on lines {}'.format(graph, ', '.join(map(str, lines))).encode('utf-8'))
        return g2p

    def _load_punc_norm_map(self):
        """Load the map table for normalizing 'down' punctuation."""
        path = os.path.join('data', 'puncnorm.csv')
        path = pkg_resources.resource_filename(__name__, path)
        with open(path, 'rb') as f:
            reader = csv.reader(f, encoding='utf-8',
                                delimiter=str(','),
                                quotechar=str('"'))
            next(reader)
            return {punc: norm for (punc, norm) in reader}

    def _construct_regex(self, g2p_keys):
        """Build a regular expression that will greadily match segments from
           the mapping table.
        """
        graphemes = sorted(g2p_keys, key=len, reverse=True)
        return re.compile(r'({})'.format(r'|'.join(graphemes)), re.I)

    def general_trans(self, text, filter_func,
                      normpunc=False, ligatures=False):
        """Transliaterates a word into IPA, filtering with filter_func

        Args:
            text (str): word to transcribe; unicode strings
            filter_func (function): function for filtering segments; takes
                                    a <segment, is_ipa> tuple and returns a
                                    boolean.
            normpunc (bool): normalize punctuation
            ligatures (bool): use precomposed ligatures instead of
                              standard IPA

        Returns:
            unicode: IPA string, filtered by filter_func.
        """
        text = unicode(text)
        text = unicodedata.normalize('NFD', text.lower())
        logging.debug('(after norm) text=' + repr(list(text)))
        text = self.strip_diacritics.process(text)
        logging.debug('(after strip) text=' + repr(list(text)))
        if self.preproc:
            text = self.preprocessor.process(text)
        logging.debug('(after preproc) text=' + repr(list(text)))
        tr_list = []
        while text:
            logging.debug('text=' + repr(list(text)))
            m = self.regexp.match(text)
            if m:
                source = m.group(0)
                try:
                    target = self.g2p[source][0]
                except KeyError:
                    logging.debug("source = '{}'".format(source))
                    logging.debug("self.g2p[source] = '{}'"
                                  .format(self.g2p[source]))
                    target = source
                except IndexError:
                    logging.debug("self.g2p[source]={}".format(self.g2p[source]))
                    target = source
                tr_list.append((target, True))
                text = text[len(source):]
            else:
                tr_list.append((text[0], False))
                self.nils[text[0]] += 2
                text = text[1:]
        text = ''.join([s for (s, _) in filter(filter_func, tr_list)])
        if self.postproc:
            text = self.postprocessor.process(text)
        if ligatures or self.ligatures:
            text = ligaturize(text)
        if normpunc:
            text = self.puncnorm.norm(text)
        return unicodedata.normalize('NFC', text)

    def transliterate(self, text, normpunc=False, ligatures=False):
        """Transliterates/transcribes a word into IPA

        Passes unmapped characters through to output unchanged.

        Args:
            word (str): word to transcribe; unicode string
            normpunc (bool): normalize punctuation
            ligatures (bool): use precomposed ligatures instead of standard IPA

        Returns:
            unicode: IPA string with unrecognized characters included
        """
        return self.general_trans(text, lambda x: True,
                                  normpunc, ligatures)

    def general_reverse_trans(self, ipa):
        """Reconstructs word from IPA. Does the reverse of transliterate().
        Ignores unmapped characters.

            Args:
                ipa (str): word transcription in ipa; unicode string

            Returns:
                unicode: reconstructed word
        """
        text = unicode(ipa)
        if self.rev_preproc:
            text = self.rev_preprocessor.process(text)
        tr_list = []
        while text:
            m = self.rev_regexp.match(text)
            if m:
                source = m.group(0)
                try:
                    target = self.rev_g2p[source][0]
                except KeyError:
                    logging.debug("source = '{}'".format(source))
                    logging.debug("self.rev_g2p[source] = '{}'"
                                  .format(self.g2p[source]))
                    target = source
                tr_list.append((target, True))
                text = text[len(source):]
            else:
                tr_list.append((text[0], False))
                self.nils[text[0]] += 2
                text = text[1:]
        text = ''.join([s for (s, _) in tr_list])
        if self.rev_postproc:
            text = self.rev_postprocessor.process(text)
        return unicodedata.normalize('NFC', text)

    def reverse_transliterate(self, ipa):
        """Reconstructs word from IPA. Does the reverse of transliterate()

        Args:
            ipa (str): word transcription in ipa; unicode string

        Returns:
            unicode: reconstructed word
        """
        if not self.rev:
            raise ValueError('This Epitran object was initialized with no reverse transliteration loaded')
        return self.general_reverse_trans(ipa)

    def strict_trans(self, text, normpunc=False, ligatures=False):
        """Transliterates/transcribes a word into IPA

        Ignores umapped characters.

        Args:
            word (str): word to transcribe; unicode string
            normpunc (bool): normalize punctuation
            ligatures (bool): use precomposed ligatures instead of standard IPA

        Returns:
            unicode: IPA string
        """
        return self.general_trans(text, lambda x: x[1],
                                  normpunc, ligatures)

    def word_to_tuples(self, word, normpunc=False):
        """Given a word, returns a list of tuples corresponding to IPA segments.

        Args:
            word (unicode): word to transliterate
            normpunc (bool): If True, normalizes punctuation to ASCII inventory

        Returns:
            list: A list of (category, lettercase, orthographic_form,
                  phonetic_form, feature_vectors) tuples.

        The "feature vectors" form a list consisting of (segment, vector)
        pairs. For IPA segments, segment is a substring of phonetic_form such
        that the concatenation of all segments in the list is equal to
        the phonetic_form. The vectors are a sequence of integers drawn from
        the set {-1, 0, 1} where -1 corresponds to '-', 0 corresponds to '0',
        and 1 corresponds to '+'.
        """
        def cat_and_cap(c):
            cat, case = tuple(unicodedata.category(c))
            case = 1 if case == 'u' else 0
            return unicode(cat), case

        def recode_ft(ft):
            try:
                return {'+': 1, '0': 0, '-': -1}[ft]
            except KeyError:
                return None

        def vec2bin(vec):
            return list(map(recode_ft, vec))

        def to_vector(seg):
            return seg, vec2bin(self.ft.segment_to_vector(seg))

        def to_vectors(phon):
            if phon == '':
                return [(-1, [0] * self.num_panphon_fts)]
            else:
                return [to_vector(seg) for seg in self.ft.ipa_segs(phon)]

        tuples = []
        word = unicode(word)
        word = self.strip_diacritics.process(word)
        word = unicodedata.normalize('NFD', word)
        if self.preproc:
            word = self.preprocessor.process(word)
        while word:
            match = self.regexp.match(word)
            if match:
                span = match.group(1)
                cat, case = cat_and_cap(span[0])
                phon = self.g2p[span.lower()][0]
                vecs = to_vectors(phon)
                tuples.append(('L', case, span, phon, vecs))
                word = word[len(span):]
            else:
                span = word[0]
                span = self.puncnorm.norm(span) if normpunc else span
                cat, case = cat_and_cap(span)
                cat = 'P' if normpunc and cat in self.puncnorm else cat
                phon = ''
                vecs = to_vectors(phon)
                tuples.append((cat, case, span, phon, vecs))
                word = word[1:]
        return tuples

    def ipa_segs(self, ipa):
        """Given an IPA string, decompose it into a list of segments

        Args:
            ipa (unicode): a Unicode IPA string

        Returns:
            list: a list of unicode strings corresponding to segments
                  (consonants and vowels) in the input string
        """
        return self.ft.ipa_segs(ipa)
<|endoftext|>


<commit_before>"""
 SlipStream Client
 =====
 Copyright (C) 2014 SixSq Sarl (sixsq.com)
 =====
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import traceback

from slipstream import util
from slipstream.SlipStreamHttpClient import SlipStreamHttpClient
from slipstream.NodeDecorator import NodeDecorator
from slipstream.exceptions import Exceptions

class NodeInfoPublisher(SlipStreamHttpClient):
    def __init__(self, configHolder):
        super(NodeInfoPublisher, self).__init__(configHolder)
        self.set_http_max_retries(-5)

    def publish(self, nodename, vm_id, vm_ip):
        self.publish_instanceid(nodename, vm_id)
        self.publish_hostname(nodename, vm_ip)

    def publish_instanceid(self, nodename, vm_id):
        self._setRuntimeParameter(nodename, 'instanceid', vm_id)

    def publish_hostname(self, nodename, vm_ip):
        self._setRuntimeParameter(nodename, 'hostname', vm_ip)

    def publish_url_ssh(self, nodename, vm_ip, username):
        url = 'ssh://%s@%s' % (username.strip(), vm_ip.strip())
        self._setRuntimeParameter(nodename, 'url.ssh', url)

    def _setRuntimeParameter(self, nodename, key, value):
        parameter = nodename + NodeDecorator.NODE_PROPERTY_SEPARATOR + key
        self.setRuntimeParameter(parameter, value, ignoreAbort=True)

class BaseWrapper(object):

    SCALE_STATE_CREATING = 'creating'
    SCALE_STATE_CREATED = 'created'

    SCALE_STATE_DISK_RESIZING = 'disk_resizing'
    SCALE_STATE_DISK_RESIZED = 'disk_resized'

    SCALE_STATE_OPERATIONAL = 'operational'

    SCALE_STATE_REMOVING = 'removing'
    SCALE_STATE_REMOVED = 'removed'
    SCALE_STATE_GONE = 'gone'

    SCALE_STATES_TERMINAL = (SCALE_STATE_OPERATIONAL, SCALE_STATE_GONE)

    SCALE_ACTION_CREATION = 'creation'
    SCALE_ACTION_REMOVAL = 'removal'
    SCALE_ACTION_DISK_RESIZE = 'disk_resize'

    STATE_TO_ACTION = {SCALE_STATE_CREATED: SCALE_ACTION_CREATION,
                       SCALE_STATE_REMOVED: SCALE_ACTION_REMOVAL,
                       SCALE_STATE_DISK_RESIZED: SCALE_ACTION_DISK_RESIZE}

    def __init__(self, config_holder):
        self._ss_client = SlipStreamHttpClient(config_holder)
        self._ss_client.set_http_max_retries(self.is_mutable() and -5 or 5)
        self._ss_client.ignoreAbort = True

        self.my_node_instance_name = self._get_my_node_instance_name(config_holder)

        self._user_info = None
        self._run_parameters = None
        self._nodes_instances = {}

    def _get_my_node_instance_name(self, config_holder):
        try:
            return config_holder.node_instance_name
        except Exception:
            raise Exceptions.ExecutionException('Failed to get the node instance name of the the current VM')

    def get_my_node_instance_name(self):
        return self.my_node_instance_name

    def get_slipstream_client(self):
        return self._ss_client

    def complete_state(self, node_instance_name=None):
        if not node_instance_name:
            node_instance_name = self.get_my_node_instance_name()
        self._ss_client.complete_state(node_instance_name)

    # TODO: LS: Can we remove this method ?
    def reset(self):
        self._ss_client.reset()

    def fail(self, message):
        key = self._qualifyKey(NodeDecorator.ABORT_KEY)
        self._fail(key, message)

    def fail_global(self, message):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.ABORT_KEY
        self._fail(key, message)

    def _fail(self, key, message):
        util.printError('Failing... %s' % message)
        traceback.print_exc()
        self._ss_client.setRuntimeParameter(key, message)

    def getState(self):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.STATE_KEY
        return self._get_runtime_parameter(key)

    def get_recovery_mode(self):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.RECOVERY_MODE_KEY
        return util.str2bool(self._get_runtime_parameter(key))

    def isAbort(self):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.ABORT_KEY
        try:
            value = self._get_runtime_parameter(key, True)
        except Exceptions.NotYetSetException:
            value = ''
        return (value and True) or False

    def get_run_category(self):
        return self._ss_client.get_run_category()

    def get_run_type(self):
        return self._ss_client.get_run_type()

    def _qualifyKey(self, key):
        """Qualify the key, if not already done, with the right nodename"""

        _key = key

        # Is this a reserved or special nodename?
        for reserved in NodeDecorator.reservedNodeNames:
            if _key.startswith(reserved + NodeDecorator.NODE_PROPERTY_SEPARATOR):
                return _key

        # Is the key namespaced (i.e. contains node/key separator: ':')?
        if NodeDecorator.NODE_PROPERTY_SEPARATOR in _key:
            # Is the nodename in the form: <nodename>.<index>?  If not, make it so
            # such that <nodename>:<property> -> <nodename>.1:<property
            parts = _key.split(NodeDecorator.NODE_PROPERTY_SEPARATOR)
            nodenamePart = parts[0]
            propertyPart = parts[1]  # safe since we've done the test in the if above
            parts = nodenamePart.split(NodeDecorator.NODE_MULTIPLICITY_SEPARATOR)
            nodename = parts[0]
            if len(parts) == 1:
                _key = nodename + \
                    NodeDecorator.NODE_MULTIPLICITY_SEPARATOR + \
                    NodeDecorator.nodeMultiplicityStartIndex + \
                    NodeDecorator.NODE_PROPERTY_SEPARATOR + \
                    propertyPart
            return _key

        _key = self.get_my_node_instance_name() + NodeDecorator.NODE_PROPERTY_SEPARATOR + _key

        return _key

    def get_cloud_instance_id(self):
        key = self._qualifyKey(NodeDecorator.INSTANCEID_KEY)
        return self._get_runtime_parameter(key)

    def get_user_ssh_pubkey(self):
        userInfo = self._get_user_info('')
        return userInfo.get_public_keys()

    def _get_runtime_parameter(self, key, ignore_abort=False):
        return self._ss_client.getRuntimeParameter(key, ignoreAbort=ignore_abort)

    def need_to_stop_images(self, ignore_on_success_run_forever=False):
        # pylint: disable=unused-argument
        return False

    def is_mutable(self):
        mutable = self._ss_client.get_run_mutable()
        return util.str2bool(mutable)

    def set_scale_state_on_node_instances(self, instance_names, scale_state):
        for instance_name in instance_names:
            key = instance_name + NodeDecorator.NODE_PROPERTY_SEPARATOR + NodeDecorator.SCALE_STATE_KEY
            self._ss_client.setRuntimeParameter(key, scale_state)

    def is_scale_state_operational(self):
        return self.SCALE_STATE_OPERATIONAL == self.get_scale_state()

    def is_scale_state_creating(self):
        return self.SCALE_STATE_CREATING == self.get_scale_state()

    def _is_scale_state_terminal(self, state):
        return state in self.SCALE_STATES_TERMINAL

    def set_scale_state_operational(self):
        self.set_scale_state(self.SCALE_STATE_OPERATIONAL)

    def set_scale_state_created(self):
        self.set_scale_state(self.SCALE_STATE_CREATED)

    def set_scale_state(self, scale_state):
        '''Set scale state for this node instances.
        '''
        key = self._qualifyKey(NodeDecorator.SCALE_STATE_KEY)
        self._ss_client.setRuntimeParameter(key, scale_state)

    def get_scale_state(self):
        '''Get scale state for this node instances.
        '''
        key = self._qualifyKey(NodeDecorator.SCALE_STATE_KEY)
        return self._get_runtime_parameter(key)

    def _get_effective_scale_state(self, node_instance_name):
        '''Get effective node instance scale state from the server.
        '''
        key = node_instance_name + NodeDecorator.NODE_PROPERTY_SEPARATOR + \
            NodeDecorator.SCALE_STATE_KEY
        return self._get_runtime_parameter(key)

    def _get_effective_scale_states(self):
        '''Extract node instances in scaling states and update their effective
        states from the server.
        Return {scale_state: [node_instance_name, ], }
        '''
        states_instances = {}
        for node_instance_name, node_instance in self._get_nodes_instances().iteritems():
            state = node_instance.get_scale_state()
            if not self._is_scale_state_terminal(state):
                state = self._get_effective_scale_state(node_instance_name)
                states_instances.setdefault(state, []).append(node_instance_name)
        return states_instances

    def _get_global_scale_state(self):
        '''Return scale state all the node instances are in, or None.
        Raise ExecutionException if there are instances in different states.
        For consistency reasons, only single scalability action is allowed.
        '''
        states_node_instances = self._get_effective_scale_states()

        if len(states_node_instances) == 0:
            return None

        if len(states_node_instances) == 1:
            return states_node_instances.keys()[0]

        msg = "Inconsistent scaling situation. Single scaling action allowed," \
            " found: %s" % states_node_instances
        raise Exceptions.ExecutionException(msg)

    def get_global_scale_action(self):
        state = self._get_global_scale_state()
        return self._state_to_action(state)

    def get_scale_action(self):
        state = self.get_scale_state()
        return self._state_to_action(state)

    def check_scale_state_consistency(self):
        states_node_instances = self._get_effective_scale_states()
        states_node_instances.pop(self.SCALE_STATE_REMOVED, None)

        if len(states_node_instances) > 1:
            msg = "Inconsistent scaling situation. Single scaling action allowed," \
                  " found: %s" % states_node_instances
            raise Exceptions.ExecutionException(msg)

    def get_scaling_node_and_instance_names(self):
        '''Return name of the node and the corresponding instances that are
        currently being scaled.
        Return tuple: node_name, [node_instance_name, ]
        '''
        node_names = set()
        node_instance_names = []

        for node_instance_name, node_instance in self._get_nodes_instances().iteritems():
            state = node_instance.get_scale_state()
            if not self._is_scale_state_terminal(state):
                node_names.add(node_instance.get_node_name())
                node_instance_names.append(node_instance_name)

        if len(node_names) != 1:
            msg = "Inconsistent scaling situation. Scaling of only single" \
                " node type is allowed, found: %s" % ', '.join(node_names)
            raise Exceptions.ExecutionException(msg)

        return node_names.pop(), node_instance_names

    def _state_to_action(self, state):
        return self.STATE_TO_ACTION.get(state, None)

    def get_node_instances_in_scale_state(self, scale_state, cloud_service_name=None):
        '''Return dict {<node_instance_name>: NodeInstance, } with the node instances
        in the scale_state.
        '''
        instances = {}

        nodes_instances = self._get_nodes_instances(cloud_service_name)
        for instance_name, instance in nodes_instances.iteritems():
            if instance.get_scale_state() == scale_state:
                instances[instance_name] = instance

        return instances

    def send_report(self, filename):
        self._ss_client.sendReport(filename)

    def set_statecustom(self, message):
        key = self._qualifyKey(NodeDecorator.STATECUSTOM_KEY)
        self._ss_client.setRuntimeParameter(key, message)

    #
    # Local cache of NodesInstances, Run, Run Parameters and User.
    #
    def discard_nodes_info_locally(self):
        self._nodes_instances = {}

    def _get_nodes_instances(self, cloud_service_name=None):
        '''Return dict {<node_instance_name>: NodeInstance, }
        '''
        if not self._nodes_instances:
            self._nodes_instances = self._ss_client.get_nodes_instances(cloud_service_name)
        return self._nodes_instances

    def get_my_node_instance(self):
        node_name = self.get_my_node_instance_name()
        return self._get_nodes_instances().get(node_name)

    def discard_run_locally(self):
        self._ss_client.discard_run()

    def discard_user_info_locally(self):
        self._user_info = None

    def _get_user_info(self, cloud_service_name):
        if self._user_info is None:
            self._user_info = self._ss_client.get_user_info(cloud_service_name)
        return self._user_info

    def discard_run_parameters_locally(self):
        self._run_parameters = None

    def _get_run_parameters(self):
        if self._run_parameters is None:
            self._run_parameters = self._ss_client.get_run_parameters()
        return self._run_parameters

    def has_to_execute_build_recipes(self):
        run_parameters = self._get_run_parameters()

        key = self.get_my_node_instance_name().rsplit(NodeDecorator.NODE_MULTIPLICITY_SEPARATOR, 1)[0] \
              + NodeDecorator.NODE_PROPERTY_SEPARATOR \
              + NodeDecorator.RUN_BUILD_RECIPES_KEY

        return util.str2bool(run_parameters.get(key))

    def clean_local_cache(self):
        self.discard_run_locally()
        self.discard_nodes_info_locally()
        self.discard_run_parameters_locally()

    #
    # Helpers
    #
    def _terminate_run_server_side(self):
        self._ss_client.terminate_run()

    def _put_new_image_id(self, url, new_image_id):
        self._ss_client.put_new_image_id(url, new_image_id)
<commit_msg>Minor.<commit_after>"""
 SlipStream Client
 =====
 Copyright (C) 2014 SixSq Sarl (sixsq.com)
 =====
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import traceback

from slipstream import util
from slipstream.SlipStreamHttpClient import SlipStreamHttpClient
from slipstream.NodeDecorator import NodeDecorator
from slipstream.exceptions import Exceptions

class NodeInfoPublisher(SlipStreamHttpClient):
    def __init__(self, config_holder):
        super(NodeInfoPublisher, self).__init__(config_holder)
        self.set_http_max_retries(-5)

    def publish(self, nodename, vm_id, vm_ip):
        self.publish_instanceid(nodename, vm_id)
        self.publish_hostname(nodename, vm_ip)

    def publish_instanceid(self, nodename, vm_id):
        self._set_runtime_parameter(nodename, 'instanceid', vm_id)

    def publish_hostname(self, nodename, vm_ip):
        self._set_runtime_parameter(nodename, 'hostname', vm_ip)

    def publish_url_ssh(self, nodename, vm_ip, username):
        url = 'ssh://%s@%s' % (username.strip(), vm_ip.strip())
        self._set_runtime_parameter(nodename, 'url.ssh', url)

    def _set_runtime_parameter(self, nodename, key, value):
        parameter = nodename + NodeDecorator.NODE_PROPERTY_SEPARATOR + key
        self.setRuntimeParameter(parameter, value, ignoreAbort=True)

class BaseWrapper(object):

    SCALE_STATE_CREATING = 'creating'
    SCALE_STATE_CREATED = 'created'

    SCALE_STATE_DISK_RESIZING = 'disk_resizing'
    SCALE_STATE_DISK_RESIZED = 'disk_resized'

    SCALE_STATE_OPERATIONAL = 'operational'

    SCALE_STATE_REMOVING = 'removing'
    SCALE_STATE_REMOVED = 'removed'
    SCALE_STATE_GONE = 'gone'

    SCALE_STATES_TERMINAL = (SCALE_STATE_OPERATIONAL, SCALE_STATE_GONE)

    SCALE_ACTION_CREATION = 'creation'
    SCALE_ACTION_REMOVAL = 'removal'
    SCALE_ACTION_DISK_RESIZE = 'disk_resize'

    STATE_TO_ACTION = {SCALE_STATE_CREATED: SCALE_ACTION_CREATION,
                       SCALE_STATE_REMOVED: SCALE_ACTION_REMOVAL,
                       SCALE_STATE_DISK_RESIZED: SCALE_ACTION_DISK_RESIZE}

    def __init__(self, config_holder):
        self._ss_client = SlipStreamHttpClient(config_holder)
        self._ss_client.set_http_max_retries(self.is_mutable() and -5 or 5)
        self._ss_client.ignoreAbort = True

        self.my_node_instance_name = self._get_my_node_instance_name(config_holder)

        self._user_info = None
        self._run_parameters = None
        self._nodes_instances = {}

    def _get_my_node_instance_name(self, config_holder):
        try:
            return config_holder.node_instance_name
        except Exception:
            raise Exceptions.ExecutionException('Failed to get the node instance name of the the current VM')

    def get_my_node_instance_name(self):
        return self.my_node_instance_name

    def get_slipstream_client(self):
        return self._ss_client

    def complete_state(self, node_instance_name=None):
        if not node_instance_name:
            node_instance_name = self.get_my_node_instance_name()
        self._ss_client.complete_state(node_instance_name)

    # TODO: LS: Can we remove this method ?
    def reset(self):
        self._ss_client.reset()

    def fail(self, message):
        key = self._qualifyKey(NodeDecorator.ABORT_KEY)
        self._fail(key, message)

    def fail_global(self, message):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.ABORT_KEY
        self._fail(key, message)

    def _fail(self, key, message):
        util.printError('Failing... %s' % message)
        traceback.print_exc()
        self._ss_client.setRuntimeParameter(key, message)

    def getState(self):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.STATE_KEY
        return self._get_runtime_parameter(key)

    def get_recovery_mode(self):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.RECOVERY_MODE_KEY
        return util.str2bool(self._get_runtime_parameter(key))

    def isAbort(self):
        key = NodeDecorator.globalNamespacePrefix + NodeDecorator.ABORT_KEY
        try:
            value = self._get_runtime_parameter(key, True)
        except Exceptions.NotYetSetException:
            value = ''
        return (value and True) or False

    def get_run_category(self):
        return self._ss_client.get_run_category()

    def get_run_type(self):
        return self._ss_client.get_run_type()

    def _qualifyKey(self, key):
        """Qualify the key, if not already done, with the right nodename"""

        _key = key

        # Is this a reserved or special nodename?
        for reserved in NodeDecorator.reservedNodeNames:
            if _key.startswith(reserved + NodeDecorator.NODE_PROPERTY_SEPARATOR):
                return _key

        # Is the key namespaced (i.e. contains node/key separator: ':')?
        if NodeDecorator.NODE_PROPERTY_SEPARATOR in _key:
            # Is the nodename in the form: <nodename>.<index>?  If not, make it so
            # such that <nodename>:<property> -> <nodename>.1:<property
            parts = _key.split(NodeDecorator.NODE_PROPERTY_SEPARATOR)
            nodenamePart = parts[0]
            propertyPart = parts[1]  # safe since we've done the test in the if above
            parts = nodenamePart.split(NodeDecorator.NODE_MULTIPLICITY_SEPARATOR)
            nodename = parts[0]
            if len(parts) == 1:
                _key = nodename + \
                    NodeDecorator.NODE_MULTIPLICITY_SEPARATOR + \
                    NodeDecorator.nodeMultiplicityStartIndex + \
                    NodeDecorator.NODE_PROPERTY_SEPARATOR + \
                    propertyPart
            return _key

        _key = self.get_my_node_instance_name() + NodeDecorator.NODE_PROPERTY_SEPARATOR + _key

        return _key

    def get_cloud_instance_id(self):
        key = self._qualifyKey(NodeDecorator.INSTANCEID_KEY)
        return self._get_runtime_parameter(key)

    def get_user_ssh_pubkey(self):
        userInfo = self._get_user_info('')
        return userInfo.get_public_keys()

    def _get_runtime_parameter(self, key, ignore_abort=False):
        return self._ss_client.getRuntimeParameter(key, ignoreAbort=ignore_abort)

    def need_to_stop_images(self, ignore_on_success_run_forever=False):
        # pylint: disable=unused-argument
        return False

    def is_mutable(self):
        mutable = self._ss_client.get_run_mutable()
        return util.str2bool(mutable)

    def set_scale_state_on_node_instances(self, instance_names, scale_state):
        for instance_name in instance_names:
            key = instance_name + NodeDecorator.NODE_PROPERTY_SEPARATOR + NodeDecorator.SCALE_STATE_KEY
            self._ss_client.setRuntimeParameter(key, scale_state)

    def is_scale_state_operational(self):
        return self.SCALE_STATE_OPERATIONAL == self.get_scale_state()

    def is_scale_state_creating(self):
        return self.SCALE_STATE_CREATING == self.get_scale_state()

    def _is_scale_state_terminal(self, state):
        return state in self.SCALE_STATES_TERMINAL

    def set_scale_state_operational(self):
        self.set_scale_state(self.SCALE_STATE_OPERATIONAL)

    def set_scale_state_created(self):
        self.set_scale_state(self.SCALE_STATE_CREATED)

    def set_scale_state(self, scale_state):
        '''Set scale state for this node instances.
        '''
        key = self._qualifyKey(NodeDecorator.SCALE_STATE_KEY)
        self._ss_client.setRuntimeParameter(key, scale_state)

    def get_scale_state(self):
        '''Get scale state for this node instances.
        '''
        key = self._qualifyKey(NodeDecorator.SCALE_STATE_KEY)
        return self._get_runtime_parameter(key)

    def _get_effective_scale_state(self, node_instance_name):
        '''Get effective node instance scale state from the server.
        '''
        key = node_instance_name + NodeDecorator.NODE_PROPERTY_SEPARATOR + \
            NodeDecorator.SCALE_STATE_KEY
        return self._get_runtime_parameter(key)

    def _get_effective_scale_states(self):
        '''Extract node instances in scaling states and update their effective
        states from the server.
        Return {scale_state: [node_instance_name, ], }
        '''
        states_instances = {}
        for node_instance_name, node_instance in self._get_nodes_instances().iteritems():
            state = node_instance.get_scale_state()
            if not self._is_scale_state_terminal(state):
                state = self._get_effective_scale_state(node_instance_name)
                states_instances.setdefault(state, []).append(node_instance_name)
        return states_instances

    def _get_global_scale_state(self):
        '''Return scale state all the node instances are in, or None.
        Raise ExecutionException if there are instances in different states.
        For consistency reasons, only single scalability action is allowed.
        '''
        states_node_instances = self._get_effective_scale_states()

        if len(states_node_instances) == 0:
            return None

        if len(states_node_instances) == 1:
            return states_node_instances.keys()[0]

        msg = "Inconsistent scaling situation. Single scaling action allowed," \
            " found: %s" % states_node_instances
        raise Exceptions.ExecutionException(msg)

    def get_global_scale_action(self):
        state = self._get_global_scale_state()
        return self._state_to_action(state)

    def get_scale_action(self):
        state = self.get_scale_state()
        return self._state_to_action(state)

    def check_scale_state_consistency(self):
        states_node_instances = self._get_effective_scale_states()
        states_node_instances.pop(self.SCALE_STATE_REMOVED, None)

        if len(states_node_instances) > 1:
            msg = "Inconsistent scaling situation. Single scaling action allowed," \
                  " found: %s" % states_node_instances
            raise Exceptions.ExecutionException(msg)

    def get_scaling_node_and_instance_names(self):
        '''Return name of the node and the corresponding instances that are
        currently being scaled.
        Return tuple: node_name, [node_instance_name, ]
        '''
        node_names = set()
        node_instance_names = []

        for node_instance_name, node_instance in self._get_nodes_instances().iteritems():
            state = node_instance.get_scale_state()
            if not self._is_scale_state_terminal(state):
                node_names.add(node_instance.get_node_name())
                node_instance_names.append(node_instance_name)

        if len(node_names) != 1:
            msg = "Inconsistent scaling situation. Scaling of only single" \
                " node type is allowed, found: %s" % ', '.join(node_names)
            raise Exceptions.ExecutionException(msg)

        return node_names.pop(), node_instance_names

    def _state_to_action(self, state):
        return self.STATE_TO_ACTION.get(state, None)

    def get_node_instances_in_scale_state(self, scale_state, cloud_service_name=None):
        '''Return dict {<node_instance_name>: NodeInstance, } with the node instances
        in the scale_state.
        '''
        instances = {}

        nodes_instances = self._get_nodes_instances(cloud_service_name)
        for instance_name, instance in nodes_instances.iteritems():
            if instance.get_scale_state() == scale_state:
                instances[instance_name] = instance

        return instances

    def send_report(self, filename):
        self._ss_client.sendReport(filename)

    def set_statecustom(self, message):
        key = self._qualifyKey(NodeDecorator.STATECUSTOM_KEY)
        self._ss_client.setRuntimeParameter(key, message)

    #
    # Local cache of NodesInstances, Run, Run Parameters and User.
    #
    def discard_nodes_info_locally(self):
        self._nodes_instances = {}

    def _get_nodes_instances(self, cloud_service_name=None):
        '''Return dict {<node_instance_name>: NodeInstance, }
        '''
        if not self._nodes_instances:
            self._nodes_instances = self._ss_client.get_nodes_instances(cloud_service_name)
        return self._nodes_instances

    def get_my_node_instance(self):
        node_name = self.get_my_node_instance_name()
        return self._get_nodes_instances().get(node_name)

    def discard_run_locally(self):
        self._ss_client.discard_run()

    def discard_user_info_locally(self):
        self._user_info = None

    def _get_user_info(self, cloud_service_name):
        if self._user_info is None:
            self._user_info = self._ss_client.get_user_info(cloud_service_name)
        return self._user_info

    def discard_run_parameters_locally(self):
        self._run_parameters = None

    def _get_run_parameters(self):
        if self._run_parameters is None:
            self._run_parameters = self._ss_client.get_run_parameters()
        return self._run_parameters

    def has_to_execute_build_recipes(self):
        run_parameters = self._get_run_parameters()

        key = self.get_my_node_instance_name().rsplit(NodeDecorator.NODE_MULTIPLICITY_SEPARATOR, 1)[0] \
              + NodeDecorator.NODE_PROPERTY_SEPARATOR \
              + NodeDecorator.RUN_BUILD_RECIPES_KEY

        return util.str2bool(run_parameters.get(key))

    def clean_local_cache(self):
        self.discard_run_locally()
        self.discard_nodes_info_locally()
        self.discard_run_parameters_locally()

    #
    # Helpers
    #
    def _terminate_run_server_side(self):
        self._ss_client.terminate_run()

    def _put_new_image_id(self, url, new_image_id):
        self._ss_client.put_new_image_id(url, new_image_id)
<|endoftext|>


<commit_before># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
import astropy.units as u
from astropy.units.physical import _physical_unit_mapping

class UnitSystem(object):
    """
    Represents a system of units. At minimum, this consists of a set of
    length, time, mass, and angle units, but may also contain preferred
    representations for composite units. For example, the base unit system
    could be ``{kpc, Myr, Msun, radian}``, but you can also specify a preferred
    speed, such as ``km/s``.

    This class functions like a dictionary with keys set by physical types.
    If a unit for a particular physical type is not specified on creation,
    a composite unit will be created with the base units. See Examples below
    for some demonstrations.

    Parameters
    ----------
    *units
        The units that define the unit system. At minimum, this must
        contain length, time, mass, and angle units.

    Examples
    --------
    If only base units are specified, any physical type specified as a key
    to this object will be composed out of the base units::

        >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian)
        >>> usys['energy']
        Unit("kg m2 / s2")

    However, custom representations for composite units can also be specified
    when initializing::

        >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian, u.erg)
        >>> usys['energy']
        Unit("erg")

    This is useful for Galactic dynamics where lengths and times are usually
    given in terms of ``kpc`` and ``Myr``, but speeds are given in ``km/s``::

        >>> usys = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, u.km/u.s)
        >>> usys['speed']
        Unit("km / s")

    """
    def __init__(self, *units):

        self._required_physical_types = ['length', 'time', 'mass', 'angle']
        self._core_units = []

        self._registry = dict()
        for unit in units:
            typ = unit.physical_type
            if typ in self._registry:
                raise ValueError("Multiple units passed in with type '{0}'".format(typ))
            self._registry[typ] = unit

        for phys_type in self._required_physical_types:
            if phys_type not in self._registry:
                raise ValueError("You must specify a unit with physical type '{0}'".format(phys_type))
            self._core_units.append(self._registry[phys_type])

    def __getitem__(self, key):

        if key in self._registry:
            return self._registry[key]

        else:
            unit = None
            for k,v in _physical_unit_mapping.items():
                if v == key:
                    unit = u.Unit(" ".join(["{}**{}".format(x,y) for x,y in k]))
                    break

            if unit is None:
                raise ValueError("Physical type '{0}' doesn't exist in unit registry.".format(key))

            unit = unit.decompose(self._core_units)
            unit._scale = 1.
            return unit

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

    def __iter__(self):
        for uu in self._core_units:
            yield uu

    def __str__(self):
        return "UnitSystem ({0})".format(",".join([str(uu) for uu in self._core_units]))

    def __repr__(self):
        return "<{0}>".format(self.__str__())

    def to_dict(self):
        return self._registry.copy()

# define galactic unit system
galactic = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian,
                      u.km/u.s, u.mas/u.yr)

# solar system units
solarsystem = UnitSystem(u.au, u.M_sun, u.yr, u.radian)
<commit_msg>docstrings and added a decompose method to unit system<commit_after># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
import astropy.units as u
from astropy.units.physical import _physical_unit_mapping

class UnitSystem(object):
    """
    Represents a system of units. At minimum, this consists of a set of
    length, time, mass, and angle units, but may also contain preferred
    representations for composite units. For example, the base unit system
    could be ``{kpc, Myr, Msun, radian}``, but you can also specify a preferred
    speed, such as ``km/s``.

    This class functions like a dictionary with keys set by physical types.
    If a unit for a particular physical type is not specified on creation,
    a composite unit will be created with the base units. See Examples below
    for some demonstrations.

    Parameters
    ----------
    *units
        The units that define the unit system. At minimum, this must
        contain length, time, mass, and angle units.

    Examples
    --------
    If only base units are specified, any physical type specified as a key
    to this object will be composed out of the base units::

        >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian)
        >>> usys['energy']
        Unit("kg m2 / s2")

    However, custom representations for composite units can also be specified
    when initializing::

        >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian, u.erg)
        >>> usys['energy']
        Unit("erg")

    This is useful for Galactic dynamics where lengths and times are usually
    given in terms of ``kpc`` and ``Myr``, but speeds are given in ``km/s``::

        >>> usys = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, u.km/u.s)
        >>> usys['speed']
        Unit("km / s")

    """
    def __init__(self, *units):

        self._required_physical_types = ['length', 'time', 'mass', 'angle']
        self._core_units = []

        self._registry = dict()
        for unit in units:
            typ = unit.physical_type
            if typ in self._registry:
                raise ValueError("Multiple units passed in with type '{0}'".format(typ))
            self._registry[typ] = unit

        for phys_type in self._required_physical_types:
            if phys_type not in self._registry:
                raise ValueError("You must specify a unit with physical type '{0}'".format(phys_type))
            self._core_units.append(self._registry[phys_type])

    def __getitem__(self, key):

        if key in self._registry:
            return self._registry[key]

        else:
            unit = None
            for k,v in _physical_unit_mapping.items():
                if v == key:
                    unit = u.Unit(" ".join(["{}**{}".format(x,y) for x,y in k]))
                    break

            if unit is None:
                raise ValueError("Physical type '{0}' doesn't exist in unit registry.".format(key))

            unit = unit.decompose(self._core_units)
            unit._scale = 1.
            return unit

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

    def __iter__(self):
        for uu in self._core_units:
            yield uu

    def __str__(self):
        return "UnitSystem ({0})".format(",".join([str(uu) for uu in self._core_units]))

    def __repr__(self):
        return "<{0}>".format(self.__str__())

    def to_dict(self):
        """
        Return a dictionary representation of the unit system with keys
        set by the physical types and values set by the unit objects.
        """
        return self._registry.copy()

    def decompose(self, q):
        """
        A thin wrapper around `astropy.units.Quantity.decompose()` that
        knows how to handle Quantities with physical types with non-default
        representations.

        Parameters
        ----------
        q : :class:`~astropy.units.Quantity`
            An instance of an astropy Quantity object.

        Returns
        -------
        q : :class:`~astropy.units.Quantity`
            A new quantity, decomposed to represented in this unit system.
        """
        try:
            ptype = q.unit.physical_type
        except AttributeError:
            raise TypeError("Object must be an astropy.units.Quantity, not "
                            "a '{}'.".format(q.__class__.__name__))

        if ptype in self._registry:
            return q.to(self._registry[ptype])
        else:
            return q.decompose(self)

# define galactic unit system
galactic = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian,
                      u.km/u.s, u.mas/u.yr)

# solar system units
solarsystem = UnitSystem(u.au, u.M_sun, u.yr, u.radian)
<|endoftext|>