<commit_before>import ast
from collections import OrderedDict
from .codegen import to_source
from .function_compiler_ast import timeshift, StandardizeDatesSimple
from dolo.compiler.recipes import recipes
from numba import njit

class NumericModel:

    calibration = None
    calibration_dict = None
    covariances = None
    markov_chain = None

    def __init__(self, symbolic_model, options=None, infos=None):

        self.symbolic = symbolic_model

        self.symbols = symbolic_model.symbols

        # compat
        if self.symbolic.definitions:
            self.symbols['auxiliaries'] = [e for e in self.symbolic.definitions.keys()]

        self.variables = sum( [tuple(e) for k,e in  self.symbols.items() if k not in ('parameters','shocks','values')], ())

        self.options = options if options is not None else {}

        self.infos = infos
        self.name = self.infos['name']
        self.model_type = self.infos['type']
        # self.model_spec
        self.__update_from_symbolic__()
        self.__compile_functions__()

    def __update_from_symbolic__(self):

        import numpy
        # updates calibration according to the symbolic definitions

        system = self.symbolic.calibration_dict

        from dolo.compiler.triangular_solver import solve_triangular_system
        self.calibration_dict = solve_triangular_system( system )

        from dolo.compiler.misc import CalibrationDict, calibration_to_vector
        calib = calibration_to_vector(self.symbols, self.calibration_dict)
        self.calibration = CalibrationDict(self.symbols, calib)

        from .symbolic_eval import NumericEval
        evaluator = NumericEval(self.calibration_dict)
        # read symbolic structure
        self.options = evaluator.eval(self.symbolic.options)
        #
        # distribution = evaluator.eval(self.symbolic.options.get('distribution'))
        # discrete_transition = evaluator.eval(self.symbolic.options.get('discrete_transition'))

        distribution = self.options.get('distribution')
        discrete_transition = self.options.get('discrete_transition')

        self.distribution = distribution
        self.discrete_transition = discrete_transition

        if distribution is None:
            self.covariances = None
        else:
            self.covariances = numpy.atleast_2d(numpy.array(distribution.sigma, dtype=float))

        markov_chain = discrete_transition

        if markov_chain is None:
            self.markov_chain = None
        else:
            markov_chain = [markov_chain.P, markov_chain.Q]
            self.markov_chain = [numpy.atleast_2d(numpy.array(tab, dtype=float)) for tab in markov_chain]

    def get_calibration(self, pname, *args):

        if isinstance(pname, list):
            return [ self.get_calibration(p) for p in pname ]
        elif isinstance(pname, tuple):
            return tuple( [ self.get_calibration(p) for p in pname ] )
        elif len(args)>0:
            pnames = (pname,) + args
            return self.get_calibration(pnames)

        group = [g for g in self.symbols.keys() if pname in self.symbols[g]]
        try:
            group = group[0]
        except Exception:
            raise Exception('Unknown symbol {}.'.format(pname))
        i = self.symbols[group].index(pname)
        v = self.calibration[group][i]

        return v

    def set_calibration(self, *args, **kwargs):

        # raise exception if unknown symbol ?

        if len(args)==2:
            pname, pvalue = args
            if isinstance(pname, str):
                self.set_calibration(**{pname:pvalue})
        else:
            # else ignore pname and pvalue
            calib =  self.symbolic.calibration_dict
            calib.update(kwargs)
            self.__update_from_symbolic__()

    def __str__(self):

        from dolo.misc.termcolor import colored
        from numpy import zeros
        s = u'''
Model:
------
name: "{name}"
type: "{type}"
file: "{filename}\n'''.format(**self.infos)

        ss = '\nEquations:\n----------\n\n'
        res = self.residuals()
        res.update({'definitions': zeros(1)})

        equations = self.symbolic.equations
        definitions = self.symbolic.definitions
        tmp = []
        for deftype in definitions:
            tmp.append(deftype + ' = ' + definitions[deftype])
        definitions = {'definitions': tmp}
        equations.update(definitions)

        # for eqgroup, eqlist in self.symbolic.equations.items():
        for eqgroup in res.keys():
            if eqgroup == 'auxiliary':
                continue
            if eqgroup == 'dynare':
                eqlist = equations
            if eqgroup == 'definitions':
                eqlist = equations[eqgroup]
                # Update the residuals section with the right number of empty
                # values. Note: adding 'zeros' was easiest (rather than empty
                # cells), since other variable types have  arrays of zeros.
                res.update({'definitions': [None for i in range(len(eqlist))]})
            else:
                eqlist = equations[eqgroup]
            ss += u"{}\n".format(eqgroup)
            for i, eq in enumerate(eqlist):
                val = res[eqgroup][i]
                if val is None:
                    ss += u" {eqn:2} : {eqs}\n".format(eqn=str(i+1), eqs=eq)
                else:
                    if abs(val) < 1e-8:
                        val = 0
                    vals = '{:.4f}'.format(val)
                    if abs(val) > 1e-8:
                        vals = colored(vals, 'red')
                    ss += u" {eqn:2} : {vals} : {eqs}\n".format(eqn=str(i+1), vals=vals, eqs=eq)
            ss += "\n"
        s += ss

        return s

    def __repr__(self):
        return self.__str__()

    def _repr_html_(self):

        from dolo.compiler.latex import eq2tex

        # general informations
        infos = self.infos
        table_infos = """

         Model

        name
        {name}

        type
        {type}

        filename
        {filename}

""".format(name=infos['name'], type=infos['type'], filename=infos['filename'].replace("<","&lt").replace(">","&gt"))

        # Equations and residuals
        resids = self.residuals()
        if self.model_type == 'dynare':
            equations = {"dynare": self.symbolic.equations}
        else:
            equations = self.symbolic.equations
            # Create definitions equations and append to equations dictionary
            definitions = self.symbolic.definitions
            tmp = []
            for deftype in definitions:
                tmp.append(deftype + ' = ' + definitions[deftype])
            definitions = {'definitions': tmp}
            equations.update(definitions)

        variables = sum([e for k,e in self.symbols.items() if k != 'parameters'], [])
        table = "
Type	Equation	Residual
\n"

        for eq_type in equations:

            eq_lines = []
            for i in range(len(equations[eq_type])):
                eq = equations[eq_type][i]
                # if eq_type in ('expectation','direct_response'):
                #     vals = ''
                if eq_type not in ('arbitrage', 'transition', 'arbitrage_exp'):
                    vals = ''
                else:
                    val = resids[eq_type][i]
                    if abs(val) > 1e-8:
                        vals = '{:.4f}'.format(val)
                    else:
                        vals = '{:.3f}'.format(val)
                if '|' in eq:
                    # keep only lhs for now
                    eq, comp = str.split(eq,'|')
                lat = eq2tex(variables, eq)
                lat = '${}$'.format(lat)
                line = [lat, vals]
                h = eq_type if i==0 else ''
                fmt_line = '
{}	{}	{}
'.format(h, *line)
        #         print(fmt_line)
                eq_lines.append(fmt_line)
            table += str.join("\n", eq_lines)
        table = "
{}
".format(table)

        return table_infos + table

    @property
    def x_bounds(self):

        if 'controls_ub' in self.functions:
            fun_lb = self.functions['controls_lb']
            fun_ub = self.functions['controls_ub']
            return [fun_lb, fun_ub]
        else:
            return None

    def residuals(self, calib=None):

        if self.model_type == 'dtcscc':
            from dolo.algos.dtcscc.steady_state import residuals
        elif self.model_type == 'dtmscc':
            from dolo.algos.dtmscc.steady_state import residuals
        elif self.model_type == 'dynare':
            from dolo.algos.dynare.steady_state import residuals
        return residuals(self, calib)

    def eval_formula(self, expr, dataframe=None, calib=None):

        from dolo.compiler.eval_formula import eval_formula
        if calib is None:
            calib = self.calibration
        return eval_formula(expr, dataframe=dataframe, context=calib)

    def get_distribution(model, **opts):
        import copy
        gg = model.symbolic.options.get('distribution')
        if gg is None:
            raise Exception("Model has no distribution.")
        d = copy.deepcopy(gg)
        d.update(opts)
        if 'type' in d: d.pop('type')
        return d.eval(model.calibration.flat)

    def get_grid(model, **dis_opts):
        import copy
        gg = model.symbolic.options.get('grid')
        if gg is None:
            raise Exception("Model has no grid.")
        d = copy.deepcopy(gg)
        gtype = dis_opts.get('type')
        if gtype:
            from dolo.compiler.language import minilang
            try:
                cls = [e for e in minilang if e.__name__.lower()==gtype.lower()][0]
            except:
                raise Exception("Unknown grid type {}.".format(gtype))
            d = cls(**d)
    #     return cc
        d.update(**dis_opts)
        if 'type' in d: d.pop('type')
        return d.eval(d=model.calibration.flat)

    def __compile_functions__(self):

        from dolo.compiler.function_compiler import compile_function_ast
        from dolo.compiler.function_compiler import standard_function

        defs = self.symbolic.definitions

        model_type = self.model_type

        recipe = recipes[model_type]
        symbols = self.symbols # should match self.symbols

        comps = []

        functions = {}
        original_functions = {}
        original_gufunctions = {}

        # for funname in recipe['specs'].keys():
        for funname in self.symbolic.equations:

            spec = recipe['specs'][funname]

            eqs = self.symbolic.equations[funname]

            if spec.get('target'):
                target = spec['target']
                rhs_only = True
            else:
                target = None
                rhs_only = False

            if spec.get('complementarities'):

                # TODO: Rewrite and simplify
                comp_spec = spec.get('complementarities')
                comp_order = comp_spec['middle']
                comp_args = comp_spec['left-right']

                comps = []
                eqs = []
                for i,eq in enumerate(self.symbolic.equations[funname]):

                    if '|' in eq:
                        control = self.symbols[comp_order[0]][i]
                        eq, comp = str.split(eq,'|')
                        lhs, rhs = decode_complementarity(comp, control)
                        comps.append([lhs, rhs])
                    else:
                        comps.append(['-inf', 'inf'])

                    eqs.append(eq)

                    comp_lhs, comp_rhs = zip(*comps)
                    # fb_names = ['{}_lb'.format(funname), '{}_ub'.format(funname)]
                    fb_names = ['controls_lb'.format(funname), 'controls_ub'.format(funname)]

                    lower_bound, gu_lower_bound = compile_function_ast(comp_lhs, symbols, comp_args, funname=fb_names[0],definitions=defs)
                    upper_bound, gu_upper_bound = compile_function_ast(comp_rhs, symbols, comp_args, funname=fb_names[1],definitions=defs)

                    n_output = len(comp_lhs)

                    functions[fb_names[0]] = standard_function(gu_lower_bound, n_output )
                    functions[fb_names[1]] = standard_function(gu_upper_bound, n_output )
                    original_functions[fb_names[0]] = lower_bound
                    original_functions[fb_names[1]] = upper_bound
                    original_gufunctions[fb_names[0]] = gu_lower_bound
                    original_gufunctions[fb_names[1]] = gu_upper_bound

            arg_names = recipe['specs'][funname]['eqs']

            fun, gufun = compile_function_ast(eqs, symbols, arg_names, output_names=target, rhs_only=rhs_only, funname=funname, definitions=defs,                                   )
            n_output = len(eqs)
            original_functions[funname] = fun
            functions[funname] = standard_function(gufun, n_output )
            original_functions[funname] = fun
            original_gufunctions[funname] = gufun

        # temporary hack to keep auxiliaries
        auxiliaries = [k for k in defs.keys()]
        symbols['auxiliaries'] = auxiliaries
        eqs = ['{} = {}'.format(k,k) for k in auxiliaries]
        if model_type == 'dtcscc':
            arg_names = [('states',0,'s'),('controls',0,'x'),('parameters',0,'p')]
        else:
            arg_names = [('markov_states',0,'m'),('states',0,'s'),('controls',0,'x'),('parameters',0,'p')]
        target = ('auxiliaries',0,'y')
        rhs_only = True
        funname = 'auxiliary'
        fun, gufun = compile_function_ast(eqs, symbols, arg_names, output_names=target,
            rhs_only=rhs_only, funname=funname, definitions=defs)
        n_output = len(eqs)
        functions[funname] = standard_function(gufun, n_output )
        original_functions[funname] = fun
        original_gufunctions[funname] = gufun

        self.__original_functions__ = original_functions
        self.__original_gufunctions__ = original_gufunctions
        self.functions = functions

import re
regex = re.compile("(.*)<=(.*)<=(.*)")

def decode_complementarity(comp, control):

    '''
    # comp can be either:
    - None
    - "a<=expr" where a is a controls
    - "expr<=a" where a is a control
    - "expr1<=a<=expr2"
    '''

    try:
        res = regex.match(comp).groups()
    except:
        raise Exception("Unable to parse complementarity condition '{}'".format(comp))
    res = [r.strip() for r in res]
    if res[1] != control:
        msg = "Complementarity condition '{}' incorrect. Expected {} instead of {}.".format(comp, control, res[1])
        raise Exception(msg)
    return [res[0], res[2]]
<commit_msg>FIX: printing definitions modify list of equations<commit_after>import ast
from collections import OrderedDict
from .codegen import to_source
from .function_compiler_ast import timeshift, StandardizeDatesSimple
from dolo.compiler.recipes import recipes
from numba import njit

class NumericModel:

    calibration = None
    calibration_dict = None
    covariances = None
    markov_chain = None

    def __init__(self, symbolic_model, options=None, infos=None):

        self.symbolic = symbolic_model

        self.symbols = symbolic_model.symbols

        # compat
        if self.symbolic.definitions:
            self.symbols['auxiliaries'] = [e for e in self.symbolic.definitions.keys()]

        self.variables = sum( [tuple(e) for k,e in  self.symbols.items() if k not in ('parameters','shocks','values')], ())

        self.options = options if options is not None else {}

        self.infos = infos
        self.name = self.infos['name']
        self.model_type = self.infos['type']
        # self.model_spec
        self.__update_from_symbolic__()
        self.__compile_functions__()

    def __update_from_symbolic__(self):

        import numpy
        # updates calibration according to the symbolic definitions

        system = self.symbolic.calibration_dict

        from dolo.compiler.triangular_solver import solve_triangular_system
        self.calibration_dict = solve_triangular_system( system )

        from dolo.compiler.misc import CalibrationDict, calibration_to_vector
        calib = calibration_to_vector(self.symbols, self.calibration_dict)
        self.calibration = CalibrationDict(self.symbols, calib)

        from .symbolic_eval import NumericEval
        evaluator = NumericEval(self.calibration_dict)
        # read symbolic structure
        self.options = evaluator.eval(self.symbolic.options)
        #
        # distribution = evaluator.eval(self.symbolic.options.get('distribution'))
        # discrete_transition = evaluator.eval(self.symbolic.options.get('discrete_transition'))

        distribution = self.options.get('distribution')
        discrete_transition = self.options.get('discrete_transition')

        self.distribution = distribution
        self.discrete_transition = discrete_transition

        if distribution is None:
            self.covariances = None
        else:
            self.covariances = numpy.atleast_2d(numpy.array(distribution.sigma, dtype=float))

        markov_chain = discrete_transition

        if markov_chain is None:
            self.markov_chain = None
        else:
            markov_chain = [markov_chain.P, markov_chain.Q]
            self.markov_chain = [numpy.atleast_2d(numpy.array(tab, dtype=float)) for tab in markov_chain]

    def get_calibration(self, pname, *args):

        if isinstance(pname, list):
            return [ self.get_calibration(p) for p in pname ]
        elif isinstance(pname, tuple):
            return tuple( [ self.get_calibration(p) for p in pname ] )
        elif len(args)>0:
            pnames = (pname,) + args
            return self.get_calibration(pnames)

        group = [g for g in self.symbols.keys() if pname in self.symbols[g]]
        try:
            group = group[0]
        except Exception:
            raise Exception('Unknown symbol {}.'.format(pname))
        i = self.symbols[group].index(pname)
        v = self.calibration[group][i]

        return v

    def set_calibration(self, *args, **kwargs):

        # raise exception if unknown symbol ?

        if len(args)==2:
            pname, pvalue = args
            if isinstance(pname, str):
                self.set_calibration(**{pname:pvalue})
        else:
            # else ignore pname and pvalue
            calib =  self.symbolic.calibration_dict
            calib.update(kwargs)
            self.__update_from_symbolic__()

    def __str__(self):

        from dolo.misc.termcolor import colored
        from numpy import zeros
        s = u'''
Model:
------
name: "{name}"
type: "{type}"
file: "{filename}\n'''.format(**self.infos)

        ss = '\nEquations:\n----------\n\n'
        res = self.residuals()
        res.update({'definitions': zeros(1)})

        equations = self.symbolic.equations.copy()
        definitions = self.symbolic.definitions
        tmp = []
        for deftype in definitions:
            tmp.append(deftype + ' = ' + definitions[deftype])
        definitions = {'definitions': tmp}
        equations.update(definitions)
        # for eqgroup, eqlist in self.symbolic.equations.items():
        for eqgroup in res.keys():
            if eqgroup == 'auxiliary':
                continue
            if eqgroup == 'dynare':
                eqlist = equations
            if eqgroup == 'definitions':
                eqlist = equations[eqgroup]
                # Update the residuals section with the right number of empty
                # values. Note: adding 'zeros' was easiest (rather than empty
                # cells), since other variable types have  arrays of zeros.
                res.update({'definitions': [None for i in range(len(eqlist))]})
            else:
                eqlist = equations[eqgroup]
            ss += u"{}\n".format(eqgroup)
            for i, eq in enumerate(eqlist):
                val = res[eqgroup][i]
                if val is None:
                    ss += u" {eqn:2} : {eqs}\n".format(eqn=str(i+1), eqs=eq)
                else:
                    if abs(val) < 1e-8:
                        val = 0
                    vals = '{:.4f}'.format(val)
                    if abs(val) > 1e-8:
                        vals = colored(vals, 'red')
                    ss += u" {eqn:2} : {vals} : {eqs}\n".format(eqn=str(i+1), vals=vals, eqs=eq)
            ss += "\n"
        s += ss

        return s

    def __repr__(self):
        return self.__str__()

    def _repr_html_(self):

        from dolo.compiler.latex import eq2tex

        # general informations
        infos = self.infos
        table_infos = """

         Model

        name
        {name}

        type
        {type}

        filename
        {filename}

""".format(name=infos['name'], type=infos['type'], filename=infos['filename'].replace("<","&lt").replace(">","&gt"))

        # Equations and residuals
        resids = self.residuals()
        if self.model_type == 'dynare':
            equations = {"dynare": self.symbolic.equations.copy()}
        else:
            equations = self.symbolic.equations.copy()
        # Create definitions equations and append to equations dictionary
        definitions = self.symbolic.definitions
        tmp = []
        for deftype in definitions:
            tmp.append(deftype + ' = ' + definitions[deftype])
        definitions = {'definitions': tmp}
        equations.update(definitions)

        variables = sum([e for k,e in self.symbols.items() if k != 'parameters'], [])
        table = "
Type	Equation	Residual
\n"

        for eq_type in equations:

            eq_lines = []
            for i in range(len(equations[eq_type])):
                eq = equations[eq_type][i]
                # if eq_type in ('expectation','direct_response'):
                #     vals = ''
                if eq_type not in ('arbitrage', 'transition', 'arbitrage_exp'):
                    vals = ''
                else:
                    val = resids[eq_type][i]
                    if abs(val) > 1e-8:
                        vals = '{:.4f}'.format(val)
                    else:
                        vals = '{:.3f}'.format(val)
                if '|' in eq:
                    # keep only lhs for now
                    eq, comp = str.split(eq,'|')
                lat = eq2tex(variables, eq)
                lat = '${}$'.format(lat)
                line = [lat, vals]
                h = eq_type if i==0 else ''
                fmt_line = '
{}	{}	{}
'.format(h, *line)
        #         print(fmt_line)
                eq_lines.append(fmt_line)
            table += str.join("\n", eq_lines)
        table = "
{}
".format(table)

        return table_infos + table

    @property
    def x_bounds(self):

        if 'controls_ub' in self.functions:
            fun_lb = self.functions['controls_lb']
            fun_ub = self.functions['controls_ub']
            return [fun_lb, fun_ub]
        else:
            return None

    def residuals(self, calib=None):

        if self.model_type == 'dtcscc':
            from dolo.algos.dtcscc.steady_state import residuals
        elif self.model_type == 'dtmscc':
            from dolo.algos.dtmscc.steady_state import residuals
        elif self.model_type == 'dynare':
            from dolo.algos.dynare.steady_state import residuals
        return residuals(self, calib)

    def eval_formula(self, expr, dataframe=None, calib=None):

        from dolo.compiler.eval_formula import eval_formula
        if calib is None:
            calib = self.calibration
        return eval_formula(expr, dataframe=dataframe, context=calib)

    def get_distribution(model, **opts):
        import copy
        gg = model.symbolic.options.get('distribution')
        if gg is None:
            raise Exception("Model has no distribution.")
        d = copy.deepcopy(gg)
        d.update(opts)
        if 'type' in d: d.pop('type')
        return d.eval(model.calibration.flat)

    def get_grid(model, **dis_opts):
        import copy
        gg = model.symbolic.options.get('grid')
        if gg is None:
            raise Exception("Model has no grid.")
        d = copy.deepcopy(gg)
        gtype = dis_opts.get('type')
        if gtype:
            from dolo.compiler.language import minilang
            try:
                cls = [e for e in minilang if e.__name__.lower()==gtype.lower()][0]
            except:
                raise Exception("Unknown grid type {}.".format(gtype))
            d = cls(**d)
    #     return cc
        d.update(**dis_opts)
        if 'type' in d: d.pop('type')
        return d.eval(d=model.calibration.flat)

    def __compile_functions__(self):

        from dolo.compiler.function_compiler import compile_function_ast
        from dolo.compiler.function_compiler import standard_function

        defs = self.symbolic.definitions

        model_type = self.model_type

        recipe = recipes[model_type]
        symbols = self.symbols # should match self.symbols

        comps = []

        functions = {}
        original_functions = {}
        original_gufunctions = {}

        # for funname in recipe['specs'].keys():
        for funname in self.symbolic.equations:

            spec = recipe['specs'][funname]

            eqs = self.symbolic.equations[funname]

            if spec.get('target'):
                target = spec['target']
                rhs_only = True
            else:
                target = None
                rhs_only = False

            if spec.get('complementarities'):

                # TODO: Rewrite and simplify
                comp_spec = spec.get('complementarities')
                comp_order = comp_spec['middle']
                comp_args = comp_spec['left-right']

                comps = []
                eqs = []
                for i,eq in enumerate(self.symbolic.equations[funname]):

                    if '|' in eq:
                        control = self.symbols[comp_order[0]][i]
                        eq, comp = str.split(eq,'|')
                        lhs, rhs = decode_complementarity(comp, control)
                        comps.append([lhs, rhs])
                    else:
                        comps.append(['-inf', 'inf'])

                    eqs.append(eq)

                    comp_lhs, comp_rhs = zip(*comps)
                    # fb_names = ['{}_lb'.format(funname), '{}_ub'.format(funname)]
                    fb_names = ['controls_lb'.format(funname), 'controls_ub'.format(funname)]

                    lower_bound, gu_lower_bound = compile_function_ast(comp_lhs, symbols, comp_args, funname=fb_names[0],definitions=defs)
                    upper_bound, gu_upper_bound = compile_function_ast(comp_rhs, symbols, comp_args, funname=fb_names[1],definitions=defs)

                    n_output = len(comp_lhs)

                    functions[fb_names[0]] = standard_function(gu_lower_bound, n_output )
                    functions[fb_names[1]] = standard_function(gu_upper_bound, n_output )
                    original_functions[fb_names[0]] = lower_bound
                    original_functions[fb_names[1]] = upper_bound
                    original_gufunctions[fb_names[0]] = gu_lower_bound
                    original_gufunctions[fb_names[1]] = gu_upper_bound

            arg_names = recipe['specs'][funname]['eqs']

            fun, gufun = compile_function_ast(eqs, symbols, arg_names, output_names=target, rhs_only=rhs_only, funname=funname, definitions=defs,                                   )
            n_output = len(eqs)
            original_functions[funname] = fun
            functions[funname] = standard_function(gufun, n_output )
            original_functions[funname] = fun
            original_gufunctions[funname] = gufun

        # temporary hack to keep auxiliaries
        auxiliaries = [k for k in defs.keys()]
        symbols['auxiliaries'] = auxiliaries
        eqs = ['{} = {}'.format(k,k) for k in auxiliaries]
        if model_type == 'dtcscc':
            arg_names = [('states',0,'s'),('controls',0,'x'),('parameters',0,'p')]
        else:
            arg_names = [('markov_states',0,'m'),('states',0,'s'),('controls',0,'x'),('parameters',0,'p')]
        target = ('auxiliaries',0,'y')
        rhs_only = True
        funname = 'auxiliary'
        fun, gufun = compile_function_ast(eqs, symbols, arg_names, output_names=target,
            rhs_only=rhs_only, funname=funname, definitions=defs)
        n_output = len(eqs)
        functions[funname] = standard_function(gufun, n_output )
        original_functions[funname] = fun
        original_gufunctions[funname] = gufun

        self.__original_functions__ = original_functions
        self.__original_gufunctions__ = original_gufunctions
        self.functions = functions

import re
regex = re.compile("(.*)<=(.*)<=(.*)")

def decode_complementarity(comp, control):

    '''
    # comp can be either:
    - None
    - "a<=expr" where a is a controls
    - "expr<=a" where a is a control
    - "expr1<=a<=expr2"
    '''

    try:
        res = regex.match(comp).groups()
    except:
        raise Exception("Unable to parse complementarity condition '{}'".format(comp))
    res = [r.strip() for r in res]
    if res[1] != control:
        msg = "Complementarity condition '{}' incorrect. Expected {} instead of {}.".format(comp, control, res[1])
        raise Exception(msg)
    return [res[0], res[2]]
<|endoftext|>


<commit_before>#!/usr/bin/env python

def ath(array):
    hash = {}
    for s in array:
        hash[s] = True
    return hash

knownHash = ath([

# bug 391976
"nsMathMLContainerFrame",
"nsMathMLmtableOuterFrame",
"nsMathMLmtdInnerFrame",
"nsMathMLmactionFrame",

# Bug 398462
"nsBaseAppShell",
"nsRunnable",

# Bug 403199
"nsSimpleNestedURI",

# Bug 467693
"nsStringBuffer"

])

# Things that are known to leak AND entrain smaller objects.
# If one of these leaks, leaks of small objects will not be reported.
knownLargeHash = ath([

# Bug 467008, bug 467647
"nsDocument",

# Bug 467686
"nsGlobalWindow",
"nsGenericElement",

# Bug 397206
"BackstagePass",

# Bug 102229 or bug 419562
"nsDNSService",

# Bug 467648
"nsHTMLDNSPrefetch",

# Bug 424418
"nsRDFResource",

# Bug 417630 and friends
"nsJVMManager",

# Bug 467873
"gfxTextRun"

])

# Large items that
# - should be reported even if things in knownLargeHash leak
# - should quell the reporting of smaller objects
otherLargeHash = ath([
"nsDocShell"
])

def amiss(logPrefix):
    currentFile = file(logPrefix + "-out", "r")
    sawLeakStats = False

    for line in currentFile:
        line = line.rstrip("\n")
        if (line == "== BloatView: ALL (cumulative) LEAK STATISTICS"):
            sawLeakStats = True
        # This line appears only if there are leaks
        if (line.endswith("Mean       StdDev")):
            break
    else:
        if sawLeakStats:
            #print "No leaks :)"
            pass
        else:
            print "Didn't see leak stats"
        currentFile.close()
        return False

    smallLeaks = ""
    largeKnownLeaks = ""
    largeOtherLeaks = ""

    for line in currentFile:
        line = line.strip("\x07").rstrip("\n").lstrip(" ")
        if (line == ""):
            break
        a = line.split(" ")[1]
        if a == "TOTAL":
            continue
        # print "Leaked at least one: " + a
        if a in knownLargeHash:
            largeKnownLeaks += "*** Leaked large object " + a + " (known)\n"
        if a in otherLargeHash:
            largeOtherLeaks += "*** Leaked large object " + a + "\n"
        if not a in knownHash:
            smallLeaks += a + "\n"

    if largeOtherLeaks != "":
        print "Leaked large objects:"
        print largeOtherLeaks
        # print "Also leaked 'known' large objects:"
        # print largeKnownLeaks
        currentFile.close()
        return True
    elif largeKnownLeaks != "":
        # print "(Known large leaks, and no other large leaks, so all leaks were ignored)"
        # print largeKnownLeaks
        currentFile.close()
        return False
    elif smallLeaks != "":
        print "Leaked:"
        print smallLeaks
        currentFile.close()
        return True
    else:
        # print "(Only known small leaks)"
        currentFile.close()
        return False

# print "detect_leaks is ready"
<commit_msg>Add 468160 to list of known large leaks<commit_after>#!/usr/bin/env python

def ath(array):
    hash = {}
    for s in array:
        hash[s] = True
    return hash

knownHash = ath([

# bug 391976
"nsMathMLContainerFrame",
"nsMathMLmtableOuterFrame",
"nsMathMLmtdInnerFrame",
"nsMathMLmactionFrame",

# Bug 398462
"nsBaseAppShell",
"nsRunnable",

# Bug 403199
"nsSimpleNestedURI",

# Bug 467693
"nsStringBuffer"

])

# Things that are known to leak AND entrain smaller objects.
# If one of these leaks, leaks of small objects will not be reported.
knownLargeHash = ath([

# Bug 467008, bug 467647
"nsDocument",

# Bug 467686
"nsGlobalWindow",
"nsGenericElement",

# Bug 397206
"BackstagePass",

# Bug 102229 or bug 419562
"nsDNSService",

# Bug 467648
"nsHTMLDNSPrefetch",

# Bug 424418
"nsRDFResource",

# Bug 417630 and friends
"nsJVMManager",

# Bug 467873
"gfxTextRun",

# Bug 468160
"imgRequest"

])

# Large items that
# - should be reported even if things in knownLargeHash leak
# - should quell the reporting of smaller objects
otherLargeHash = ath([
"nsDocShell"
])

def amiss(logPrefix):
    currentFile = file(logPrefix + "-out", "r")
    sawLeakStats = False

    for line in currentFile:
        line = line.rstrip("\n")
        if (line == "== BloatView: ALL (cumulative) LEAK STATISTICS"):
            sawLeakStats = True
        # This line appears only if there are leaks
        if (line.endswith("Mean       StdDev")):
            break
    else:
        if sawLeakStats:
            #print "No leaks :)"
            pass
        else:
            print "Didn't see leak stats"
        currentFile.close()
        return False

    smallLeaks = ""
    largeKnownLeaks = ""
    largeOtherLeaks = ""

    for line in currentFile:
        line = line.strip("\x07").rstrip("\n").lstrip(" ")
        if (line == ""):
            break
        a = line.split(" ")[1]
        if a == "TOTAL":
            continue
        # print "Leaked at least one: " + a
        if a in knownLargeHash:
            largeKnownLeaks += "*** Leaked large object " + a + " (known)\n"
        if a in otherLargeHash:
            largeOtherLeaks += "*** Leaked large object " + a + "\n"
        if not a in knownHash:
            smallLeaks += a + "\n"

    if largeOtherLeaks != "":
        print "Leaked large objects:"
        print largeOtherLeaks
        # print "Also leaked 'known' large objects:"
        # print largeKnownLeaks
        currentFile.close()
        return True
    elif largeKnownLeaks != "":
        # print "(Known large leaks, and no other large leaks, so all leaks were ignored)"
        # print largeKnownLeaks
        currentFile.close()
        return False
    elif smallLeaks != "":
        print "Leaked:"
        print smallLeaks
        currentFile.close()
        return True
    else:
        # print "(Only known small leaks)"
        currentFile.close()
        return False

# print "detect_leaks is ready"
<|endoftext|>


<commit_before>'''
===============
=== Purpose ===
===============

Stores flu-related search volume using private Google APIs (Trends and Health).
In theory, this data could be used to reconstruct Google Flu Trends.

The program checks (up to) all U.S. states, DC, and the entire US (52) over the
following inclusive epiweek range:
  From: 4 weeks prior to the most recently stored date
  To: the actual week at runtime

See: google_health_trends.py

=======================
=== Data Dictionary ===
=======================

`ght` is the table where Google health trends metadata is stored.
+----------+-------------+------+-----+---------+----------------+
| Field    | Type        | Null | Key | Default | Extra          |
+----------+-------------+------+-----+---------+----------------+
| id       | int(11)     | NO   | PRI | NULL    | auto_increment |
| query    | varchar(64) | NO   | MUL | NULL    |                |
| location | varchar(8)  | NO   | MUL | NULL    |                |
| epiweek  | int(11)     | NO   | MUL | NULL    |                |
| value    | float       | NO   |     | NULL    |                |
+----------+-------------+------+-----+---------+----------------+
id: unique identifier for each record
query: a search string or topic ID (see https://www.freebase.com/)
location: two-letter U.S. state abbreviation, or 'US' for entire US
epiweek: the epiweek during which the queries were executed
value: relative search volume; the exact definition is still unclear

=================
=== Changelog ===
=================

2017-02-16:
  * use secrets
2016-06-20:
  * updated list of terms (matching 2010-04-19 through 2015-05-05)
2016-06-13:
  * don't give up so easily (now 5 retries)
2016-06-10:
  + automatic retry on API failure, up to 3 retries
  + ability to update all terms in a single run
  * store values of zero
  * consolidate reports of missing values
2016-04-14:
  + first and last epiweek override
  * don't wrap terms in quotation marks
  * allow multiple locations
2016-01-29:
  + command line arguments for updating specific locations/terms
  - moved GHT class to google_health_trends.py
2015-12-04
  * not including quotes when storing to database
2015-12-03
  * fixed multiple-word queries (surround with quotes)
2015-12-01
  * Original version
'''

# standard library
import argparse
import time

# third party
import mysql.connector
from apiclient.discovery import build

# first party
from .google_health_trends import GHT
from .google_health_trends import NO_LOCATION_STR
import delphi.operations.secrets as secrets
from delphi.utils.epidate import EpiDate
import delphi.utils.epiweek as flu

# secret key for accessing Google's health trends APIs
# see: https://console.developers.google.com/apis/credentials?project=delphi-epi-trends
API_KEY = secrets.googletrends.apikey

# search strings and topics that are most highly correlated with wILI between
# 2010-04-19 and 2015-05-05
# see: https://www.google.com/trends/correlate
TERMS = [
  '/m/0cycc',
  'influenza type a',
  'flu duration',
  'flu fever',
  'treating flu',
  'fever flu',
  'flu recovery',
  'braun thermoscan',
  'oscillococcinum',
  'treating the flu',
  'cold or flu',
  'flu versus cold',
  'flu remedies',
  'contagious flu',
  'type a influenza',
  'flu or cold',
  'duration of flu',
  'cold versus flu',
  'flu cough',
  'flu headache',
  'thermoscan',
  'influenza incubation period',
  'flu lasts',
  'length of flu',
  'flu stomach',
  'cold vs flu',
  'flu and fever',
  'getting over the flu',
  'influenza a',
  'treatment for flu',
  'flu length',
  'treatment for the flu',
  'influenza symptoms',
  'over the counter flu',
  'flu complications',
  'cold and flu symptoms',
  'influenza incubation',
  'treatment of flu',
  'human temperature',
  'low body',
  'flu contagious',
  'robitussin ac',
  'flu how long',
  'ear thermometer',
  'flu contagious period',
  'treat flu',
  'cough flu',
  'low body temperature',
  'expectorant',
  'flu and cold',
  'rapid flu',
  'flu vs. cold',
  'how to treat the flu',
  'how long does the flu last?',
  'viral pneumonia',
  'flu in kids',
  'type a flu',
  'influenza treatment',
  'fighting the flu',
  'flu relief',
  'treat the flu',
  'flu medicine',
  'dangerous fever',
  'what is influenza',
  'tussin',
  'low body temp',
  'flu care',
  'flu in infants',
  'flu dizziness',
  'feed a fever',
  'flu vs cold',
  'flu vomiting',
  'bacterial pneumonia',
  'flu activity',
  'flu chills',
  'anas barbariae',
  'flu germs',
  'tylenol cold',
  'how to get over the flu',
  'flu in children',
  'influenza a and b',
  'duration of the flu',
  'cold symptoms',
  'flu report',
  'rapid flu test',
  'flu relapse',
  'get over the flu',
  'flu during pregnancy',
  'flu recovery time',
  'cure for flu',
  'tamiflu and breastfeeding',
  'flu chest pain',
  'flu treatment',
  'flu nausea',
  'remedies for the flu',
  'tamiflu in pregnancy',
  'side effects of tamiflu',
  'how to treat flu',
  'viral bronchitis',
  'flu how long contagious',
  'flu remedy',
]

# a list of all US states, including DC and the US as a whole
LOCATIONS = [
  'US',
  'AL',
  'AK',
  'AZ',
  'AR',
  'CA',
  'CO',
  'CT',
  'DC',
  'DE',
  'FL',
  'GA',
  'HI',
  'ID',
  'IL',
  'IN',
  'IA',
  'KS',
  'KY',
  'LA',
  'ME',
  'MD',
  'MA',
  'MI',
  'MN',
  'MS',
  'MO',
  'MT',
  'NE',
  'NV',
  'NH',
  'NJ',
  'NM',
  'NY',
  'NC',
  'ND',
  'OH',
  'OK',
  'OR',
  'PA',
  'RI',
  'SC',
  'SD',
  'TN',
  'TX',
  'UT',
  'VT',
  'VA',
  'WA',
  'WV',
  'WI',
  'WY',
]

def update(locations, terms, first=None, last=None, country='US'):
  # connect to the database
  u, p = secrets.db.epi
  cnx = mysql.connector.connect(user=u, password=p, database='epidata')
  cur = cnx.cursor()

  def get_num_rows():
    cur.execute('SELECT count(1) `num` FROM `ght`')
    for (num,) in cur:
      pass
    return num

  # check from 4 weeks preceeding the last week with data through this week
  cur.execute('SELECT max(`epiweek`) `ew0`, yearweek(now(), 6) `ew1` FROM `ght`')
  for (ew0, ew1) in cur:
    ew0 = 200401 if ew0 is None else flu.add_epiweeks(ew0, -4)
  ew0 = ew0 if first is None else first
  ew1 = ew1 if last is None else last
  print('Checking epiweeks between %d and %d...' % (ew0, ew1))

  # keep track of how many rows were added
  rows_before = get_num_rows()

  # check Google Trends for new and/or revised data
  sql = '''
    INSERT INTO
      `ght` (`query`, `location`, `epiweek`, `value`)
    VALUES
      (%s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE
      `value` = %s
  '''
  total_rows = 0
  ght = GHT(API_KEY)
  for term in terms:
    print(' [%s] using term' % term)
    for location in locations:
      try:
        #term2 = ('"%s"' % term) if ' ' in term else term
        term2 = term
        attempt = 0
        while True:
          attempt += 1
          try:
            result = ght.get_data(ew0, ew1, location, term2, country=country)
            break
          except Exception as ex:
            if attempt >= 5:
              raise ex
            else:
              delay = 2 ** attempt
              print(' [%s|%s] caught exception (will retry in %ds):' % (term, location, delay), ex)
              time.sleep(delay)
        values = [p['value'] for p in result['data']['lines'][0]['points']]
        ew = result['start_week']
        num_missing = 0
        for v in values:
          # Default SQL location value for US country for backwards compatibility
          # i.e. California's location is still stored as 'CA',
          # and having location == 'US' is still stored as 'US'
          sql_location = location if location != NO_LOCATION_STR else country

          # Change SQL location for non-US countries
          if country != 'US':
            # Underscore added to distinguish countries from 2-letter US states
            sql_location = country + "_"
            if location != NO_LOCATION_STR:
              sql_location = sql_location + location
          sql_data = (term, sql_location, ew, v, v)
          cur.execute(sql, sql_data)
          total_rows += 1
          if v == 0:
            num_missing += 1
            #print(' [%s|%s|%d] missing value' % (term, location, ew))
          ew = flu.add_epiweeks(ew, 1)
        if num_missing > 0:
          print(' [%s|%s] missing %d/%d value(s)' % (term, location, num_missing, len(values)))
      except Exception as ex:
        print(' [%s|%s] caught exception (will NOT retry):' % (term, location), ex)

  # keep track of how many rows were added
  rows_after = get_num_rows()
  print('Inserted %d/%d row(s)'%(rows_after - rows_before, total_rows))

  # cleanup
  cur.close()
  cnx.commit()
  cnx.close()

def main():
  # args and usage
  parser = argparse.ArgumentParser()
  parser.add_argument('location', action='store', type=str, default=None, help='location(s) (ex: all; US; TX; CA,LA,WY)')
  parser.add_argument('term', action='store', type=str, default=None, help='term/query/topic (ex: all; /m/0cycc; "flu fever")')
  parser.add_argument('--first', '-f', default=None, type=int, help='first epiweek override')
  parser.add_argument('--last', '-l', default=None, type=int, help='last epiweek override')
  parser.add_argument('--country', '-c', default='US', type=str, help='location country (ex: US; BR)')
  args = parser.parse_args()

  # sanity check
  first, last = args.first, args.last
  if first is not None:
    flu.check_epiweek(first)
  if last is not None:
    flu.check_epiweek(last)
  if first is not None and last is not None and first > last:
    raise Exception('epiweeks in the wrong order')

  # decide what to update
  if args.location.lower() == 'all':
    locations = LOCATIONS
  elif args.location.lower() == 'none':
    locations = [NO_LOCATION_STR]
  else:
    locations = args.location.upper().split(',')
  if args.term.lower() == 'all':
    terms = TERMS
  else:
    terms = [args.term]

  # country argument
  # Check that country follows ISO 1366 Alpha-2 code.
  # See https://www.iso.org/obp/ui/#search.
  if len(args.country) != 2:
    raise Exception('country name must be two letters (ISO 1366 Alpha-2)')
  country = args.country.upper()

  # run the update
  update(locations, terms, first, last, country)

if __name__ == '__main__':
  main()
<commit_msg>GHT update for multiple countries at once<commit_after>'''
===============
=== Purpose ===
===============

Stores flu-related search volume using private Google APIs (Trends and Health).
In theory, this data could be used to reconstruct Google Flu Trends.

The program checks (up to) all U.S. states, DC, and the entire US (52) over the
following inclusive epiweek range:
  From: 4 weeks prior to the most recently stored date
  To: the actual week at runtime

See: google_health_trends.py

=======================
=== Data Dictionary ===
=======================

`ght` is the table where Google health trends metadata is stored.
+----------+-------------+------+-----+---------+----------------+
| Field    | Type        | Null | Key | Default | Extra          |
+----------+-------------+------+-----+---------+----------------+
| id       | int(11)     | NO   | PRI | NULL    | auto_increment |
| query    | varchar(64) | NO   | MUL | NULL    |                |
| location | varchar(8)  | NO   | MUL | NULL    |                |
| epiweek  | int(11)     | NO   | MUL | NULL    |                |
| value    | float       | NO   |     | NULL    |                |
+----------+-------------+------+-----+---------+----------------+
id: unique identifier for each record
query: a search string or topic ID (see https://www.freebase.com/)
location: two-letter U.S. state abbreviation, or 'US' for entire US
epiweek: the epiweek during which the queries were executed
value: relative search volume; the exact definition is still unclear

=================
=== Changelog ===
=================

2017-02-16:
  * use secrets
2016-06-20:
  * updated list of terms (matching 2010-04-19 through 2015-05-05)
2016-06-13:
  * don't give up so easily (now 5 retries)
2016-06-10:
  + automatic retry on API failure, up to 3 retries
  + ability to update all terms in a single run
  * store values of zero
  * consolidate reports of missing values
2016-04-14:
  + first and last epiweek override
  * don't wrap terms in quotation marks
  * allow multiple locations
2016-01-29:
  + command line arguments for updating specific locations/terms
  - moved GHT class to google_health_trends.py
2015-12-04
  * not including quotes when storing to database
2015-12-03
  * fixed multiple-word queries (surround with quotes)
2015-12-01
  * Original version
'''

# standard library
import argparse
import time

# third party
import mysql.connector
from apiclient.discovery import build

# first party
from .google_health_trends import GHT
from .google_health_trends import NO_LOCATION_STR
import delphi.operations.secrets as secrets
from delphi.utils.epidate import EpiDate
import delphi.utils.epiweek as flu

# secret key for accessing Google's health trends APIs
# see: https://console.developers.google.com/apis/credentials?project=delphi-epi-trends
API_KEY = secrets.googletrends.apikey

# search strings and topics that are most highly correlated with wILI between
# 2010-04-19 and 2015-05-05
# see: https://www.google.com/trends/correlate
TERMS = [
  '/m/0cycc',
  'influenza type a',
  'flu duration',
  'flu fever',
  'treating flu',
  'fever flu',
  'flu recovery',
  'braun thermoscan',
  'oscillococcinum',
  'treating the flu',
  'cold or flu',
  'flu versus cold',
  'flu remedies',
  'contagious flu',
  'type a influenza',
  'flu or cold',
  'duration of flu',
  'cold versus flu',
  'flu cough',
  'flu headache',
  'thermoscan',
  'influenza incubation period',
  'flu lasts',
  'length of flu',
  'flu stomach',
  'cold vs flu',
  'flu and fever',
  'getting over the flu',
  'influenza a',
  'treatment for flu',
  'flu length',
  'treatment for the flu',
  'influenza symptoms',
  'over the counter flu',
  'flu complications',
  'cold and flu symptoms',
  'influenza incubation',
  'treatment of flu',
  'human temperature',
  'low body',
  'flu contagious',
  'robitussin ac',
  'flu how long',
  'ear thermometer',
  'flu contagious period',
  'treat flu',
  'cough flu',
  'low body temperature',
  'expectorant',
  'flu and cold',
  'rapid flu',
  'flu vs. cold',
  'how to treat the flu',
  'how long does the flu last?',
  'viral pneumonia',
  'flu in kids',
  'type a flu',
  'influenza treatment',
  'fighting the flu',
  'flu relief',
  'treat the flu',
  'flu medicine',
  'dangerous fever',
  'what is influenza',
  'tussin',
  'low body temp',
  'flu care',
  'flu in infants',
  'flu dizziness',
  'feed a fever',
  'flu vs cold',
  'flu vomiting',
  'bacterial pneumonia',
  'flu activity',
  'flu chills',
  'anas barbariae',
  'flu germs',
  'tylenol cold',
  'how to get over the flu',
  'flu in children',
  'influenza a and b',
  'duration of the flu',
  'cold symptoms',
  'flu report',
  'rapid flu test',
  'flu relapse',
  'get over the flu',
  'flu during pregnancy',
  'flu recovery time',
  'cure for flu',
  'tamiflu and breastfeeding',
  'flu chest pain',
  'flu treatment',
  'flu nausea',
  'remedies for the flu',
  'tamiflu in pregnancy',
  'side effects of tamiflu',
  'how to treat flu',
  'viral bronchitis',
  'flu how long contagious',
  'flu remedy',
]

# a list of all US states, including DC and the US as a whole
LOCATIONS = [
  'US',
  'AL',
  'AK',
  'AZ',
  'AR',
  'CA',
  'CO',
  'CT',
  'DC',
  'DE',
  'FL',
  'GA',
  'HI',
  'ID',
  'IL',
  'IN',
  'IA',
  'KS',
  'KY',
  'LA',
  'ME',
  'MD',
  'MA',
  'MI',
  'MN',
  'MS',
  'MO',
  'MT',
  'NE',
  'NV',
  'NH',
  'NJ',
  'NM',
  'NY',
  'NC',
  'ND',
  'OH',
  'OK',
  'OR',
  'PA',
  'RI',
  'SC',
  'SD',
  'TN',
  'TX',
  'UT',
  'VT',
  'VA',
  'WA',
  'WV',
  'WI',
  'WY',
]

def update(locations, terms, first=None, last=None, countries=['US']):
  # connect to the database
  u, p = secrets.db.epi
  cnx = mysql.connector.connect(user=u, password=p, database='epidata')
  cur = cnx.cursor()

  def get_num_rows():
    cur.execute('SELECT count(1) `num` FROM `ght`')
    for (num,) in cur:
      pass
    return num

  # check from 4 weeks preceeding the last week with data through this week
  cur.execute('SELECT max(`epiweek`) `ew0`, yearweek(now(), 6) `ew1` FROM `ght`')
  for (ew0, ew1) in cur:
    ew0 = 200401 if ew0 is None else flu.add_epiweeks(ew0, -4)
  ew0 = ew0 if first is None else first
  ew1 = ew1 if last is None else last
  print('Checking epiweeks between %d and %d...' % (ew0, ew1))

  # keep track of how many rows were added
  rows_before = get_num_rows()

  # check Google Trends for new and/or revised data
  sql = '''
    INSERT INTO
      `ght` (`query`, `location`, `epiweek`, `value`)
    VALUES
      (%s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE
      `value` = %s
  '''
  total_rows = 0
  ght = GHT(API_KEY)
  for term in terms:
    print(' [%s] using term' % term)
    ll, cl = len(locations), len(countries)
    for i in range(max(ll,cl)):
      location = locations[i] if i < ll else locations[0]
      country = countries[i] if i < cl else countries[0]
      try:
        #term2 = ('"%s"' % term) if ' ' in term else term
        term2 = term
        attempt = 0
        while True:
          attempt += 1
          try:
            result = ght.get_data(ew0, ew1, location, term2, country=country)
            break
          except Exception as ex:
            if attempt >= 5:
              raise ex
            else:
              delay = 2 ** attempt
              print(' [%s|%s] caught exception (will retry in %ds):' % (term, location, delay), ex)
              time.sleep(delay)
        values = [p['value'] for p in result['data']['lines'][0]['points']]
        ew = result['start_week']
        num_missing = 0
        for v in values:
          # Default SQL location value for US country for backwards compatibility
          # i.e. California's location is still stored as 'CA',
          # and having location == 'US' is still stored as 'US'
          sql_location = location if location != NO_LOCATION_STR else country

          # Change SQL location for non-US countries
          if country != 'US':
            # Underscore added to distinguish countries from 2-letter US states
            sql_location = country + "_"
            if location != NO_LOCATION_STR:
              sql_location = sql_location + location
          sql_data = (term, sql_location, ew, v, v)
          cur.execute(sql, sql_data)
          total_rows += 1
          if v == 0:
            num_missing += 1
            #print(' [%s|%s|%d] missing value' % (term, location, ew))
          ew = flu.add_epiweeks(ew, 1)
        if num_missing > 0:
          print(' [%s|%s] missing %d/%d value(s)' % (term, location, num_missing, len(values)))
      except Exception as ex:
        print(' [%s|%s] caught exception (will NOT retry):' % (term, location), ex)

  # keep track of how many rows were added
  rows_after = get_num_rows()
  print('Inserted %d/%d row(s)'%(rows_after - rows_before, total_rows))

  # cleanup
  cur.close()
  cnx.commit()
  cnx.close()

def main():
  # args and usage
  parser = argparse.ArgumentParser()
  parser.add_argument('location', action='store', type=str, default=None, help='location(s) (ex: all; US; TX; CA,LA,WY)')
  parser.add_argument('term', action='store', type=str, default=None, help='term/query/topic (ex: all; /m/0cycc; "flu fever")')
  parser.add_argument('--first', '-f', default=None, type=int, help='first epiweek override')
  parser.add_argument('--last', '-l', default=None, type=int, help='last epiweek override')
  parser.add_argument('--country', '-c', default='US', type=str, help='location country (ex: US; BR)')
  args = parser.parse_args()

  # sanity check
  first, last = args.first, args.last
  if first is not None:
    flu.check_epiweek(first)
  if last is not None:
    flu.check_epiweek(last)
  if first is not None and last is not None and first > last:
    raise Exception('epiweeks in the wrong order')

  # decide what to update
  if args.location.lower() == 'all':
    locations = LOCATIONS
  elif args.location.lower() == 'none':
    locations = [NO_LOCATION_STR]
  else:
    locations = args.location.upper().split(',')
  if args.term.lower() == 'all':
    terms = TERMS
  else:
    terms = [args.term]

  # country argument
  # Check that country follows ISO 1366 Alpha-2 code.
  # See https://www.iso.org/obp/ui/#search.
  countries = args.country.upper().split(',')
  if not all(map(lambda x: len(x) == 2, countries)):
    raise Exception('country name must be two letters (ISO 1366 Alpha-2)')

  # if length of locations and countries is > 1, need to be the same
  if len(locations) > 1 and len(countries) > 1 and len(locations) != len(countries):
    raise Exception('locations and countries must be length 1, or same length')

  # run the update
  update(locations, terms, first, last, countries)

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


<commit_before>"""
   Contains information regarding an ArcGIS Server Feature Server
"""
from __future__ import absolute_import
from __future__ import print_function
from re import search
from .._abstract.abstract import BaseAGSServer, BaseSecurityHandler
from ..security import security
from ..agol.services import FeatureLayer
import json
from ..common.geometry import SpatialReference
from ..common.general import FeatureSet
from ..common.filters import LayerDefinitionFilter, GeometryFilter, TimeFilter
########################################################################
class FeatureService(BaseAGSServer):
    """ contains information about a feature service """
    _url = None
    _currentVersion = None
    _serviceDescription = None
    _hasVersionedData = None
    _supportsDisconnectedEditing = None
    _hasStaticData = None
    _maxRecordCount = None
    _supportedQueryFormats = None
    _capabilities = None
    _description = None
    _copyrightText = None
    _spatialReference = None
    _initialExtent = None
    _fullExtent = None
    _allowGeometryUpdates = None
    _units = None
    _syncEnabled = None
    _syncCapabilities = None
    _editorTrackingInfo = None
    _documentInfo = None
    _layers = None
    _tables = None
    _enableZDefaults = None
    _zDefault = None
    _proxy_url = None
    _proxy_port = None
    _securityHandler = None
    _json = None
    _json_dict = None
    #----------------------------------------------------------------------
    def __init__(self, url, securityHandler=None,
                 initialize=False, proxy_url=None, proxy_port=None):
        """Constructor"""
        self._proxy_url = proxy_url
        self._proxy_port = proxy_port
        self._url = url
        if securityHandler is not None:
            self._securityHandler = securityHandler

        elif securityHandler is None:
            pass
        else:
            raise AttributeError("Invalid Security Handler")
        if not securityHandler is None and \
           hasattr(securityHandler, 'referer_url'):
            self._referer_url = securityHandler.referer_url
        if initialize:
            self.__init()
    #----------------------------------------------------------------------
    def __init(self):
        """ loads the data into the class """
        params = {"f": "json"}
        json_dict = self._get(self._url, params,
                                 securityHandler=self._securityHandler,
                                 proxy_port=self._proxy_port,
                                 proxy_url=self._proxy_url)
        self._json_dict = json_dict
        self._json = json.dumps(self._json_dict)
        attributes = [attr for attr in dir(self)
                      if not attr.startswith('__') and \
                      not attr.startswith('_')]
        for k,v in json_dict.items():
            if k in attributes:
                setattr(self, "_"+ k, v)
            else:
                print("%s - attribute not implemented for Feature Service." % k)
    #----------------------------------------------------------------------
    @property
    def administration(self):
        """returns the service admin object (if accessible)"""
        from ..manageags._services import AGSService
        url = self._url
        res = search("/rest/", url).span()
        addText = "/admin/"
        part1 = url[:res[1]].lower().replace('/rest/', '')
        part2 = url[res[1]:].lower().replace('/featureserver', ".mapserver")
        adminURL = "%s%s%s" % (part1, addText, part2)
        return AGSService(url=adminURL,
                          securityHandler=self._securityHandler,
                          proxy_url=self._proxy_url,
                          proxy_port=self._proxy_port,
                          initialize=False)
    #----------------------------------------------------------------------
    @property
    def itemInfo(self):
        """gets the item's info"""
        params = {"f" : "json"}
        url = self._url + "/info/iteminfo"
        return self._get(url=url, param_dict=params,
                            securityHandler=self._securityHandler,
                            proxy_url=self._proxy_url,
                            proxy_port=self._proxy_port)
    #----------------------------------------------------------------------
    def downloadThumbnail(self, outPath):
        """downloads the items's thumbnail"""
        url = self._url + "/info/thumbnail"
        params = {}
        return self._get(url=url,
                         out_folder=outPath,
                         securityHandler=self._securityHandler,
                         file_name=None,
                         param_dict=params,
                         proxy_url=self._proxy_url,
                         proxy_port=self._proxy_port)
    #----------------------------------------------------------------------
    def downloadMetadataFile(self, outPath):
        """downloads the metadata file to a given path"""
        fileName = "metadata.xml"
        url = self._url + "/info/metadata"
        params = {}
        return self._get(url=url,
                         out_folder=outPath,
                         file_name=fileName,
                         param_dict=params,
                         securityHandler=self._securityHandler,
                         proxy_url=self._proxy_url,
                         proxy_port=self._proxy_port)
    #----------------------------------------------------------------------
    def __str__(self):
        """returns object as a string"""
        if self._json is None:
            self.__init()
        return self._json
    #----------------------------------------------------------------------
    def __iter__(self):
        """returns the JSON response in key/value pairs"""
        if self._json_dict is None:
            self.__init()
        for k,v in self._json_dict.items():
            yield [k,v]
    #----------------------------------------------------------------------
    @property
    def securityHandler(self):
        """ gets the security handler """
        return self._securityHandler
    #----------------------------------------------------------------------
    @securityHandler.setter
    def securityHandler(self, value):
        """ sets the security handler """
        if isinstance(value, BaseSecurityHandler):
            if isinstance(value, security.AGSTokenSecurityHandler):
                self._securityHandler = value
            else:
                pass
        elif value is None:
            self._securityHandler = None
            self._token = None
    #----------------------------------------------------------------------
    @property
    def maxRecordCount(self):
        """returns the max record count"""
        if self._maxRecordCount is None:
            self.__init()
        return self._maxRecordCount
    #----------------------------------------------------------------------
    @property
    def supportedQueryFormats(self):
        """"""
        if self._supportedQueryFormats is None:
            self.__init()
        return self._supportedQueryFormats
    #----------------------------------------------------------------------
    @property
    def capabilities(self):
        """ returns a list of capabilities """
        if self._capabilities is None:
            self.__init()
        return self._capabilities
    #----------------------------------------------------------------------
    @property
    def description(self):
        """ returns the service description """
        if self._description is None:
            self.__init()
        return self._description
    #----------------------------------------------------------------------
    @property
    def copyrightText(self):
        """ returns the copyright text """
        if self._copyrightText is None:
            self.__init()
        return self._copyrightText
    #----------------------------------------------------------------------
    @property
    def spatialReference(self):
        """ returns the spatial reference """
        if self._spatialReference is None:
            self.__init()
        return self._spatialReference
    #----------------------------------------------------------------------
    @property
    def initialExtent(self):
        """ returns the initial extent of the feature service """
        if self._initialExtent is None:
            self.__init()
        return self._initialExtent
    #----------------------------------------------------------------------
    @property
    def fullExtent(self):
        """ returns the full extent of the feature service """
        if self._fullExtent is None:
            self.__init()
        return self._fullExtent
    #----------------------------------------------------------------------
    @property
    def allowGeometryUpdates(self):
        """ informs the user if the data allows geometry updates """
        if self._allowGeometryUpdates is None:
            self.__init()
        return self._allowGeometryUpdates
    #----------------------------------------------------------------------
    @property
    def units(self):
        """ returns the measurement unit """
        if self._units is None:
            self.__init()
        return self._units
    #----------------------------------------------------------------------
    @property
    def syncEnabled(self):
        """ informs the user if sync of data can be performed """
        if self._syncEnabled is None:
            self.__init()
        return self._syncEnabled
    #----------------------------------------------------------------------
    @property
    def syncCapabilities(self):
        """ type of sync that can be performed """
        if self._syncCapabilities is None:
            self.__init()
        return self._syncCapabilities
    #----------------------------------------------------------------------
    @property
    def editorTrackingInfo(self):
        """"""
        if self._editorTrackingInfo is None:
            self.__init()
        return self._editorTrackingInfo
    #----------------------------------------------------------------------
    @property
    def documentInfo(self):
        """"""
        if self._documentInfo is None:
            self.__init()
        return self._documentInfo
    #----------------------------------------------------------------------
    @property
    def layers(self):
        """ gets the layers for the feature service """
        if self._layers is None:
            self.__init()
        self._getLayers()
        return self._layers
    #----------------------------------------------------------------------
    def _getLayers(self):
        """ gets layers for the featuer service """
        params = {"f": "json"}

        json_dict = self._get(self._url, params,
                                 securityHandler=self._securityHandler,
                                 proxy_url=self._proxy_url,
                                 proxy_port=self._proxy_port)
        self._layers = []
        if 'layers' in json_dict:
            for l in json_dict["layers"]:
                self._layers.append(
                    FeatureLayer(url=self._url + "/%s" % l['id'],
                                       securityHandler=self._securityHandler,
                                       proxy_port=self._proxy_port,
                                       proxy_url=self._proxy_url)
                )
    #----------------------------------------------------------------------
    @property
    def tables(self):
        """lists the tables on the feature service"""
        if self._tables is None:
            self.__init()
        return self._tables
    #----------------------------------------------------------------------
    @property
    def enableZDefaults(self):
        """"""
        if self._enableZDefaults is None:
            self.__init()
        return self._enableZDefaults
    #----------------------------------------------------------------------
    @property
    def zDefault(self):
        """"""
        if self._zDefault is None:
            self.__init()
        return self._zDefault
    #----------------------------------------------------------------------
    @property
    def hasStaticData(self):
        """"""
        if self._hasStaticData is None:
            self.__init()
        return self._hasStaticData

    #----------------------------------------------------------------------
    @property
    def currentVersion(self):
        """ returns the map service current version """
        if self._currentVersion is None:
            self.__init()
        return self._currentVersion
    #----------------------------------------------------------------------
    @property
    def serviceDescription(self):
        """ returns the serviceDescription of the map service """
        if self._serviceDescription is None:
            self.__init()
        return self._serviceDescription
    #----------------------------------------------------------------------
    @property
    def hasVersionedData(self):
        """ returns boolean for versioned data """
        if self._hasVersionedData is None:
            self.__init()
        return self._hasVersionedData
    #----------------------------------------------------------------------
    @property
    def supportsDisconnectedEditing(self):
        """ returns boolean is disconnecting editted supported """
        if self._supportsDisconnectedEditing is None:
            self.__init()
        return self._supportsDisconnectedEditing
    #----------------------------------------------------------------------
    def query(self,
              layerDefsFilter=None,
              geometryFilter=None,
              timeFilter=None,
              returnGeometry=True,
              returnIdsOnly=False,
              returnCountOnly=False,
              returnZ=False,
              returnM=False,
              outSR=None
              ):
        """
           The Query operation is performed on a feature service resource
        """
        qurl = self._url + "/query"
        params = {"f": "json",
                  "returnGeometry": returnGeometry,
                  "returnIdsOnly": returnIdsOnly,
                  "returnCountOnly": returnCountOnly,
                  "returnZ": returnZ,
                  "returnM" : returnM}
        if not layerDefsFilter is None and \
           isinstance(layerDefsFilter, LayerDefinitionFilter):
            params['layerDefs'] = layerDefsFilter.filter
        if not geometryFilter is None and \
           isinstance(geometryFilter, GeometryFilter):
            gf = geometryFilter.filter
            params['geometryType'] = gf['geometryType']
            params['spatialRel'] = gf['spatialRel']
            params['geometry'] = gf['geometry']
            params['inSR'] = gf['inSR']
        if not outSR is None and \
           isinstance(outSR, SpatialReference):
            params['outSR'] = outSR.asDictionary
        if not timeFilter is None and \
           isinstance(timeFilter, TimeFilter):
            params['time'] = timeFilter.filter

        res = self._post(url=qurl,
                           param_dict=params,
                           securityHandler=self._securityHandler,
                           proxy_url=self._proxy_url,
                           proxy_port=self._proxy_port)
        if returnIdsOnly == False and returnCountOnly == False:
            if isinstance(res, str):
                jd = json.loads(res)
                return [FeatureSet.fromJSON(json.dumps(lyr)) for lyr in jd['layers']]
            elif isinstance(res, dict):
                return [FeatureSet.fromJSON(json.dumps(lyr)) for lyr in res['layers']]
            else:
                return res
        return res
<commit_msg>added missing properties for feature service.<commit_after>"""
   Contains information regarding an ArcGIS Server Feature Server
"""
from __future__ import absolute_import
from __future__ import print_function
from re import search
from .._abstract.abstract import BaseAGSServer, BaseSecurityHandler
from ..security import security
from ..agol.services import FeatureLayer
import json
from ..common.geometry import SpatialReference
from ..common.general import FeatureSet
from ..common.filters import LayerDefinitionFilter, GeometryFilter, TimeFilter
########################################################################
class FeatureService(BaseAGSServer):
    """ contains information about a feature service """
    _url = None
    _serviceItemId = None
    _supportsApplyEditsWithGlobalIds = None
    _currentVersion = None
    _serviceDescription = None
    _hasVersionedData = None
    _supportsDisconnectedEditing = None
    _hasStaticData = None
    _maxRecordCount = None
    _supportedQueryFormats = None
    _capabilities = None
    _description = None
    _copyrightText = None
    _spatialReference = None
    _initialExtent = None
    _fullExtent = None
    _allowGeometryUpdates = None
    _units = None
    _syncEnabled = None
    _syncCapabilities = None
    _editorTrackingInfo = None
    _documentInfo = None
    _layers = None
    _tables = None
    _enableZDefaults = None
    _zDefault = None
    _proxy_url = None
    _proxy_port = None
    _securityHandler = None
    _json = None
    _json_dict = None
    #----------------------------------------------------------------------
    def __init__(self, url, securityHandler=None,
                 initialize=False, proxy_url=None, proxy_port=None):
        """Constructor"""
        self._proxy_url = proxy_url
        self._proxy_port = proxy_port
        self._url = url
        if securityHandler is not None:
            self._securityHandler = securityHandler

        elif securityHandler is None:
            pass
        else:
            raise AttributeError("Invalid Security Handler")
        if not securityHandler is None and \
           hasattr(securityHandler, 'referer_url'):
            self._referer_url = securityHandler.referer_url
        if initialize:
            self.__init()
    #----------------------------------------------------------------------
    def __init(self):
        """ loads the data into the class """
        params = {"f": "json"}
        json_dict = self._get(self._url, params,
                                 securityHandler=self._securityHandler,
                                 proxy_port=self._proxy_port,
                                 proxy_url=self._proxy_url)
        self._json_dict = json_dict
        self._json = json.dumps(self._json_dict)
        attributes = [attr for attr in dir(self)
                      if not attr.startswith('__') and \
                      not attr.startswith('_')]
        for k,v in json_dict.items():
            if k in attributes:
                setattr(self, "_"+ k, v)
            else:
                print("%s - attribute not implemented for Feature Service." % k)
    #----------------------------------------------------------------------
    @property
    def administration(self):
        """returns the service admin object (if accessible)"""
        from ..manageags._services import AGSService
        url = self._url
        res = search("/rest/", url).span()
        addText = "/admin/"
        part1 = url[:res[1]].lower().replace('/rest/', '')
        part2 = url[res[1]:].lower().replace('/featureserver', ".mapserver")
        adminURL = "%s%s%s" % (part1, addText, part2)
        return AGSService(url=adminURL,
                          securityHandler=self._securityHandler,
                          proxy_url=self._proxy_url,
                          proxy_port=self._proxy_port,
                          initialize=False)
    #----------------------------------------------------------------------
    @property
    def itemInfo(self):
        """gets the item's info"""
        params = {"f" : "json"}
        url = self._url + "/info/iteminfo"
        return self._get(url=url, param_dict=params,
                            securityHandler=self._securityHandler,
                            proxy_url=self._proxy_url,
                            proxy_port=self._proxy_port)
    #----------------------------------------------------------------------
    def downloadThumbnail(self, outPath):
        """downloads the items's thumbnail"""
        url = self._url + "/info/thumbnail"
        params = {}
        return self._get(url=url,
                         out_folder=outPath,
                         securityHandler=self._securityHandler,
                         file_name=None,
                         param_dict=params,
                         proxy_url=self._proxy_url,
                         proxy_port=self._proxy_port)
    #----------------------------------------------------------------------
    def downloadMetadataFile(self, outPath):
        """downloads the metadata file to a given path"""
        fileName = "metadata.xml"
        url = self._url + "/info/metadata"
        params = {}
        return self._get(url=url,
                         out_folder=outPath,
                         file_name=fileName,
                         param_dict=params,
                         securityHandler=self._securityHandler,
                         proxy_url=self._proxy_url,
                         proxy_port=self._proxy_port)
    #----------------------------------------------------------------------
    def __str__(self):
        """returns object as a string"""
        if self._json is None:
            self.__init()
        return self._json
    #----------------------------------------------------------------------
    def __iter__(self):
        """returns the JSON response in key/value pairs"""
        if self._json_dict is None:
            self.__init()
        for k,v in self._json_dict.items():
            yield [k,v]
    #----------------------------------------------------------------------
    @property
    def securityHandler(self):
        """ gets the security handler """
        return self._securityHandler
    #----------------------------------------------------------------------
    @securityHandler.setter
    def securityHandler(self, value):
        """ sets the security handler """
        if isinstance(value, BaseSecurityHandler):
            if isinstance(value, security.AGSTokenSecurityHandler):
                self._securityHandler = value
            else:
                pass
        elif value is None:
            self._securityHandler = None
            self._token = None
    #----------------------------------------------------------------------
    @property
    def maxRecordCount(self):
        """returns the max record count"""
        if self._maxRecordCount is None:
            self.__init()
        return self._maxRecordCount
    #----------------------------------------------------------------------
    @property
    def supportedQueryFormats(self):
        """"""
        if self._supportedQueryFormats is None:
            self.__init()
        return self._supportedQueryFormats
    #----------------------------------------------------------------------
    @property
    def capabilities(self):
        """ returns a list of capabilities """
        if self._capabilities is None:
            self.__init()
        return self._capabilities
    #----------------------------------------------------------------------
    @property
    def description(self):
        """ returns the service description """
        if self._description is None:
            self.__init()
        return self._description
    #----------------------------------------------------------------------
    @property
    def copyrightText(self):
        """ returns the copyright text """
        if self._copyrightText is None:
            self.__init()
        return self._copyrightText
    #----------------------------------------------------------------------
    @property
    def spatialReference(self):
        """ returns the spatial reference """
        if self._spatialReference is None:
            self.__init()
        return self._spatialReference
    #----------------------------------------------------------------------
    @property
    def initialExtent(self):
        """ returns the initial extent of the feature service """
        if self._initialExtent is None:
            self.__init()
        return self._initialExtent
    #----------------------------------------------------------------------
    @property
    def fullExtent(self):
        """ returns the full extent of the feature service """
        if self._fullExtent is None:
            self.__init()
        return self._fullExtent
    #----------------------------------------------------------------------
    @property
    def allowGeometryUpdates(self):
        """ informs the user if the data allows geometry updates """
        if self._allowGeometryUpdates is None:
            self.__init()
        return self._allowGeometryUpdates
    #----------------------------------------------------------------------
    @property
    def units(self):
        """ returns the measurement unit """
        if self._units is None:
            self.__init()
        return self._units
    #----------------------------------------------------------------------
    @property
    def syncEnabled(self):
        """ informs the user if sync of data can be performed """
        if self._syncEnabled is None:
            self.__init()
        return self._syncEnabled
    #----------------------------------------------------------------------
    @property
    def syncCapabilities(self):
        """ type of sync that can be performed """
        if self._syncCapabilities is None:
            self.__init()
        return self._syncCapabilities
    #----------------------------------------------------------------------
    @property
    def editorTrackingInfo(self):
        """"""
        if self._editorTrackingInfo is None:
            self.__init()
        return self._editorTrackingInfo
    #----------------------------------------------------------------------
    @property
    def documentInfo(self):
        """"""
        if self._documentInfo is None:
            self.__init()
        return self._documentInfo
    #----------------------------------------------------------------------
    @property
    def layers(self):
        """ gets the layers for the feature service """
        if self._layers is None:
            self.__init()
        self._getLayers()
        return self._layers
    #----------------------------------------------------------------------
    def _getLayers(self):
        """ gets layers for the featuer service """
        params = {"f": "json"}

        json_dict = self._get(self._url, params,
                                 securityHandler=self._securityHandler,
                                 proxy_url=self._proxy_url,
                                 proxy_port=self._proxy_port)
        self._layers = []
        if 'layers' in json_dict:
            for l in json_dict["layers"]:
                self._layers.append(
                    FeatureLayer(url=self._url + "/%s" % l['id'],
                                       securityHandler=self._securityHandler,
                                       proxy_port=self._proxy_port,
                                       proxy_url=self._proxy_url)
                )
    #----------------------------------------------------------------------
    @property
    def tables(self):
        """lists the tables on the feature service"""
        if self._tables is None:
            self.__init()
        return self._tables
    #----------------------------------------------------------------------
    @property
    def enableZDefaults(self):
        """"""
        if self._enableZDefaults is None:
            self.__init()
        return self._enableZDefaults
    #----------------------------------------------------------------------
    @property
    def zDefault(self):
        """"""
        if self._zDefault is None:
            self.__init()
        return self._zDefault
    #----------------------------------------------------------------------
    @property
    def hasStaticData(self):
        """"""
        if self._hasStaticData is None:
            self.__init()
        return self._hasStaticData
    @property
    def supportsApplyEditsWithGlobalIds(self):
        """returns the supportsApplyEditsWithGlobalIds value"""
        if self._supportsApplyEditsWithGlobalIds is None:
            self.__init()
        return self._supportsApplyEditsWithGlobalIds
    @property
    def serviceItemId(self):
        """returns the serviceItemId value"""
        if self._serviceItemId is None:
            self.__init()
        return self._serviceItemId
    #----------------------------------------------------------------------
    @property
    def currentVersion(self):
        """ returns the map service current version """
        if self._currentVersion is None:
            self.__init()
        return self._currentVersion
    #----------------------------------------------------------------------
    @property
    def serviceDescription(self):
        """ returns the serviceDescription of the map service """
        if self._serviceDescription is None:
            self.__init()
        return self._serviceDescription
    #----------------------------------------------------------------------
    @property
    def hasVersionedData(self):
        """ returns boolean for versioned data """
        if self._hasVersionedData is None:
            self.__init()
        return self._hasVersionedData
    #----------------------------------------------------------------------
    @property
    def supportsDisconnectedEditing(self):
        """ returns boolean is disconnecting editted supported """
        if self._supportsDisconnectedEditing is None:
            self.__init()
        return self._supportsDisconnectedEditing
    #----------------------------------------------------------------------
    def query(self,
              layerDefsFilter=None,
              geometryFilter=None,
              timeFilter=None,
              returnGeometry=True,
              returnIdsOnly=False,
              returnCountOnly=False,
              returnZ=False,
              returnM=False,
              outSR=None
              ):
        """
           The Query operation is performed on a feature service resource
        """
        qurl = self._url + "/query"
        params = {"f": "json",
                  "returnGeometry": returnGeometry,
                  "returnIdsOnly": returnIdsOnly,
                  "returnCountOnly": returnCountOnly,
                  "returnZ": returnZ,
                  "returnM" : returnM}
        if not layerDefsFilter is None and \
           isinstance(layerDefsFilter, LayerDefinitionFilter):
            params['layerDefs'] = layerDefsFilter.filter
        if not geometryFilter is None and \
           isinstance(geometryFilter, GeometryFilter):
            gf = geometryFilter.filter
            params['geometryType'] = gf['geometryType']
            params['spatialRel'] = gf['spatialRel']
            params['geometry'] = gf['geometry']
            params['inSR'] = gf['inSR']
        if not outSR is None and \
           isinstance(outSR, SpatialReference):
            params['outSR'] = outSR.asDictionary
        if not timeFilter is None and \
           isinstance(timeFilter, TimeFilter):
            params['time'] = timeFilter.filter

        res = self._post(url=qurl,
                           param_dict=params,
                           securityHandler=self._securityHandler,
                           proxy_url=self._proxy_url,
                           proxy_port=self._proxy_port)
        if returnIdsOnly == False and returnCountOnly == False:
            if isinstance(res, str):
                jd = json.loads(res)
                return [FeatureSet.fromJSON(json.dumps(lyr)) for lyr in jd['layers']]
            elif isinstance(res, dict):
                return [FeatureSet.fromJSON(json.dumps(lyr)) for lyr in res['layers']]
            else:
                return res
        return res
<|endoftext|>


<commit_before># Copyright 2016 Raytheon BBN Technologies
#
# 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

__all__ = ['ElementwiseFilter']

import queue
import h5py
import itertools
import numpy as np
import os.path
import time

from auspex.stream import InputConnector, OutputConnector
from auspex.log import logger
import auspex.config as config
from .filter import Filter

class ElementwiseFilter(Filter):
    """Perform elementwise operations on multiple streams:
    e.g. multiply or add all streams element-by-element"""

    sink        = InputConnector()
    source      = OutputConnector()
    filter_name = "GenericElementwise" # To identify subclasses when naming data streams

    def __init__(self, filter_name=None, **kwargs):
        super(ElementwiseFilter, self).__init__(filter_name=filter_name, **kwargs)
        self.sink.max_input_streams = 100
        self.quince_parameters = []

    def operation(self):
        """Must be overridden with the desired mathematical function"""
        pass

    def unit(self, base_unit):
        """Must be overridden accoriding the desired mathematical function
        e.g. return base_unit + "^{}".format(len(self.sink.input_streams))"""
        pass

    def update_descriptors(self):
        """Must be overridden depending on the desired mathematical function"""
        logger.debug('Updating %s "%s" descriptors based on input descriptor: %s.', self.filter_name, self.filter_name, self.sink.descriptor)

        # Sometimes not all of the input descriptors have been updated... pause here until they are:
        if None in [ss.descriptor for ss in self.sink.input_streams]:
            logger.debug('%s "%s" waiting for all input streams to be updated.', self.filter_name, self.name)
            return

        self.descriptor = self.sink.descriptor.copy()
        if self.filter_name:
            self.descriptor.data_name = self.filter_name
        if self.descriptor.unit:
            self.descriptor.unit = self.descriptor.unit + "^{}".format(len(self.sink.input_streams))
        self.source.descriptor = self.descriptor
        self.source.update_descriptors()

    def main(self):
        self.done.clear()
        streams = self.sink.input_streams

        for s in streams[1:]:
            if not np.all(s.descriptor.expected_tuples() == streams[0].descriptor.expected_tuples()):
                raise ValueError("Multiple streams connected to correlator must have matching descriptors.")

        # Buffers for stream data
        stream_data = {s: np.zeros(0, dtype=self.sink.descriptor.dtype) for s in streams}

        # Store whether streams are done
        streams_done      = {s: False for s in streams}
        points_per_stream = {s: 0 for s in streams}

        while not self.exit.is_set():

            # Try to pull all messages in the queue. queue.empty() is not reliable, so we
            # ask for forgiveness rather than permission.
            msgs_by_stream = {s: [] for s in streams}

            for stream in streams[::-1]:
                while not self.exit.is_set():
                    try:
                        msgs_by_stream[stream].append(stream.queue.get(False))
                    except queue.Empty as e:
                        time.sleep(0.002)
                        break

            # Process many messages for each stream
            for stream, messages in msgs_by_stream.items():
                for message in messages:
                    message_type = message['type']
                    message_data = message['data']
                    message_data = message_data if hasattr(message_data, 'size') else np.array([message_data])
                    if message_type == 'event':
                        if message['event_type'] == 'done':
                            streams_done[stream] = True
                        elif message['event_type'] == 'refine':
                            logger.warning("ElementwiseFilter doesn't handle refinement yet!")
                        self.push_to_all(message)
                    elif message_type == 'data':
                        # Add any old data...
                        points_per_stream[stream] += len(message_data.flatten())
                        stream_data[stream] = np.concatenate((stream_data[stream], message_data.flatten()))

            # Now process the data with the elementwise operation
            smallest_length = min([d.size for d in stream_data.values()])
            new_data = [d[:smallest_length] for d in stream_data.values()]
            result = new_data[0]
            for nd in new_data[1:]:
                result = self.operation()(result, nd)
            if result.size > 0:
                self.source.push(result)

            # Add data to carry_data if necessary
            for stream in stream_data.keys():
                if stream_data[stream].size > smallest_length:
                    stream_data[stream] = stream_data[stream][smallest_length:]
                else:
                    stream_data[stream] = np.zeros(0, dtype=self.sink.descriptor.dtype)

            # If the amount of data processed is equal to the num points in the stream, we are done
            if np.all([streams_done[stream] for stream in streams]):
                self.done.set()
                break
<commit_msg>h5py no longer necessary<commit_after># Copyright 2016 Raytheon BBN Technologies
#
# 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

__all__ = ['ElementwiseFilter']

import queue
import itertools
import numpy as np
import os.path
import time

from auspex.stream import InputConnector, OutputConnector
from auspex.log import logger
import auspex.config as config
from .filter import Filter

class ElementwiseFilter(Filter):
    """Perform elementwise operations on multiple streams:
    e.g. multiply or add all streams element-by-element"""

    sink        = InputConnector()
    source      = OutputConnector()
    filter_name = "GenericElementwise" # To identify subclasses when naming data streams

    def __init__(self, filter_name=None, **kwargs):
        super(ElementwiseFilter, self).__init__(filter_name=filter_name, **kwargs)
        self.sink.max_input_streams = 100
        self.quince_parameters = []

    def operation(self):
        """Must be overridden with the desired mathematical function"""
        pass

    def unit(self, base_unit):
        """Must be overridden accoriding the desired mathematical function
        e.g. return base_unit + "^{}".format(len(self.sink.input_streams))"""
        pass

    def update_descriptors(self):
        """Must be overridden depending on the desired mathematical function"""
        logger.debug('Updating %s "%s" descriptors based on input descriptor: %s.', self.filter_name, self.filter_name, self.sink.descriptor)

        # Sometimes not all of the input descriptors have been updated... pause here until they are:
        if None in [ss.descriptor for ss in self.sink.input_streams]:
            logger.debug('%s "%s" waiting for all input streams to be updated.', self.filter_name, self.name)
            return

        self.descriptor = self.sink.descriptor.copy()
        if self.filter_name:
            self.descriptor.data_name = self.filter_name
        if self.descriptor.unit:
            self.descriptor.unit = self.descriptor.unit + "^{}".format(len(self.sink.input_streams))
        self.source.descriptor = self.descriptor
        self.source.update_descriptors()

    def main(self):
        self.done.clear()
        streams = self.sink.input_streams

        for s in streams[1:]:
            if not np.all(s.descriptor.expected_tuples() == streams[0].descriptor.expected_tuples()):
                raise ValueError("Multiple streams connected to correlator must have matching descriptors.")

        # Buffers for stream data
        stream_data = {s: np.zeros(0, dtype=self.sink.descriptor.dtype) for s in streams}

        # Store whether streams are done
        streams_done      = {s: False for s in streams}
        points_per_stream = {s: 0 for s in streams}

        while not self.exit.is_set():

            # Try to pull all messages in the queue. queue.empty() is not reliable, so we
            # ask for forgiveness rather than permission.
            msgs_by_stream = {s: [] for s in streams}

            for stream in streams[::-1]:
                while not self.exit.is_set():
                    try:
                        msgs_by_stream[stream].append(stream.queue.get(False))
                    except queue.Empty as e:
                        time.sleep(0.002)
                        break

            # Process many messages for each stream
            for stream, messages in msgs_by_stream.items():
                for message in messages:
                    message_type = message['type']
                    message_data = message['data']
                    message_data = message_data if hasattr(message_data, 'size') else np.array([message_data])
                    if message_type == 'event':
                        if message['event_type'] == 'done':
                            streams_done[stream] = True
                        elif message['event_type'] == 'refine':
                            logger.warning("ElementwiseFilter doesn't handle refinement yet!")
                        self.push_to_all(message)
                    elif message_type == 'data':
                        # Add any old data...
                        points_per_stream[stream] += len(message_data.flatten())
                        stream_data[stream] = np.concatenate((stream_data[stream], message_data.flatten()))

            # Now process the data with the elementwise operation
            smallest_length = min([d.size for d in stream_data.values()])
            new_data = [d[:smallest_length] for d in stream_data.values()]
            result = new_data[0]
            for nd in new_data[1:]:
                result = self.operation()(result, nd)
            if result.size > 0:
                self.source.push(result)

            # Add data to carry_data if necessary
            for stream in stream_data.keys():
                if stream_data[stream].size > smallest_length:
                    stream_data[stream] = stream_data[stream][smallest_length:]
                else:
                    stream_data[stream] = np.zeros(0, dtype=self.sink.descriptor.dtype)

            # If the amount of data processed is equal to the num points in the stream, we are done
            if np.all([streams_done[stream] for stream in streams]):
                self.done.set()
                break
<|endoftext|>