<commit_before>import pysb
import sympy
from re import sub
import warnings

class KappaGenerator(object):

    # Dialect can be either 'complx' or 'kasim' (default)
    def __init__(self, model, dialect='kasim'):
        self.model = model
        self.__content = None
        self.dialect = dialect

    def get_content(self):
        if self.__content == None:
            self.generate_content()
        return self.__content

    def generate_content(self):
        self.__content = ''
        #self.generate_compartments()

        # Agent declarations appear to be required in kasim
        # but prohibited in complx
        if (self.dialect == 'kasim'):
            self.generate_molecule_types()
            # Parameters, variables, and expressions are allowed in kasim
            self.generate_parameters()

        self.generate_reaction_rules()
        self.generate_observables()
        self.generate_species()

    def generate_parameters(self):
        names = [x.name for x in self.model.parameters] + \
                [x.name for x in self.model.expressions]
        for p in self.model.parameters:
            self.__content += ("%%var: \'%s\' %e\n") % (p.name, p.value)
        for e in self.model.expressions:
            sym_names = [x.name for x in e.expr.atoms(sympy.Symbol)]
            str_expr = str(sympy_to_muparser(e.expr))
            for n in sym_names:
                str_expr = str_expr.replace(n,"\'%s\'"%n)
            self.__content += ("%%var: \'%s\' %s\n") % \
                              (e.name, str_expr.replace('**','^'))
        self.__content += "\n"

    #def generate_compartments(self):
    #    self.__content += "begin compartments\n"
    #    for c in self.model.compartments:
    #        if c.parent is None:
    #            parent_name = ''
    #        else:
    #            parent_name = c.parent.name
    #        self.__content += ("  %s  %d  %f  %s\n") % \
    #                          (c.name, c.dimension, c.size, parent_name)
    #    self.__content += "end compartments\n\n"

    def generate_molecule_types(self):
        for m in self.model.monomers:
            site_code = ','.join([format_monomer_site(m, s) for s in m.sites])
            self.__content += "%%agent: %s(%s)\n" % (m.name, site_code)
        self.__content += "\n"

    def generate_reaction_rules(self):
        # +1 for the colon
        #max_length = max(len(r.name) for r in self.model.rules) + 1
        for r in self.model.rules:
            label = '\'' + r.name + '\''
            reactants_code = format_reactionpattern(r.reactant_pattern)
            products_code  = format_reactionpattern(r.product_pattern)
            arrow = '->'

            if isinstance(r.rate_forward,pysb.core.Expression) and \
               self.dialect == 'complx':
                raise KappaException("Expressions are not supported by complx.")
            # Get the rate code depending on the dialect
            if self.dialect == 'kasim':
                f_rate_code = "\'" + r.rate_forward.name + "\'"
            else:
                f_rate_code = float(r.rate_forward.value)

            self.__content += ("%s %s %s %s @ %s") % \
                (label, reactants_code, arrow, products_code, f_rate_code)
            self.__content += "\n"

            # Add the reverse reaction
            if r.is_reversible:
                if isinstance(r.rate_reverse,pysb.core.Expression) and \
                   self.dialect == 'complx':
                    raise KappaException("Expressions are not supported by "
                                         "complx.")
                # Get the rate code depending on the dialect
                if self.dialect == 'kasim':
                    r_rate_code = "\'" + r.rate_reverse.name + "\'"
                else:
                    r_rate_code = float(r.rate_reverse.value)

                label = '\'' + r.name + '_rev' + '\''
                self.__content += ("%s %s %s %s @ %s") % \
                      (label, products_code, arrow, reactants_code, r_rate_code)
                self.__content += "\n"

        self.__content += "\n"

    def generate_observables(self):
        if not self.model.observables:
            return
        for obs in self.model.observables:
            name = '\'' + obs.name + '\''
            observable_code = format_reactionpattern(obs.reaction_pattern)
            # In the near future (KaSim 4.0), the observable syntax will
            # require pipe characters around the expression. However, for the
            # time being we'll stick with the old syntax for backwards
            # compatibility
            #if self.dialect == 'kasim':
            #    self.__content += ("%%obs: %s |%s|\n") % \
            #                      (name, observable_code)
            #else:
            #    self.__content += ("%%obs: %s %s\n") % (name, observable_code)
            self.__content += ("%%obs: %s %s\n") % (name, observable_code)

        self.__content += "\n"

    def generate_species(self):
        if not self.model.initial_conditions:
            warnings.warn("Warning: No initial conditions.")

        species_codes = [format_complexpattern(cp)
                         for cp, param in self.model.initial_conditions]
        #max_length = max(len(code) for code in species_codes)
        for i, code in enumerate(species_codes):
            param = self.model.initial_conditions[i][1]
            #self.__content += ("%%init:  %-" + str(max_length) + \
            #                  "s   %s\n") % (code, param.name)
            if (self.dialect == 'kasim'):
                # Switched from %g (float) to %d (int) because kappa didn't
                # like scientific notation for large integers
                self.__content += ("%%init: \'%s\' %s\n") % (param.name, code)
                #self.__content += ("%%init: %10g %s\n") % (param.value, code)
            else:
                if isinstance(param,pysb.core.Expression):
                    raise KappaException("complx does not support Expressions.")
                # Switched from %g (float) to %d (int) because kappa didn't
                # like scientific notation for large integers
                self.__content += ("%%init: %10d * %s\n") % \
                                  (float(param.value), code)
                #self.__content += ("%%init: %10g * %s\n") % (param.value, code)

        self.__content += "\n"

def format_monomer_site(monomer, site):
    ret = site
    if site in monomer.site_states:
        for state in monomer.site_states[site]:
            ret += '~' + state
    return ret

def format_reactionpattern(rp):
    return ','.join([format_complexpattern(cp) for cp in rp.complex_patterns])

def format_complexpattern(cp):
    ret = ','.join([format_monomerpattern(mp) for mp in cp.monomer_patterns])
    if cp.compartment is not None:
        ret = '@%s:%s' % (cp.compartment.name, ret)
    return ret

def format_monomerpattern(mp):
    # sort sites in the same order given in the original Monomer
    site_conditions = sorted(mp.site_conditions.items(),
                             key=lambda x: mp.monomer.sites.index(x[0]))
    site_pattern_code = ','.join([format_site_condition(site, state)
                                  for (site, state) in site_conditions])
    ret = '%s(%s)' % (mp.monomer.name, site_pattern_code)
    if mp.compartment is not None:
        ret = '%s@%s' % (ret, mp.compartment.name)
    return ret

def format_site_condition(site, state):
    # If the state/bond is unspecified
    if state == None:
        state_code = ''
    # If there is a bond number
    elif type(state) == int:
        state_code = '!' + str(state)
    # If there is a lists of bonds to the site (not supported by Kappa)
    elif type(state) == list:
        raise KappaException("Kappa generator does not support multiple bonds "
                              "to a single site.")
    # Site with state
    elif type(state) == str:
        state_code = '~' + state
    # Site with state and a bond
    elif type(state) == tuple:
        # If the bond is ANY
        if state[1] == pysb.ANY:
            state_code = '~%s!_' % state[0]
        # If the bond is WILD
        elif state[1] == pysb.WILD:
            state_code = '~%s?' % state[0]
        # If the bond is a number
        elif type(state[1]) == int:
            state_code = '~%s!%s' % state
        # If it's something else, raise an Exception
        else:
            raise KappaException("Kappa generator has encountered an unknown "
                                 "element in a site state/bond tuple: (%s, %s)"
                                 % state)
    # Site bound to ANY
    elif state == pysb.ANY:
        state_code = '!_'
    # Site bound to WILD
    elif state == pysb.WILD:
        state_code = '?'
    # Something else
    else:
        raise KappaException("Kappa generator has encountered an unknown "
                        "element in a rule pattern site condition.")
    return '%s%s' % (site, state_code)

def sympy_to_muparser(expr):
    code = sympy.fcode(expr)
    code = code.replace('\n @', '')
    code = code.replace('**', '^')
    # kasim syntax cannot handle Fortran scientific notation (must use 'e'
    # instead of 'd')
    code = sub('(?<=[0-9])d','e',code)
    return code

class KappaException(Exception):
    pass
<commit_msg>Updated sympy expression func for kappa generator<commit_after>import pysb
import sympy
from re import sub
import warnings

class KappaGenerator(object):

    # Dialect can be either 'complx' or 'kasim' (default)
    def __init__(self, model, dialect='kasim'):
        self.model = model
        self.__content = None
        self.dialect = dialect

    def get_content(self):
        if self.__content == None:
            self.generate_content()
        return self.__content

    def generate_content(self):
        self.__content = ''
        #self.generate_compartments()

        # Agent declarations appear to be required in kasim
        # but prohibited in complx
        if (self.dialect == 'kasim'):
            self.generate_molecule_types()
            # Parameters, variables, and expressions are allowed in kasim
            self.generate_parameters()

        self.generate_reaction_rules()
        self.generate_observables()
        self.generate_species()

    def generate_parameters(self):
        names = [x.name for x in self.model.parameters] + \
                [x.name for x in self.model.expressions]
        for p in self.model.parameters:
            self.__content += ("%%var: \'%s\' %e\n") % (p.name, p.value)
        for e in self.model.expressions:
            sym_names = [x.name for x in e.expr.atoms(sympy.Symbol)]
            str_expr = str(sympy_to_muparser(e.expr))
            for n in sym_names:
                str_expr = str_expr.replace(n,"\'%s\'"%n)
            self.__content += ("%%var: \'%s\' %s\n") % \
                              (e.name, str_expr.replace('**','^'))
        self.__content += "\n"

    #def generate_compartments(self):
    #    self.__content += "begin compartments\n"
    #    for c in self.model.compartments:
    #        if c.parent is None:
    #            parent_name = ''
    #        else:
    #            parent_name = c.parent.name
    #        self.__content += ("  %s  %d  %f  %s\n") % \
    #                          (c.name, c.dimension, c.size, parent_name)
    #    self.__content += "end compartments\n\n"

    def generate_molecule_types(self):
        for m in self.model.monomers:
            site_code = ','.join([format_monomer_site(m, s) for s in m.sites])
            self.__content += "%%agent: %s(%s)\n" % (m.name, site_code)
        self.__content += "\n"

    def generate_reaction_rules(self):
        # +1 for the colon
        #max_length = max(len(r.name) for r in self.model.rules) + 1
        for r in self.model.rules:
            label = '\'' + r.name + '\''
            reactants_code = format_reactionpattern(r.reactant_pattern)
            products_code  = format_reactionpattern(r.product_pattern)
            arrow = '->'

            if isinstance(r.rate_forward,pysb.core.Expression) and \
               self.dialect == 'complx':
                raise KappaException("Expressions are not supported by complx.")
            # Get the rate code depending on the dialect
            if self.dialect == 'kasim':
                f_rate_code = "\'" + r.rate_forward.name + "\'"
            else:
                f_rate_code = float(r.rate_forward.value)

            self.__content += ("%s %s %s %s @ %s") % \
                (label, reactants_code, arrow, products_code, f_rate_code)
            self.__content += "\n"

            # Add the reverse reaction
            if r.is_reversible:
                if isinstance(r.rate_reverse,pysb.core.Expression) and \
                   self.dialect == 'complx':
                    raise KappaException("Expressions are not supported by "
                                         "complx.")
                # Get the rate code depending on the dialect
                if self.dialect == 'kasim':
                    r_rate_code = "\'" + r.rate_reverse.name + "\'"
                else:
                    r_rate_code = float(r.rate_reverse.value)

                label = '\'' + r.name + '_rev' + '\''
                self.__content += ("%s %s %s %s @ %s") % \
                      (label, products_code, arrow, reactants_code, r_rate_code)
                self.__content += "\n"

        self.__content += "\n"

    def generate_observables(self):
        if not self.model.observables:
            return
        for obs in self.model.observables:
            name = '\'' + obs.name + '\''
            observable_code = format_reactionpattern(obs.reaction_pattern)
            # In the near future (KaSim 4.0), the observable syntax will
            # require pipe characters around the expression. However, for the
            # time being we'll stick with the old syntax for backwards
            # compatibility
            #if self.dialect == 'kasim':
            #    self.__content += ("%%obs: %s |%s|\n") % \
            #                      (name, observable_code)
            #else:
            #    self.__content += ("%%obs: %s %s\n") % (name, observable_code)
            self.__content += ("%%obs: %s %s\n") % (name, observable_code)

        self.__content += "\n"

    def generate_species(self):
        if not self.model.initial_conditions:
            warnings.warn("Warning: No initial conditions.")

        species_codes = [format_complexpattern(cp)
                         for cp, param in self.model.initial_conditions]
        #max_length = max(len(code) for code in species_codes)
        for i, code in enumerate(species_codes):
            param = self.model.initial_conditions[i][1]
            #self.__content += ("%%init:  %-" + str(max_length) + \
            #                  "s   %s\n") % (code, param.name)
            if (self.dialect == 'kasim'):
                # Switched from %g (float) to %d (int) because kappa didn't
                # like scientific notation for large integers
                self.__content += ("%%init: \'%s\' %s\n") % (param.name, code)
                #self.__content += ("%%init: %10g %s\n") % (param.value, code)
            else:
                if isinstance(param,pysb.core.Expression):
                    raise KappaException("complx does not support Expressions.")
                # Switched from %g (float) to %d (int) because kappa didn't
                # like scientific notation for large integers
                self.__content += ("%%init: %10d * %s\n") % \
                                  (float(param.value), code)
                #self.__content += ("%%init: %10g * %s\n") % (param.value, code)

        self.__content += "\n"

def format_monomer_site(monomer, site):
    ret = site
    if site in monomer.site_states:
        for state in monomer.site_states[site]:
            ret += '~' + state
    return ret

def format_reactionpattern(rp):
    return ','.join([format_complexpattern(cp) for cp in rp.complex_patterns])

def format_complexpattern(cp):
    ret = ','.join([format_monomerpattern(mp) for mp in cp.monomer_patterns])
    if cp.compartment is not None:
        ret = '@%s:%s' % (cp.compartment.name, ret)
    return ret

def format_monomerpattern(mp):
    # sort sites in the same order given in the original Monomer
    site_conditions = sorted(mp.site_conditions.items(),
                             key=lambda x: mp.monomer.sites.index(x[0]))
    site_pattern_code = ','.join([format_site_condition(site, state)
                                  for (site, state) in site_conditions])
    ret = '%s(%s)' % (mp.monomer.name, site_pattern_code)
    if mp.compartment is not None:
        ret = '%s@%s' % (ret, mp.compartment.name)
    return ret

def format_site_condition(site, state):
    # If the state/bond is unspecified
    if state == None:
        state_code = ''
    # If there is a bond number
    elif type(state) == int:
        state_code = '!' + str(state)
    # If there is a lists of bonds to the site (not supported by Kappa)
    elif type(state) == list:
        raise KappaException("Kappa generator does not support multiple bonds "
                              "to a single site.")
    # Site with state
    elif type(state) == str:
        state_code = '~' + state
    # Site with state and a bond
    elif type(state) == tuple:
        # If the bond is ANY
        if state[1] == pysb.ANY:
            state_code = '~%s!_' % state[0]
        # If the bond is WILD
        elif state[1] == pysb.WILD:
            state_code = '~%s?' % state[0]
        # If the bond is a number
        elif type(state[1]) == int:
            state_code = '~%s!%s' % state
        # If it's something else, raise an Exception
        else:
            raise KappaException("Kappa generator has encountered an unknown "
                                 "element in a site state/bond tuple: (%s, %s)"
                                 % state)
    # Site bound to ANY
    elif state == pysb.ANY:
        state_code = '!_'
    # Site bound to WILD
    elif state == pysb.WILD:
        state_code = '?'
    # Something else
    else:
        raise KappaException("Kappa generator has encountered an unknown "
                        "element in a rule pattern site condition.")
    return '%s%s' % (site, state_code)

def sympy_to_kasim_expression(expr):
    """Render the Expression as a Kappa compatible string."""
    # sympy.printing.sstr is the preferred way to render an Expression as a
    # string (rather than, e.g., str(Expression.expr) or repr(Expression.expr).
    # Note: "For large expressions where speed is a concern, use the setting
    # order='none'"
    code = sympy.printing.sstr(expression.expr, order='none')
    code = code.replace('\n @', '')
    code = code.replace('**', '^')
    # kasim syntax cannot handle Fortran scientific notation (must use 'e'
    # instead of 'd')
    code = sub('(?<=[0-9])d', 'e', code)
    return code

class KappaException(Exception):
    pass
<|endoftext|>


<commit_before>"""
These classes are responsible for managing state during a client/service
request, particularly to minimize the burden on plugins to manage context
"""
from . import common
from . import wsgi
import requests

def service(service, operation, request_body):
    process = ServiceProcessor(service, operation, request_body)
    return process()

def client(client, operation, request_body):
    process = ClientProcessor(client, operation, request_body)
    return process()

class Processor(object):
    def __init__(self, obj, operation):
        """
        A python class with __init__ and 1 or 2 functions is usually an
        anti-pattern, but in this case we're using it to simplify the
        chaining contract for plugin authors.  This allows them to
        use context.process_request() without managing the state to pass
        to the next plugin, or in any way knowing if there is another plugin
        """
        self.obj = obj
        # Don't rely on context's operation to be immutable
        self.operation = operation

        self.context = common.Context(self)
        self.context.operation = operation

        self.request = common.Container()
        self.request_body = None
        self.response = common.Container()
        self.response_body = None

        # request -> operation -> function
        self.state = "request"
        self.transitions = {
            "request": "operation",
            "operation": "function",
            "function": None
        }
        self.index = -1

    @property
    def result(self):
        raise NotImplementedError("Subclasses must define result.")

    def __call__(self):
        """ Entry point for external callers """
        if self.state is None:
            raise ValueError("Already processed request")
        self.continue_execution()
        return self.result

    def continue_execution(self):
        # If this is the first time we've been in this state,
        # give subclasses a chance to (de)serialize, load/dump containers, etc
        call_after_state = False
        if self.index == -1:
            self.before_state(self.state)
            # Since we're about to invoke the first plugin for this state,
            # the recursive call will go through all remaining plugins for this
            # state.  After that recursive call, this scope will be the first
            # point after all plugins at this state ran, and we should call
            # self.after_state(self.state).  However, self.state will have
            # mutated (since this is recursive).  Therefore we store the state
            # to invoke after_state with in the variable.
            call_after_state = self.state

        # We don't recurse for the next plugin during 'function' since there
        # shouldn't be any function-scope plugins.
        if self.state in ["request", "operation"]:
            self._continue()
        # We've made it!  Make the remote call or invoke the service handler
        elif self.state == "function":
            self._execute()

        # If we called before_state in this scope, then we also need to call
        # after_state in this scope (the rest of this scope was executed in)
        # one of the recursive calls above
        if call_after_state:
            self.after_state(call_after_state)

    def _continue(self):
        ''' Call the next plugin '''
        self.index += 1
        # When state is function, there won't be any plugins
        plugins = self.obj.plugins.get(self.state, [])
        n = len(plugins)

        # We're still working through the plugins for this scope
        if self.index < n:
            if self.state == "request":
                plugins[self.index]()
            elif self.state == "operation":
                plugins[self.index](self.request, self.response, self.context)
            else:
                raise ValueError("Unexpected state '{}'".format(self.state))

        # We just processed the last plugin for this state, either roll over
        # to the next state or invoke the function underneath it all
        elif self.index == n:

            # Determine the next state and continue execution
            if self.transitions[self.state]:
                self.state = self.transitions[self.state]
                self.index = -1
                self.continue_execution()

            # Otherwise this is a continue_execution call during the 'function'
            # state and we can safely ignore it
            else:
                pass

    def before_state(self, state):
        ''' The state whose execution is about to begin '''
        pass

    def after_state(self, state):
        ''' The state whose execution just finished '''
        pass

class ClientProcessor(Processor):
    def __init__(self, client, operation, request):
        super().__init__(client, operation)
        self.context.client = client
        self.request.update(request)

    @property
    def result(self):
        return self.response

    def _execute(self):
        '''
        1. Pack the request
        2. Construct the endpoint
        3. Unpack the response
        4. Raise native errors if necessary
        '''

        self.request_body = common.serialize(self.request)

        pattern = self.obj.api["endpoint"]["client_pattern"]
        uri = pattern.format(operation=self.operation)
        data = self.request_body
        timeout = self.obj.api["timeout"]
        response = requests.post(uri, data=data, timeout=timeout)

        self.handle_http_error(response)
        self.response_body = response.text
        common.deserialize(self.response_body, self.response)
        self.handle_service_exception()

    def handle_http_error(self, response):
        if wsgi.is_request_exception(response):
            message = "{} {}".format(response.status_code, response.reason)
            self.raise_exception({
                "cls": "RequestException",
                "args": (message,)
            })

    def handle_service_exception(self):
        exception = self.response.get("__exception__", None)
        if exception:
            # Don't leak incomplete operation state if we're not debugging
            if not self.obj.api["debug"]:
                self.response.clear()
            self.raise_exception(exception)

    def raise_exception(self, exception):
        name = exception["cls"]
        args = exception["args"]
        raise getattr(self.obj.exceptions, name)(*args)

class ServiceProcessor(Processor):
    def __init__(self, service, operation, request_body):
        super().__init__(service, operation)
        self.context.service = service
        self.request_body = request_body

    def __call__(self):
        '''
        Wrap exceptions so they can be serialized back to the client.

        Exceptions that are thrown anywhere in the plugin chain may be allowed
        to serialize back to the client, so we have to try/except the entire
        call chain.
        '''
        try:
            super().__call__()
        except Exception as exception:
            self.raise_exception(exception)
        finally:
            return self.result

    def before_state(self, state):
        # Unpack request_body so it's available to operation scoped plugins
        if state == "operation":
            common.deserialize(self.request_body, self.request)

    def after_state(self, state):
        # Pack response into response body so we can ship it back on the wire
        # It's important to do this before the request-scope plugins clean up,
        # since their state may be required to serialize the response body
        if state == "operation":
            self.response_body = common.serialize(self.response)

    @property
    def result(self):
        return self.response_body

    def _execute(self):
        '''
        Invoke the service's function for the current operation
        '''
        func = self.obj.functions[self.operation]
        func(self.request, self.response, self.context)

    def raise_exception(self, exception):
        '''
        Because this can occurr after after_state('operation') has already
        fired, we have to make sure we re-serialize the response body.
        '''
        name = exception.__class__.__name__
        args = exception.args

        # Don't let non-whitelisted exceptions escape if we're not debugging
        whitelisted = name in self.obj.api["exceptions"]
        debugging = self.obj.api["debug"]
        if not whitelisted and not debugging:
            raise wsgi.INTERNAL_ERROR

        # Don't leak incomplete operation state
        self.response.clear()
        self.response["__exception__"] = {
            "cls": name,
            "args": args
        }
        self.response_body = common.serialize(self.response)
<commit_msg>Reorganize Processor class layout<commit_after>"""
These classes are responsible for managing state during a client/service
request, particularly to minimize the burden on plugins to manage context
"""
from . import common
from . import wsgi
import requests

def service(service, operation, request_body):
    process = ServiceProcessor(service, operation, request_body)
    return process()

def client(client, operation, request_body):
    process = ClientProcessor(client, operation, request_body)
    return process()

class Processor(object):
    def __init__(self, obj, operation):
        """
        A python class with __init__ and 1 or 2 functions is usually an
        anti-pattern, but in this case we're using it to simplify the
        chaining contract for plugin authors.  This allows them to
        use context.process_request() without managing the state to pass
        to the next plugin, or in any way knowing if there is another plugin
        """
        self.obj = obj
        # Don't rely on context's operation to be immutable
        self.operation = operation

        self.context = common.Context(self)
        self.context.operation = operation

        self.request = common.Container()
        self.request_body = None
        self.response = common.Container()
        self.response_body = None

        # request -> operation -> function
        self.state = "request"
        self.transitions = {
            "request": "operation",
            "operation": "function",
            "function": None
        }
        self.index = -1

    def __call__(self):
        """ Entry point for external callers """
        if self.state is None:
            raise ValueError("Already processed request")
        self.continue_execution()
        return self.result

    def continue_execution(self):
        # If this is the first time we've been in this state,
        # give subclasses a chance to (de)serialize, load/dump containers, etc
        call_after_state = False
        if self.index == -1:
            self.before_state(self.state)
            # Since we're about to invoke the first plugin for this state,
            # the recursive call will go through all remaining plugins for this
            # state.  After that recursive call, this scope will be the first
            # point after all plugins at this state ran, and we should call
            # self.after_state(self.state).  However, self.state will have
            # mutated (since this is recursive).  Therefore we store the state
            # to invoke after_state with in the variable.
            call_after_state = self.state

        # We don't recurse for the next plugin during 'function' since there
        # shouldn't be any function-scope plugins.
        if self.state in ["request", "operation"]:
            self._continue()
        # We've made it!  Make the remote call or invoke the service handler
        elif self.state == "function":
            self._execute()

        # If we called before_state in this scope, then we also need to call
        # after_state in this scope (the rest of this scope was executed in)
        # one of the recursive calls above
        if call_after_state:
            self.after_state(call_after_state)

    def _continue(self):
        ''' Call the next plugin '''
        self.index += 1
        # When state is function, there won't be any plugins
        plugins = self.obj.plugins.get(self.state, [])
        n = len(plugins)

        # We're still working through the plugins for this scope
        if self.index < n:
            if self.state == "request":
                plugins[self.index]()
            elif self.state == "operation":
                plugins[self.index](self.request, self.response, self.context)
            else:
                raise ValueError("Unexpected state '{}'".format(self.state))

        # We just processed the last plugin for this state, either roll over
        # to the next state or invoke the function underneath it all
        elif self.index == n:

            # Determine the next state and continue execution
            if self.transitions[self.state]:
                self.state = self.transitions[self.state]
                self.index = -1
                self.continue_execution()

            # Otherwise this is a continue_execution call during the 'function'
            # state and we can safely ignore it
            else:
                pass

    def _execute(self):
        raise NotImplementedError("Subclasses must implement _execute.")

    def before_state(self, state):
        ''' The state whose execution is about to begin '''
        pass

    def after_state(self, state):
        ''' The state whose execution just finished '''
        pass

    @property
    def result(self):
        raise NotImplementedError("Subclasses must define result.")

class ClientProcessor(Processor):
    def __init__(self, client, operation, request):
        super().__init__(client, operation)
        self.context.client = client
        self.request.update(request)

    def _execute(self):
        '''
        1. Pack the request
        2. Construct the endpoint
        3. Unpack the response
        4. Raise native errors if necessary
        '''

        self.request_body = common.serialize(self.request)

        pattern = self.obj.api["endpoint"]["client_pattern"]
        uri = pattern.format(operation=self.operation)
        data = self.request_body
        timeout = self.obj.api["timeout"]
        response = requests.post(uri, data=data, timeout=timeout)

        self.handle_http_error(response)
        self.response_body = response.text
        common.deserialize(self.response_body, self.response)
        self.handle_service_exception()

    @property
    def result(self):
        return self.response

    def handle_http_error(self, response):
        if wsgi.is_request_exception(response):
            message = "{} {}".format(response.status_code, response.reason)
            self.raise_exception({
                "cls": "RequestException",
                "args": (message,)
            })

    def handle_service_exception(self):
        exception = self.response.get("__exception__", None)
        if exception:
            # Don't leak incomplete operation state if we're not debugging
            if not self.obj.api["debug"]:
                self.response.clear()
            self.raise_exception(exception)

    def raise_exception(self, exception):
        name = exception["cls"]
        args = exception["args"]
        raise getattr(self.obj.exceptions, name)(*args)

class ServiceProcessor(Processor):
    def __init__(self, service, operation, request_body):
        super().__init__(service, operation)
        self.context.service = service
        self.request_body = request_body

    def __call__(self):
        '''
        Wrap exceptions so they can be serialized back to the client.

        Exceptions that are thrown anywhere in the plugin chain may be allowed
        to serialize back to the client, so we have to try/except the entire
        call chain.
        '''
        try:
            super().__call__()
        except Exception as exception:
            self.raise_exception(exception)
        finally:
            return self.result

    def _execute(self):
        '''
        Invoke the service's function for the current operation
        '''
        func = self.obj.functions[self.operation]
        func(self.request, self.response, self.context)

    def before_state(self, state):
        # Unpack request_body so it's available to operation scoped plugins
        if state == "operation":
            common.deserialize(self.request_body, self.request)

    def after_state(self, state):
        # Pack response into response body so we can ship it back on the wire
        # It's important to do this before the request-scope plugins clean up,
        # since their state may be required to serialize the response body
        if state == "operation":
            self.response_body = common.serialize(self.response)

    @property
    def result(self):
        return self.response_body

    def raise_exception(self, exception):
        '''
        Because this can occurr after after_state('operation') has already
        fired, we have to make sure we re-serialize the response body.
        '''
        name = exception.__class__.__name__
        args = exception.args

        # Don't let non-whitelisted exceptions escape if we're not debugging
        whitelisted = name in self.obj.api["exceptions"]
        debugging = self.obj.api["debug"]
        if not whitelisted and not debugging:
            raise wsgi.INTERNAL_ERROR

        # Don't leak incomplete operation state
        self.response.clear()
        self.response["__exception__"] = {
            "cls": name,
            "args": args
        }
        self.response_body = common.serialize(self.response)
<|endoftext|>


<commit_before>"""SNMP v3 Message Processing and Dispatching (RFC3412)"""
import sys
from pyasn1.compat.octets import null
from pysnmp.smi import builder, instrum
from pysnmp.proto import errind, error, cache
from pysnmp.proto.api import verdec # XXX
from pysnmp.error import PySnmpError
from pysnmp import nextid, debug

class MsgAndPduDispatcher:
    """SNMP engine PDU & message dispatcher. Exchanges SNMP PDU's with
       applications and serialized messages with transport level.
    """
    def __init__(self, mibInstrumController=None):
        if mibInstrumController is None:
            self.mibInstrumController = instrum.MibInstrumController(
                builder.MibBuilder()
                )
        else:
            self.mibInstrumController = mibInstrumController

        self.mibInstrumController.mibBuilder.loadModules(
            'SNMPv2-MIB', 'SNMP-MPD-MIB', 'SNMP-COMMUNITY-MIB',
            'SNMP-TARGET-MIB', 'SNMP-USER-BASED-SM-MIB'
            )

        # Requests cache
        self.__cache = cache.Cache()

        # Registered context engine IDs
        self.__appsRegistration = {}

        # Source of sendPduHandle and cache of requesting apps
        self.__sendPduHandle = nextid.Integer(0xffffff)

        # To pass transport info to app
        self.__transportInfo = {}

    def getTransportInfo(self, stateReference):
        if stateReference in self.__transportInfo:
            return self.__transportInfo[stateReference]
        else:
            raise error.ProtocolError(
                'No data for stateReference %s' % stateReference
                )

    # Application registration with dispatcher

    # 4.3.1
    def registerContextEngineId(self, contextEngineId, pduTypes, processPdu):
        """Register application with dispatcher"""
        # 4.3.2 -> noop

        # 4.3.3
        for pduType in pduTypes:
            k = (contextEngineId, pduType)
            if k in self.__appsRegistration:
                raise error.ProtocolError(
                    'Duplicate registration %s/%s' % (contextEngineId, pduType)
                    )

            # 4.3.4
            self.__appsRegistration[k] = processPdu

        debug.logger & debug.flagDsp and debug.logger('registerContextEngineId: contextEngineId %r pduTypes %s' % (contextEngineId, pduTypes))
    # 4.4.1
    def unregisterContextEngineId(self, contextEngineId, pduTypes):
        """Unregister application with dispatcher"""
        # 4.3.4
        if contextEngineId is None:
            # Default to local snmpEngineId
            contextEngineId, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-FRAMEWORK-MIB', 'snmpEngineID')

        for pduType in pduTypes:
            k = (contextEngineId, pduType)
            if k in self.__appsRegistration:
                del self.__appsRegistration[k]

        debug.logger & debug.flagDsp and debug.logger('unregisterContextEngineId: contextEngineId %r pduTypes %s' % (contextEngineId, pduTypes))

    def getRegisteredApp(self, contextEngineId, pduType):
        k = (contextEngineId, pduType)
        if k in self.__appsRegistration:
            return self.__appsRegistration[k]
        k = ( null, pduType )
        if k in self.__appsRegistration:
            return self.__appsRegistration[k] # wildcard

    # Dispatcher <-> application API

    # 4.1.1

    def sendPdu(
        self,
        snmpEngine,
        transportDomain,
        transportAddress,
        messageProcessingModel,
        securityModel,
        securityName,
        securityLevel,
        contextEngineId,
        contextName,
        pduVersion,
        PDU,
        expectResponse,
        timeout=0,    # timeout expressed in dispatcher ticks
        cbFun=None,
        cbCtx=None
        ):
        """PDU dispatcher -- prepare and serialize a request or notification"""
        # 4.1.1.2
        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
        else:
            raise error.StatusInformation(
                errorIndication=errind.unsupportedMsgProcessingModel
                )

        debug.logger & debug.flagDsp and debug.logger('sendPdu: securityName %s, PDU\n%s' % (securityName, PDU.prettyPrint()))

        # 4.1.1.3
        sendPduHandle = self.__sendPduHandle()
        if expectResponse:
            self.__cache.add(
                sendPduHandle,
                messageProcessingModel=messageProcessingModel,
                sendPduHandle=sendPduHandle,
                timeout=timeout+snmpEngine.transportDispatcher.getTimerTicks(),
                cbFun=cbFun,
                cbCtx=cbCtx
                )
            debug.logger & debug.flagDsp and debug.logger('sendPdu: current time in ticks %d' % (snmpEngine.transportDispatcher.getTimerTicks(),))

        debug.logger & debug.flagDsp and debug.logger('sendPdu: new sendPduHandle %s, timeout %s, cbFun %s' % (sendPduHandle, timeout, cbFun))

        # 4.1.1.4 & 4.1.1.5
        try:
            ( destTransportDomain,
              destTransportAddress,
              outgoingMessage ) = mpHandler.prepareOutgoingMessage(
                snmpEngine,
                transportDomain,
                transportAddress,
                messageProcessingModel,
                securityModel,
                securityName,
                securityLevel,
                contextEngineId,
                contextName,
                pduVersion,
                PDU,
                expectResponse,
                sendPduHandle
                )
            debug.logger & debug.flagDsp and debug.logger('sendPdu: MP succeeded')
        except error.StatusInformation:
# XXX is it still needed here?
#            self.releaseStateInformation(snmpEngine, sendPduHandle, messageProcessingModel)
            raise

        # 4.1.1.6
        if snmpEngine.transportDispatcher is None:
            raise error.PySnmpError('Transport dispatcher not set')
        snmpEngine.transportDispatcher.sendMessage(
            outgoingMessage, destTransportDomain, destTransportAddress
            )

        # Update cache with orignal req params (used for retrying)
        if expectResponse:
            self.__cache.update(
                sendPduHandle,
                transportDomain=transportDomain,
                transportAddress=transportAddress,
                securityModel=securityModel,
                securityName=securityName,
                securityLevel=securityLevel,
                contextEngineId=contextEngineId,
                contextName=contextName,
                pduVersion=pduVersion,
                PDU=PDU
                )

        return sendPduHandle

    # 4.1.2.1
    def returnResponsePdu(
        self,
        snmpEngine,
        messageProcessingModel,
        securityModel,
        securityName,
        securityLevel,
        contextEngineId,
        contextName,
        pduVersion,
        PDU,
        maxSizeResponseScopedPDU,
        stateReference,
        statusInformation
        ):
        """PDU dispatcher -- prepare and serialize a response"""
        # Extract input values and initialize defaults
        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
        else:
            raise error.StatusInformation(
                errorIndication=errind.unsupportedMsgProcessingModel
                )

        debug.logger & debug.flagDsp and debug.logger('returnResponsePdu: PDU %s' % (PDU and PDU.prettyPrint() or "<empty>",))

        # 4.1.2.2
        try:
            ( destTransportDomain,
              destTransportAddress,
              outgoingMessage ) = mpHandler.prepareResponseMessage(
                snmpEngine,
                messageProcessingModel,
                securityModel,
                securityName,
                securityLevel,
                contextEngineId,
                contextName,
                pduVersion,
                PDU,
                maxSizeResponseScopedPDU,
                stateReference,
                statusInformation
                )
            debug.logger & debug.flagDsp and debug.logger('returnResponsePdu: MP suceeded')
        except error.StatusInformation:
            # 4.1.2.3
            raise

        # Handle oversized messages XXX transport constrains?
        snmpEngineMaxMessageSize, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-FRAMEWORK-MIB', 'snmpEngineMaxMessageSize')
        if snmpEngineMaxMessageSize.syntax and \
               len(outgoingMessage) > snmpEngineMaxMessageSize.syntax:
            snmpSilentDrops, = self.mibInstrumController.mibBuilder.importSymbols('__SNMPv2-MIB', 'snmpSilentDrops')
            snmpSilentDrops.syntax = snmpSilentDrops.syntax + 1
            raise error.MessageTooBigError()

        # 4.1.2.4
        snmpEngine.transportDispatcher.sendMessage(
            outgoingMessage,
            destTransportDomain,
            destTransportAddress
            )

    # 4.2.1
    def receiveMessage(
        self,
        snmpEngine,
        transportDomain,
        transportAddress,
        wholeMsg
        ):
        """Message dispatcher -- de-serialize message into PDU"""
        # 4.2.1.1
        snmpInPkts, = self.mibInstrumController.mibBuilder.importSymbols(
            '__SNMPv2-MIB', 'snmpInPkts'
            )
        snmpInPkts.syntax = snmpInPkts.syntax + 1

        # 4.2.1.2
        try:
            restOfWholeMsg = null # XXX fix decoder non-recursive return
            msgVersion = verdec.decodeMessageVersion(wholeMsg)
        except error.ProtocolError:
            snmpInASNParseErrs, = self.mibInstrumController.mibBuilder.importSymbols('__SNMPv2-MIB', 'snmpInASNParseErrs')
            snmpInASNParseErrs.syntax = snmpInASNParseErrs.syntax + 1
            return null  # n.b the whole buffer gets dropped

        debug.logger & debug.flagDsp and debug.logger('receiveMessage: msgVersion %s, msg decoded' % msgVersion)

        messageProcessingModel = msgVersion

        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
        else:
            snmpInBadVersions, = self.mibInstrumController.mibBuilder.importSymbols('__SNMPv2-MIB', 'snmpInBadVersions')
            snmpInBadVersions.syntax = snmpInBadVersions.syntax + 1
            return restOfWholeMsg

        # 4.2.1.3 -- no-op

        # 4.2.1.4
        try:
            ( messageProcessingModel,
              securityModel,
              securityName,
              securityLevel,
              contextEngineId,
              contextName,
              pduVersion,
              PDU,
              pduType,
              sendPduHandle,
              maxSizeResponseScopedPDU,
              statusInformation,
              stateReference ) = mpHandler.prepareDataElements(
                snmpEngine,
                transportDomain,
                transportAddress,
                wholeMsg
                )
            debug.logger & debug.flagDsp and debug.logger('receiveMessage: MP succeded')
        except error.StatusInformation:
            statusInformation = sys.exc_info()[1]
            if 'sendPduHandle' in statusInformation:
                # Dropped REPORT -- re-run pending reqs queue as some
                # of them may be waiting for this REPORT
                debug.logger & debug.flagDsp and debug.logger('receiveMessage: MP failed, statusInformation %s' % statusInformation)
                self.__expireRequest(
                    statusInformation['sendPduHandle'],
                    self.__cache.pop(statusInformation['sendPduHandle']),
                    snmpEngine,
                    statusInformation
                    )
            return restOfWholeMsg

        debug.logger & debug.flagDsp and debug.logger('receiveMessage: PDU %s' % PDU.prettyPrint())

        # 4.2.2
        if sendPduHandle is None:
            # 4.2.2.1 (request or notification)

            debug.logger & debug.flagDsp and debug.logger('receiveMessage: pduType %s' % pduType)
            # 4.2.2.1.1
            processPdu = self.getRegisteredApp(contextEngineId, pduType)

            # 4.2.2.1.2
            if processPdu is None:
                # 4.2.2.1.2.a
                snmpUnknownPDUHandlers, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-MPD-MIB', 'snmpUnknownPDUHandlers')
                snmpUnknownPDUHandlers.syntax = snmpUnknownPDUHandlers.syntax+1

                # 4.2.2.1.2.b
                statusInformation = {
                    'errorIndication': errind.unknownPDUHandler,
                    'oid': snmpUnknownPDUHandlers.name,
                    'val': snmpUnknownPDUHandlers.syntax
                    }

                debug.logger & debug.flagDsp and debug.logger('receiveMessage: unhandled PDU type')

                # XXX fails on unknown PDU

                try:
                    ( destTransportDomain,
                      destTransportAddress,
                      outgoingMessage ) = mpHandler.prepareResponseMessage(
                        snmpEngine,
                        messageProcessingModel,
                        securityModel,
                        securityName,
                        securityLevel,
                        contextEngineId,
                        contextName,
                        pduVersion,
                        PDU,
                        maxSizeResponseScopedPDU,
                        stateReference,
                        statusInformation
                        )
                except error.StatusInformation:
                    debug.logger & debug.flagDsp and debug.logger('receiveMessage: report failed, statusInformation %s' % sys.exc_info()[1])
                    return restOfWholeMsg

                # 4.2.2.1.2.c
                try:
                    snmpEngine.transportDispatcher.sendMessage(
                        outgoingMessage,
                        destTransportDomain,
                        destTransportAddress
                        )
                except PySnmpError: # XXX
                    pass

                debug.logger & debug.flagDsp and debug.logger('receiveMessage: reporting succeeded')

                # 4.2.2.1.2.d
                return restOfWholeMsg
            else:
                # Pass transport info to app
                if stateReference is not None:
                    self.__transportInfo[stateReference] = (
                        transportDomain, transportAddress
                        )
                # 4.2.2.1.3
                processPdu(
                    snmpEngine,
                    messageProcessingModel,
                    securityModel,
                    securityName,
                    securityLevel,
                    contextEngineId,
                    contextName,
                    pduVersion,
                    PDU,
                    maxSizeResponseScopedPDU,
                    stateReference
                    )
                if stateReference is not None:
                    del self.__transportInfo[stateReference]
                debug.logger & debug.flagDsp and debug.logger('receiveMessage: processPdu succeeded')
                return restOfWholeMsg
        else:
            # 4.2.2.2 (response)

            # 4.2.2.2.1
            cachedParams = self.__cache.pop(sendPduHandle)

            # 4.2.2.2.2
            if cachedParams is None:
                snmpUnknownPDUHandlers, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-MPD-MIB', 'snmpUnknownPDUHandlers')
                snmpUnknownPDUHandlers.syntax = snmpUnknownPDUHandlers.syntax+1
                return restOfWholeMsg

            debug.logger & debug.flagDsp and debug.logger('receiveMessage: cache read by sendPduHandle %s' % sendPduHandle)

            # 4.2.2.2.3
            # no-op ? XXX

            # 4.2.2.2.4
            processResponsePdu = cachedParams['cbFun']
            processResponsePdu(
                snmpEngine,
                messageProcessingModel,
                securityModel,
                securityName,
                securityLevel,
                contextEngineId,
                contextName,
                pduVersion,
                PDU,
                statusInformation,
                cachedParams['sendPduHandle'],
                cachedParams['cbCtx']
                )
            debug.logger & debug.flagDsp and debug.logger('receiveMessage: processResponsePdu succeeded')
            return restOfWholeMsg

    def releaseStateInformation(
        self, snmpEngine, sendPduHandle, messageProcessingModel
        ):
        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
            mpHandler.releaseStateInformation(sendPduHandle)

    # Cache expiration stuff

    def __expireRequest(self, cacheKey, cachedParams, snmpEngine,
                        statusInformation=None):
        timeNow = snmpEngine.transportDispatcher.getTimerTicks()
        timeoutAt = cachedParams['timeout']

        if statusInformation is None and timeNow < timeoutAt:
            return

        processResponsePdu = cachedParams['cbFun']

        debug.logger & debug.flagDsp and debug.logger('__expireRequest: req cachedParams %s' % cachedParams)

        # Fail timed-out requests
        if not statusInformation:
            statusInformation = error.StatusInformation(
                errorIndication=errind.requestTimedOut
                )
        self.releaseStateInformation(
            snmpEngine,
            cachedParams['sendPduHandle'],
            cachedParams['messageProcessingModel']
            )
        processResponsePdu(
            snmpEngine,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            statusInformation,
            cachedParams['sendPduHandle'],
            cachedParams['cbCtx']
            )
        return 1

    def receiveTimerTick(self, snmpEngine, timeNow):
        self.__cache.expire(self.__expireRequest, snmpEngine)
<commit_msg>comment added<commit_after>"""SNMP v3 Message Processing and Dispatching (RFC3412)"""
import sys
from pyasn1.compat.octets import null
from pysnmp.smi import builder, instrum
from pysnmp.proto import errind, error, cache
from pysnmp.proto.api import verdec # XXX
from pysnmp.error import PySnmpError
from pysnmp import nextid, debug

class MsgAndPduDispatcher:
    """SNMP engine PDU & message dispatcher. Exchanges SNMP PDU's with
       applications and serialized messages with transport level.
    """
    def __init__(self, mibInstrumController=None):
        if mibInstrumController is None:
            self.mibInstrumController = instrum.MibInstrumController(
                builder.MibBuilder()
                )
        else:
            self.mibInstrumController = mibInstrumController

        self.mibInstrumController.mibBuilder.loadModules(
            'SNMPv2-MIB', 'SNMP-MPD-MIB', 'SNMP-COMMUNITY-MIB',
            'SNMP-TARGET-MIB', 'SNMP-USER-BASED-SM-MIB'
            )

        # Requests cache
        self.__cache = cache.Cache()

        # Registered context engine IDs
        self.__appsRegistration = {}

        # Source of sendPduHandle and cache of requesting apps
        self.__sendPduHandle = nextid.Integer(0xffffff)

        # To pass transport info to app
        self.__transportInfo = {}

    def getTransportInfo(self, stateReference):
        if stateReference in self.__transportInfo:
            return self.__transportInfo[stateReference]
        else:
            raise error.ProtocolError(
                'No data for stateReference %s' % stateReference
                )

    # Application registration with dispatcher

    # 4.3.1
    def registerContextEngineId(self, contextEngineId, pduTypes, processPdu):
        """Register application with dispatcher"""
        # 4.3.2 -> noop

        # 4.3.3
        for pduType in pduTypes:
            k = (contextEngineId, pduType)
            if k in self.__appsRegistration:
                raise error.ProtocolError(
                    'Duplicate registration %s/%s' % (contextEngineId, pduType)
                    )

            # 4.3.4
            self.__appsRegistration[k] = processPdu

        debug.logger & debug.flagDsp and debug.logger('registerContextEngineId: contextEngineId %r pduTypes %s' % (contextEngineId, pduTypes))
    # 4.4.1
    def unregisterContextEngineId(self, contextEngineId, pduTypes):
        """Unregister application with dispatcher"""
        # 4.3.4
        if contextEngineId is None:
            # Default to local snmpEngineId
            contextEngineId, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-FRAMEWORK-MIB', 'snmpEngineID')

        for pduType in pduTypes:
            k = (contextEngineId, pduType)
            if k in self.__appsRegistration:
                del self.__appsRegistration[k]

        debug.logger & debug.flagDsp and debug.logger('unregisterContextEngineId: contextEngineId %r pduTypes %s' % (contextEngineId, pduTypes))

    def getRegisteredApp(self, contextEngineId, pduType):
        k = (contextEngineId, pduType)
        if k in self.__appsRegistration:
            return self.__appsRegistration[k]
        k = ( null, pduType )
        if k in self.__appsRegistration:
            return self.__appsRegistration[k] # wildcard

    # Dispatcher <-> application API

    # 4.1.1

    def sendPdu(
        self,
        snmpEngine,
        transportDomain,
        transportAddress,
        messageProcessingModel,
        securityModel,
        securityName,
        securityLevel,
        contextEngineId,
        contextName,
        pduVersion,
        PDU,
        expectResponse,
        timeout=0,    # timeout expressed in dispatcher ticks
        cbFun=None,
        cbCtx=None
        ):
        """PDU dispatcher -- prepare and serialize a request or notification"""
        # 4.1.1.2
        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
        else:
            raise error.StatusInformation(
                errorIndication=errind.unsupportedMsgProcessingModel
                )

        debug.logger & debug.flagDsp and debug.logger('sendPdu: securityName %s, PDU\n%s' % (securityName, PDU.prettyPrint()))

        # 4.1.1.3
        sendPduHandle = self.__sendPduHandle()
        if expectResponse:
            self.__cache.add(
                sendPduHandle,
                messageProcessingModel=messageProcessingModel,
                sendPduHandle=sendPduHandle,
                timeout=timeout+snmpEngine.transportDispatcher.getTimerTicks(),
                cbFun=cbFun,
                cbCtx=cbCtx
                )
            debug.logger & debug.flagDsp and debug.logger('sendPdu: current time in ticks %d' % (snmpEngine.transportDispatcher.getTimerTicks(),))

        debug.logger & debug.flagDsp and debug.logger('sendPdu: new sendPduHandle %s, timeout %s, cbFun %s' % (sendPduHandle, timeout, cbFun))

        # 4.1.1.4 & 4.1.1.5
        try:
            ( destTransportDomain,
              destTransportAddress,
              outgoingMessage ) = mpHandler.prepareOutgoingMessage(
                snmpEngine,
                transportDomain,
                transportAddress,
                messageProcessingModel,
                securityModel,
                securityName,
                securityLevel,
                contextEngineId,
                contextName,
                pduVersion,
                PDU,
                expectResponse,
                sendPduHandle
                )
            debug.logger & debug.flagDsp and debug.logger('sendPdu: MP succeeded')
        except error.StatusInformation:
# XXX is it still needed here?
#            self.releaseStateInformation(snmpEngine, sendPduHandle, messageProcessingModel)
            raise

        # 4.1.1.6
        if snmpEngine.transportDispatcher is None:
            raise error.PySnmpError('Transport dispatcher not set')
        snmpEngine.transportDispatcher.sendMessage(
            outgoingMessage, destTransportDomain, destTransportAddress
            )

        # Update cache with orignal req params (used for retrying)
        if expectResponse:
            self.__cache.update(
                sendPduHandle,
                transportDomain=transportDomain,
                transportAddress=transportAddress,
                securityModel=securityModel,
                securityName=securityName,
                securityLevel=securityLevel,
                contextEngineId=contextEngineId,
                contextName=contextName,
                pduVersion=pduVersion,
                PDU=PDU
                )

        return sendPduHandle

    # 4.1.2.1
    def returnResponsePdu(
        self,
        snmpEngine,
        messageProcessingModel,
        securityModel,
        securityName,
        securityLevel,
        contextEngineId,
        contextName,
        pduVersion,
        PDU,
        maxSizeResponseScopedPDU,
        stateReference,
        statusInformation
        ):
        """PDU dispatcher -- prepare and serialize a response"""
        # Extract input values and initialize defaults
        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
        else:
            raise error.StatusInformation(
                errorIndication=errind.unsupportedMsgProcessingModel
                )

        debug.logger & debug.flagDsp and debug.logger('returnResponsePdu: PDU %s' % (PDU and PDU.prettyPrint() or "<empty>",))

        # 4.1.2.2
        try:
            ( destTransportDomain,
              destTransportAddress,
              outgoingMessage ) = mpHandler.prepareResponseMessage(
                snmpEngine,
                messageProcessingModel,
                securityModel,
                securityName,
                securityLevel,
                contextEngineId,
                contextName,
                pduVersion,
                PDU,
                maxSizeResponseScopedPDU,
                stateReference,
                statusInformation
                )
            debug.logger & debug.flagDsp and debug.logger('returnResponsePdu: MP suceeded')
        except error.StatusInformation:
            # 4.1.2.3
            raise

        # Handle oversized messages XXX transport constrains?
        snmpEngineMaxMessageSize, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-FRAMEWORK-MIB', 'snmpEngineMaxMessageSize')
        if snmpEngineMaxMessageSize.syntax and \
               len(outgoingMessage) > snmpEngineMaxMessageSize.syntax:
            snmpSilentDrops, = self.mibInstrumController.mibBuilder.importSymbols('__SNMPv2-MIB', 'snmpSilentDrops')
            snmpSilentDrops.syntax = snmpSilentDrops.syntax + 1
            raise error.MessageTooBigError()

        # 4.1.2.4
        snmpEngine.transportDispatcher.sendMessage(
            outgoingMessage,
            destTransportDomain,
            destTransportAddress
            )

    # 4.2.1
    def receiveMessage(
        self,
        snmpEngine,
        transportDomain,
        transportAddress,
        wholeMsg
        ):
        """Message dispatcher -- de-serialize message into PDU"""
        # 4.2.1.1
        snmpInPkts, = self.mibInstrumController.mibBuilder.importSymbols(
            '__SNMPv2-MIB', 'snmpInPkts'
            )
        snmpInPkts.syntax = snmpInPkts.syntax + 1

        # 4.2.1.2
        try:
            restOfWholeMsg = null # XXX fix decoder non-recursive return
            msgVersion = verdec.decodeMessageVersion(wholeMsg)
        except error.ProtocolError:
            snmpInASNParseErrs, = self.mibInstrumController.mibBuilder.importSymbols('__SNMPv2-MIB', 'snmpInASNParseErrs')
            snmpInASNParseErrs.syntax = snmpInASNParseErrs.syntax + 1
            return null  # n.b the whole buffer gets dropped

        debug.logger & debug.flagDsp and debug.logger('receiveMessage: msgVersion %s, msg decoded' % msgVersion)

        messageProcessingModel = msgVersion

        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
        else:
            snmpInBadVersions, = self.mibInstrumController.mibBuilder.importSymbols('__SNMPv2-MIB', 'snmpInBadVersions')
            snmpInBadVersions.syntax = snmpInBadVersions.syntax + 1
            return restOfWholeMsg

        # 4.2.1.3 -- no-op

        # 4.2.1.4
        try:
            ( messageProcessingModel,
              securityModel,
              securityName,
              securityLevel,
              contextEngineId,
              contextName,
              pduVersion,
              PDU,
              pduType,
              sendPduHandle,
              maxSizeResponseScopedPDU,
              statusInformation,
              stateReference ) = mpHandler.prepareDataElements(
                snmpEngine,
                transportDomain,
                transportAddress,
                wholeMsg
                )
            debug.logger & debug.flagDsp and debug.logger('receiveMessage: MP succeded')
        except error.StatusInformation:
            statusInformation = sys.exc_info()[1]
            if 'sendPduHandle' in statusInformation:
                # Dropped REPORT -- re-run pending reqs queue as some
                # of them may be waiting for this REPORT
                debug.logger & debug.flagDsp and debug.logger('receiveMessage: MP failed, statusInformation %s, forcing a retry' % statusInformation)
                self.__expireRequest(
                    statusInformation['sendPduHandle'],
                    self.__cache.pop(statusInformation['sendPduHandle']),
                    snmpEngine,
                    statusInformation
                    )
            return restOfWholeMsg

        debug.logger & debug.flagDsp and debug.logger('receiveMessage: PDU %s' % PDU.prettyPrint())

        # 4.2.2
        if sendPduHandle is None:
            # 4.2.2.1 (request or notification)

            debug.logger & debug.flagDsp and debug.logger('receiveMessage: pduType %s' % pduType)
            # 4.2.2.1.1
            processPdu = self.getRegisteredApp(contextEngineId, pduType)

            # 4.2.2.1.2
            if processPdu is None:
                # 4.2.2.1.2.a
                snmpUnknownPDUHandlers, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-MPD-MIB', 'snmpUnknownPDUHandlers')
                snmpUnknownPDUHandlers.syntax = snmpUnknownPDUHandlers.syntax+1

                # 4.2.2.1.2.b
                statusInformation = {
                    'errorIndication': errind.unknownPDUHandler,
                    'oid': snmpUnknownPDUHandlers.name,
                    'val': snmpUnknownPDUHandlers.syntax
                    }

                debug.logger & debug.flagDsp and debug.logger('receiveMessage: unhandled PDU type')

                # XXX fails on unknown PDU

                try:
                    ( destTransportDomain,
                      destTransportAddress,
                      outgoingMessage ) = mpHandler.prepareResponseMessage(
                        snmpEngine,
                        messageProcessingModel,
                        securityModel,
                        securityName,
                        securityLevel,
                        contextEngineId,
                        contextName,
                        pduVersion,
                        PDU,
                        maxSizeResponseScopedPDU,
                        stateReference,
                        statusInformation
                        )
                except error.StatusInformation:
                    debug.logger & debug.flagDsp and debug.logger('receiveMessage: report failed, statusInformation %s' % sys.exc_info()[1])
                    return restOfWholeMsg

                # 4.2.2.1.2.c
                try:
                    snmpEngine.transportDispatcher.sendMessage(
                        outgoingMessage,
                        destTransportDomain,
                        destTransportAddress
                        )
                except PySnmpError: # XXX
                    pass

                debug.logger & debug.flagDsp and debug.logger('receiveMessage: reporting succeeded')

                # 4.2.2.1.2.d
                return restOfWholeMsg
            else:
                # Pass transport info to app
                if stateReference is not None:
                    self.__transportInfo[stateReference] = (
                        transportDomain, transportAddress
                        )
                # 4.2.2.1.3
                processPdu(
                    snmpEngine,
                    messageProcessingModel,
                    securityModel,
                    securityName,
                    securityLevel,
                    contextEngineId,
                    contextName,
                    pduVersion,
                    PDU,
                    maxSizeResponseScopedPDU,
                    stateReference
                    )
                if stateReference is not None:
                    del self.__transportInfo[stateReference]
                debug.logger & debug.flagDsp and debug.logger('receiveMessage: processPdu succeeded')
                return restOfWholeMsg
        else:
            # 4.2.2.2 (response)

            # 4.2.2.2.1
            cachedParams = self.__cache.pop(sendPduHandle)

            # 4.2.2.2.2
            if cachedParams is None:
                snmpUnknownPDUHandlers, = self.mibInstrumController.mibBuilder.importSymbols('__SNMP-MPD-MIB', 'snmpUnknownPDUHandlers')
                snmpUnknownPDUHandlers.syntax = snmpUnknownPDUHandlers.syntax+1
                return restOfWholeMsg

            debug.logger & debug.flagDsp and debug.logger('receiveMessage: cache read by sendPduHandle %s' % sendPduHandle)

            # 4.2.2.2.3
            # no-op ? XXX

            # 4.2.2.2.4
            processResponsePdu = cachedParams['cbFun']
            processResponsePdu(
                snmpEngine,
                messageProcessingModel,
                securityModel,
                securityName,
                securityLevel,
                contextEngineId,
                contextName,
                pduVersion,
                PDU,
                statusInformation,
                cachedParams['sendPduHandle'],
                cachedParams['cbCtx']
                )
            debug.logger & debug.flagDsp and debug.logger('receiveMessage: processResponsePdu succeeded')
            return restOfWholeMsg

    def releaseStateInformation(
        self, snmpEngine, sendPduHandle, messageProcessingModel
        ):
        k = int(messageProcessingModel)
        if k in snmpEngine.messageProcessingSubsystems:
            mpHandler = snmpEngine.messageProcessingSubsystems[k]
            mpHandler.releaseStateInformation(sendPduHandle)

    # Cache expiration stuff

    def __expireRequest(self, cacheKey, cachedParams, snmpEngine,
                        statusInformation=None):
        timeNow = snmpEngine.transportDispatcher.getTimerTicks()
        timeoutAt = cachedParams['timeout']

        if statusInformation is None and timeNow < timeoutAt:
            return

        processResponsePdu = cachedParams['cbFun']

        debug.logger & debug.flagDsp and debug.logger('__expireRequest: req cachedParams %s' % cachedParams)

        # Fail timed-out requests
        if not statusInformation:
            statusInformation = error.StatusInformation(
                errorIndication=errind.requestTimedOut
                )
        self.releaseStateInformation(
            snmpEngine,
            cachedParams['sendPduHandle'],
            cachedParams['messageProcessingModel']
            )
        processResponsePdu(
            snmpEngine,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            statusInformation,
            cachedParams['sendPduHandle'],
            cachedParams['cbCtx']
            )
        return 1

    def receiveTimerTick(self, snmpEngine, timeNow):
        self.__cache.expire(self.__expireRequest, snmpEngine)
<|endoftext|>


<commit_before>from errno import EEXIST
import os
import sys
from functools import partial

if sys.version_info < (3,):
    def makedirs(name, mode=0777, exist_ok=False):
        try:
            os.makedirs(name=name, mode=mode)
        except OSError as e:
            if not exist_ok or e.errno != EEXIST or not os.path.isdir(name):
                raise
else:
    makedirs = partial(os.makedirs)
<commit_msg>Fix python_two compat module mode arg<commit_after>from errno import EEXIST
import os
import sys
from functools import partial

if sys.version_info < (3,):
    def makedirs(name, exist_ok=False):
        try:
            os.makedirs(name=name)
        except OSError as e:
            if not exist_ok or e.errno != EEXIST or not os.path.isdir(name):
                raise
else:
    makedirs = partial(os.makedirs)
<|endoftext|>


<commit_before>import inspect
import os.path
from functools import partial

try:
    # Pytest 6.x
    from _pytest.skipping import xfailed_key as evalxfail_key
    from _pytest.skipping import evaluate_xfail_marks as mark_eval
except ImportError:
    # Pytest 5.x
    from _pytest.mark.evaluate import MarkEvaluator
    mark_eval = partial(MarkEvaluator, name="xfail")
    try:
        from _pytest.skipping import evalxfail_key
    except ImportError:
        # And pytest 3-4.x
        evalxfail_key = ""

from six import reraise as raise_

import pytest

try:
    from py.io import saferepr
except ImportError:
    saferepr = repr

_FAILED_ASSUMPTIONS = []

class Assumption(object):
    __slots__ = ["entry", "tb", "locals"]

    def __init__(self, entry, tb, locals=None):
        self.entry = entry
        # TODO: trim the TB at init?
        self.tb = tb
        self.locals = locals

    def longrepr(self):
        output = [self.entry, "Locals:"]
        output.extend(self.locals)

        return "\n".join(output)

    def repr(self):
        return self.entry

class FailedAssumption(AssertionError):
    pass

class AssumeContextManager(object):
    """Context manager whose objects can be used for *soft-assertions*

    This context manager can be accessed directly through `pytest.assume`.

    Checks the expression, if it's false, add it to the
    list of failed assumptions. Also, add the locals at each failed
    assumption, if showlocals is set.

    When used as a context manager::

        with pytest.assume:
            assert expr, msg

    When used directly, it also provides a return value::

        ret = pytest.assume(expr, msg)

    :param expr: Expression to 'assert' on.
    :param msg: Message to display if the assertion fails.
    :return: True or False, acording to `expr`
    """

    def __init__(self):
        self._enter_from_call = False

    def __enter__(self):
        __tracebackhide__ = True
        self._last_status = None
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        __tracebackhide__ = True
        pretty_locals = None
        entry = None
        tb = None
        stack_level = 2 if self._enter_from_call else 1
        (frame, filename, line, funcname, contextlist) = inspect.stack()[stack_level][0:5]
        # get filename, line, and context
        try:
            filename = os.path.relpath(filename)
        except ValueError:
            pass  # filename is on a different mount than the current dir (Windows)

        context = "" if contextlist is None else contextlist[0].lstrip()

        if exc_type is None:
            # format entry
            entry = u"{filename}:{line}: AssumptionSuccess\n>>\t{context}".format(**locals())
            pytest._hook_assume_pass(lineno=line, entry=entry)

            self._last_status = True
            return True

        elif issubclass(exc_type, AssertionError):
            if exc_val:
                context += "{}: {}\n\n".format(exc_type.__name__, exc_val)

            # format entry
            entry = u"{filename}:{line}: AssumptionFailure\n>>\t{context}".format(**locals())

            # Debatable whether we should display locals for
            # every failed assertion, or just the final one.
            # I'm defaulting to per-assumption, just because vars
            # can easily change between assumptions.
            pretty_locals = [
                "\t%-10s = %s" % (name, saferepr(val)) for name, val in frame.f_locals.items()
            ]

            pytest._hook_assume_fail(lineno=line, entry=entry)
            _FAILED_ASSUMPTIONS.append(Assumption(entry, exc_tb, pretty_locals))

            self._last_status = False
            return True

        else:
            # Another type of exception, let it rise uncaught
            return

    def __call__(self, expr, msg=""):
        __tracebackhide__ = True
        self._enter_from_call = True
        with self:
            if msg:
                assert expr, msg
            else:
                assert expr
        self._enter_from_call = False
        return self._last_status

assume = AssumeContextManager()

def pytest_addhooks(pluginmanager):
    """ This example assumes the hooks are grouped in the 'hooks' module. """

    from . import hooks

    pluginmanager.add_hookspecs(hooks)

def pytest_configure(config):
    """
    Add tracking lists to the pytest namespace, so we can
    always access it, as well as the 'assume' function itself.

    :return: Dictionary of name: values added to the pytest namespace.
    """
    pytest.assume = assume
    pytest._showlocals = config.getoption("showlocals")

    # As per pytest documentation: https://docs.pytest.org/en/latest/deprecations.html
    # The pytest.config global object is deprecated. Instead use request.config (via the request fixture)
    # or if you are a plugin author use the pytest_configure(config) hook.
    pytest._hook_assume_fail = config.pluginmanager.hook.pytest_assume_fail
    pytest._hook_assume_pass = config.pluginmanager.hook.pytest_assume_pass
    pytest._hook_assume_summary_report = config.pluginmanager.hook.pytest_assume_summary_report

@pytest.hookimpl(tryfirst=True)
def pytest_assume_fail(lineno, entry):
    pass

@pytest.hookimpl(tryfirst=True)
def pytest_assume_pass(lineno, entry):
    pass

@pytest.hookimpl(tryfirst=True)
def pytest_assume_summary_report(failed_assumptions):
    if getattr(pytest, "_showlocals"):
        content = "".join(x.longrepr() for x in failed_assumptions)
    else:
        content = "".join(x.repr() for x in failed_assumptions)

    return content

def restore_xfail(item):
    # Restore the xfail marker, as it's removed by the strict xfail checking, and will need to be
    # there for later xfail/xpass checking to work.
    if hasattr(item, "_store"):
        item._store[evalxfail_key] = mark_eval(item)
    else:
        item._evalxfail = mark_eval(item)

@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_call(item):
    """
    Using pyfunc_call to be as 'close' to the actual call of the test as possible.

    This is executed immediately after the test itself is called.

    Note: I'm not happy with exception handling in here.
    """
    __tracebackhide__ = True
    outcome = None
    try:
        outcome = yield
    finally:
        failed_assumptions = _FAILED_ASSUMPTIONS
        if failed_assumptions:
            failed_count = len(failed_assumptions)
            root_msg = "\n%s Failed Assumptions:\n" % failed_count

            content = pytest._hook_assume_summary_report(failed_assumptions=failed_assumptions)

            # Pluggy module returns list for multiple implementation of hooks
            # The user, while implementing custom hook pytest_assume_summary_report, will return "string"
            # Default hook is always present as list element 0
            if len(content) == 1:  # default length
                # Uses default hook
                content = content[0]
            else:
                # User created hook, if any
                content = content[1]

            last_tb = failed_assumptions[-1].tb

            del _FAILED_ASSUMPTIONS[:]
            if outcome and outcome.excinfo:
                # Xfailed test, but with strict=True. This is done via the pytest_pyfunc_call() hook, which
                # is before our hook.
                if "[XPASS(strict)]" in str(outcome.excinfo[1]):
                    restore_xfail(item)
                    raise_(FailedAssumption, FailedAssumption("%s\n%s" % (root_msg, content)), last_tb)
                root_msg = "\nOriginal Failure:\n\n>> %s\n" % repr(outcome.excinfo[1]) + root_msg
                raise_(
                    FailedAssumption,
                    FailedAssumption(root_msg + "\n" + content),
                    outcome.excinfo[2],
                )
            else:
                exc = FailedAssumption(root_msg + "\n" + content)
                # Note: raising here so that we guarantee a failure.
                raise_(FailedAssumption, exc, last_tb)
<commit_msg>Added pyfunc_call hook back, still compatible with unittests<commit_after>import inspect
import os.path
from functools import partial

try:
    # Pytest 6.x
    from _pytest.skipping import xfailed_key as evalxfail_key
    from _pytest.skipping import evaluate_xfail_marks as mark_eval
except ImportError:
    # Pytest 5.x
    from _pytest.mark.evaluate import MarkEvaluator
    mark_eval = partial(MarkEvaluator, name="xfail")
    try:
        from _pytest.skipping import evalxfail_key
    except ImportError:
        # And pytest 3-4.x
        evalxfail_key = ""

from six import reraise as raise_

import pytest

try:
    from py.io import saferepr
except ImportError:
    saferepr = repr

_FAILED_ASSUMPTIONS = []

class Assumption(object):
    __slots__ = ["entry", "tb", "locals"]

    def __init__(self, entry, tb, locals=None):
        self.entry = entry
        # TODO: trim the TB at init?
        self.tb = tb
        self.locals = locals

    def longrepr(self):
        output = [self.entry, "Locals:"]
        output.extend(self.locals)

        return "\n".join(output)

    def repr(self):
        return self.entry

class FailedAssumption(AssertionError):
    pass

class AssumeContextManager(object):
    """Context manager whose objects can be used for *soft-assertions*

    This context manager can be accessed directly through `pytest.assume`.

    Checks the expression, if it's false, add it to the
    list of failed assumptions. Also, add the locals at each failed
    assumption, if showlocals is set.

    When used as a context manager::

        with pytest.assume:
            assert expr, msg

    When used directly, it also provides a return value::

        ret = pytest.assume(expr, msg)

    :param expr: Expression to 'assert' on.
    :param msg: Message to display if the assertion fails.
    :return: True or False, acording to `expr`
    """

    def __init__(self):
        self._enter_from_call = False

    def __enter__(self):
        __tracebackhide__ = True
        self._last_status = None
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        __tracebackhide__ = True
        pretty_locals = None
        entry = None
        tb = None
        stack_level = 2 if self._enter_from_call else 1
        (frame, filename, line, funcname, contextlist) = inspect.stack()[stack_level][0:5]
        # get filename, line, and context
        try:
            filename = os.path.relpath(filename)
        except ValueError:
            pass  # filename is on a different mount than the current dir (Windows)

        context = "" if contextlist is None else contextlist[0].lstrip()

        if exc_type is None:
            # format entry
            entry = u"{filename}:{line}: AssumptionSuccess\n>>\t{context}".format(**locals())
            pytest._hook_assume_pass(lineno=line, entry=entry)

            self._last_status = True
            return True

        elif issubclass(exc_type, AssertionError):
            if exc_val:
                context += "{}: {}\n\n".format(exc_type.__name__, exc_val)

            # format entry
            entry = u"{filename}:{line}: AssumptionFailure\n>>\t{context}".format(**locals())

            # Debatable whether we should display locals for
            # every failed assertion, or just the final one.
            # I'm defaulting to per-assumption, just because vars
            # can easily change between assumptions.
            pretty_locals = [
                "\t%-10s = %s" % (name, saferepr(val)) for name, val in frame.f_locals.items()
            ]

            pytest._hook_assume_fail(lineno=line, entry=entry)
            _FAILED_ASSUMPTIONS.append(Assumption(entry, exc_tb, pretty_locals))

            self._last_status = False
            return True

        else:
            # Another type of exception, let it rise uncaught
            return

    def __call__(self, expr, msg=""):
        __tracebackhide__ = True
        self._enter_from_call = True
        with self:
            if msg:
                assert expr, msg
            else:
                assert expr
        self._enter_from_call = False
        return self._last_status

assume = AssumeContextManager()

def pytest_addhooks(pluginmanager):
    """ This example assumes the hooks are grouped in the 'hooks' module. """

    from . import hooks

    pluginmanager.add_hookspecs(hooks)

def pytest_configure(config):
    """
    Add tracking lists to the pytest namespace, so we can
    always access it, as well as the 'assume' function itself.

    :return: Dictionary of name: values added to the pytest namespace.
    """
    pytest.assume = assume
    pytest._showlocals = config.getoption("showlocals")

    # As per pytest documentation: https://docs.pytest.org/en/latest/deprecations.html
    # The pytest.config global object is deprecated. Instead use request.config (via the request fixture)
    # or if you are a plugin author use the pytest_configure(config) hook.
    pytest._hook_assume_fail = config.pluginmanager.hook.pytest_assume_fail
    pytest._hook_assume_pass = config.pluginmanager.hook.pytest_assume_pass
    pytest._hook_assume_summary_report = config.pluginmanager.hook.pytest_assume_summary_report

@pytest.hookimpl(tryfirst=True)
def pytest_assume_fail(lineno, entry):
    pass

@pytest.hookimpl(tryfirst=True)
def pytest_assume_pass(lineno, entry):
    pass

@pytest.hookimpl(tryfirst=True)
def pytest_assume_summary_report(failed_assumptions):
    if getattr(pytest, "_showlocals"):
        content = "".join(x.longrepr() for x in failed_assumptions)
    else:
        content = "".join(x.repr() for x in failed_assumptions)

    return content

def restore_xfail(item):
    # Restore the xfail marker, as it's removed by the strict xfail checking, and will need to be
    # there for later xfail/xpass checking to work.
    if hasattr(item, "_store"):
        item._store[evalxfail_key] = mark_eval(item)
    else:
        item._evalxfail = mark_eval(item)

@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_call(item):
    """
    Using runtest_call to be compatible with unittest tests, as pyfunccall is not
    called at all.
    """
    __tracebackhide__ = True
    outcome = None
    try:
        outcome = yield
    finally:
        # Check if item was assumed already in pyfunccall wrapper
        if getattr(item, "__assumed", None) is None:
            assume_outcome(outcome, item)

@pytest.hookimpl(hookwrapper=True)
def pytest_pyfunc_call(pyfuncitem):
    """
    Using pyfunc_call to be as 'close' to the actual call of the test as possible.
    This improves compatibility with other plugins.

    This is executed immediately after the test itself is called.
    """
    __tracebackhide__ = True
    outcome = None
    try:
        outcome = yield
    finally:
        pyfuncitem.__assumed = True
        assume_outcome(outcome, pyfuncitem)

def assume_outcome(outcome, item):
    """
    This is executed immediately after the test itself is called.
    Note: I'm not happy with exception handling in here.
    """
    __tracebackhide__ = True
    failed_assumptions = _FAILED_ASSUMPTIONS
    if failed_assumptions:
        failed_count = len(failed_assumptions)
        root_msg = "\n%s Failed Assumptions:\n" % failed_count

        content = pytest._hook_assume_summary_report(failed_assumptions=failed_assumptions)

        # Pluggy module returns list for multiple implementation of hooks
        # The user, while implementing custom hook pytest_assume_summary_report, will return "string"
        # Default hook is always present as list element 0
        if len(content) == 1:  # default length
            # Uses default hook
            content = content[0]
        else:
            # User created hook, if any
            content = content[1]

        last_tb = failed_assumptions[-1].tb

        del _FAILED_ASSUMPTIONS[:]
        if outcome and outcome.excinfo:
            # Xfailed test, but with strict=True. This is done via the pytest_pyfunc_call() hook, which
            # is before our hook.
            if "[XPASS(strict)]" in str(outcome.excinfo[1]):
                restore_xfail(item)
                raise_(FailedAssumption, FailedAssumption("%s\n%s" % (root_msg, content)), last_tb)
            root_msg = "\nOriginal Failure:\n\n>> %s\n" % repr(outcome.excinfo[1]) + root_msg
            raise_(
                FailedAssumption,
                FailedAssumption(root_msg + "\n" + content),
                outcome.excinfo[2],
            )
        else:
            exc = FailedAssumption(root_msg + "\n" + content)
            # Note: raising here so that we guarantee a failure.
            raise_(FailedAssumption, exc, last_tb)
<|endoftext|>


<commit_before>import pytest
import inspect
import os.path

def pytest_namespace():
    def assume(expr, msg=''):
        if not expr:
            # get filename, line, and context
            (filename, line, funcname, contextlist) = inspect.stack()[1][1:5]
            filename = os.path.relpath(filename)
            context = contextlist[0]
            # format entry
            msg = '%s\n' % msg if msg else ''
            entry = '>%s%s%s:%s\n--------' % (context, msg, filename, line)
            # add entry
            pytest._failed_assumptions.append(entry)

    return {'_failed_assumptions': [],
            'assume': assume}

def pytest_runtest_setup(item):
    del pytest._failed_assumptions[:]

@pytest.mark.hookwrapper
def pytest_runtest_makereport(item, call):
    outcome = yield
    report = outcome.get_result()
    if call.when == "call" and pytest._failed_assumptions:
        summary = 'Failed Assumptions:%s' % len(pytest._failed_assumptions)
        pytest._failed_assumptions.append(summary)
        if report.longrepr:
            report.longrepr = str(report.longrepr) + \
                              '\n--------\n' + ('\n'.join(pytest._failed_assumptions))
        else:
            report.longrepr = '\n'.join(pytest._failed_assumptions)
        report.outcome = "failed"

<commit_msg>Added locals to assume output<commit_after>import pytest
import inspect
import os.path
from py.io import saferepr

def pytest_namespace():
    def assume(expr, msg=''):
        if not expr:
            # get filename, line, and context
            (frame, filename, line, funcname, contextlist) = inspect.stack()[1][0:5]
            filename = os.path.relpath(filename)
            context = contextlist[0]
            # format entry
            msg = '%s\n' % msg if msg else ''
            entry = '>%s%s%s:%s\n' % (context, msg, filename, line)
            # add entry
            pytest._failed_assumptions.append(entry)
            pretty_locals = ["%-10s = %s" %(name, saferepr(val))
                             for name, val in frame.f_locals.items()]
            pytest._assumption_locals.append(pretty_locals)

    return {'_assumption_locals': [],
            '_failed_assumptions': [],
            'assume': assume}

def pytest_runtest_setup(item):
    del pytest._failed_assumptions[:]
    del pytest._assumption_locals[:]

@pytest.mark.hookwrapper
def pytest_runtest_makereport(item, call):
    outcome = yield
    report = outcome.get_result()
    if call.when == "call" and pytest._failed_assumptions:
        summary = 'Failed Assumptions:%s' % len(pytest._failed_assumptions)
        if report.longrepr:
            report.sections.append((summary, "\n".join(pytest._failed_assumptions)))
        else:
            assume_data = zip(pytest._failed_assumptions, pytest._assumption_locals)
            longrepr = ["{}\n{}".format(assumption, "\n".join(locals))
                        for assumption, locals in assume_data]
            longrepr.append("-" * 40)
            longrepr.append(summary)
            report.longrepr = '\n'.join(longrepr)
        report.outcome = "failed"

<|endoftext|>