<commit_before>from pythoscope.compat import all, set, sorted
from pythoscope.generator.code_string import CodeString, combine
from pythoscope.generator.constructor import constructor_as_string, call_as_string_for
from pythoscope.serializer import BuiltinException, CompositeObject,\
    ImmutableObject, MapObject, UnknownObject, SequenceObject, SerializedObject
from pythoscope.side_effect import BuiltinMethodWithPositionArgsSideEffect
from pythoscope.store import Call, FunctionCall, UserObject, MethodCall,\
    GeneratorObject, GeneratorObjectInvocation
from pythoscope.util import counted, flatten, key_for_value

# :: SerializedObject | [SerializedObject] -> bool
def can_be_constructed(obj):
    if isinstance(obj, list):
        return all(map(can_be_constructed, obj))
    elif isinstance(obj, SequenceObject):
        return all(map(can_be_constructed, obj.contained_objects))
    elif isinstance(obj, GeneratorObject):
        return obj.is_activated()
    return not isinstance(obj, UnknownObject)

# :: Call -> Call
def top_caller(call):
    if call.caller is None:
        return call
    return top_caller(call.caller)

# :: (Call, int) -> [Call]
def subcalls_before_timestamp(call, reference_timestamp):
    for c in call.subcalls:
        if c.timestamp < reference_timestamp:
            yield c
            for sc in subcalls_before_timestamp(c, reference_timestamp):
                yield sc

# :: Call -> [Call]
def calls_before(call):
    """Go up the call graph and return all calls that happened before
    the given one.

    >>> class Call(object):
    ...     def __init__(self, caller, timestamp):
    ...         self.subcalls = []
    ...         self.caller = caller
    ...         self.timestamp = timestamp
    ...         if caller:
    ...             caller.subcalls.append(self)
    >>> top = Call(None, 1)
    >>> branch1 = Call(top, 2)
    >>> leaf1 = Call(branch1, 3)
    >>> branch2 = Call(top, 4)
    >>> leaf2 = Call(branch2, 5)
    >>> leaf3 = Call(branch2, 6)
    >>> leaf4 = Call(branch2, 7)
    >>> branch3 = Call(top, 8)
    >>> calls_before(branch3) == [top, branch1, leaf1, branch2, leaf2, leaf3, leaf4]
    True
    >>> calls_before(leaf3) == [top, branch1, leaf1, branch2, leaf2]
    True
    >>> calls_before(branch2) == [top, branch1, leaf1]
    True
    >>> calls_before(branch1) == [top]
    True
    """
    top = top_caller(call)
    return [top] + list(subcalls_before_timestamp(top, call.timestamp))

# :: Call -> [SideEffect]
def side_effects_before(call):
    return flatten(map(lambda c: c.side_effects, calls_before(call)))

# :: SerializedObject | Call | [SerializedObject] | [Call] -> [SerializedObject]
def get_contained_objects(obj):
    """Return a list of SerializedObjects this object requires during testing.

    This function will descend recursively if objects contained within given
    object are composite themselves.
    """
    if isinstance(obj, list):
        return flatten(map(get_contained_objects, obj))
    elif isinstance(obj, ImmutableObject):
        # ImmutableObjects are self-sufficient.
        return []
    elif isinstance(obj, UnknownObject):
        return []
    elif isinstance(obj, SequenceObject):
        return get_those_and_contained_objects(obj.contained_objects)
    elif isinstance(obj, MapObject):
        return get_those_and_contained_objects(flatten(obj.mapping))
    elif isinstance(obj, BuiltinException):
        return get_those_and_contained_objects(obj.args)
    elif isinstance(obj, UserObject):
        return get_contained_objects(obj.get_init_and_external_calls())
    elif isinstance(obj, (FunctionCall, MethodCall, GeneratorObjectInvocation)):
        if obj.raised_exception():
            output = obj.exception
        else:
            output = obj.output
        return get_those_and_contained_objects(obj.input.values() + [output])
    elif isinstance(obj, GeneratorObject):
        if obj.is_activated():
            return get_those_and_contained_objects(obj.args.values()) +\
                get_contained_objects(obj.calls)
        else:
            return []
    else:
        raise TypeError("Wrong argument to get_contained_objects: %r." % obj)

# :: [SerializedObject] -> [SerializedObject]
def get_those_and_contained_objects(objs):
    """Return a list containing given objects and all objects contained within
    them.
    """
    return objs + get_contained_objects(objs)

# :: [SerializedObject|Call] -> [SerializedObject|Call]
def sorted_by_timestamp(objects):
    return sorted(objects, key=lambda o: o.timestamp)

# :: [SideEffect] -> [SerializedObject]
def objects_referenced_by_side_effects(side_effects):
    return flatten(map(lambda se: se.referenced_objects, side_effects))

# :: [SideEffect] -> [SerializedObject]
def objects_affected_by_side_effects(side_effects):
    return flatten(map(lambda se: se.affected_objects, side_effects))

# :: ([SideEffect], set([SerializedObject])) -> [SideEffect]
def side_effects_that_affect_objects(side_effects, objects):
    "Filter out side effects that are irrelevant to given set of objects."
    for side_effect in side_effects:
        for obj in side_effect.affected_objects:
            if obj in objects:
                yield side_effect

class Dependencies(object):
    def __init__(self, call):
        self.objects = []
        self.side_effects = set()

    def _calculate(self, objects, relevant_side_effects):
        """
        First, we look at all objects required for the call's input/output
        (pre- and post- call dependencies ech do one of those). Next, we look
        at all side effects that affected those objects before the call. Those
        side effects can reference more objects, which in turn can be affected
        by more side effects, so we do this back and forth until we have
        a complete set of objects and side effects that have direct or indirect
        relationship to the call.
        """
        def update(objects):
            self.objects.extend(objects)

            # We have some objects, let's see how many side effects affect them.
            new_side_effects = set(side_effects_that_affect_objects(relevant_side_effects, objects))
            previous_side_effects_count = len(self.side_effects)
            self.side_effects.update(new_side_effects)
            if len(self.side_effects) == previous_side_effects_count:
                return

            # Similarly, new side effects may yield some new objects, let's recur.
            update(get_those_and_contained_objects(objects_referenced_by_side_effects(new_side_effects)))
        # We start with objects required for the call itself.
        update(objects)

    def fold(self):
        return self # TODO

    def remove_objects_unworthty_of_naming(self):
        affected_objects = objects_affected_by_side_effects(self.side_effects)
        for obj, usage_count in counted(self.objects):
            # ImmutableObjects don't need to be named, as their identity is
            # always unambiguous.
            if not isinstance(obj, ImmutableObject):
                # Anything mentioned more than once have to be named.
                if usage_count > 1:
                    continue
                # Anything affected by side effects is also worth naming.
                if obj in affected_objects:
                    continue
            for i in range(usage_count):
                self.objects.remove(obj)
        return self

    def all(self):
        return list(set(self.objects).union(self.side_effects))

    def sorted(self):
        """
        Since data we have was gathered during real execution there is no way setup
        dependencies are cyclic, i.e. there is a strict order of object creation.
        We've chosen to sort objects by their creation timestamp.
        """
        return sorted_by_timestamp(self.all())

    def unique_objects(self):
        return set(self.objects)

class PreCallDependencies(Dependencies):
    """Dependencies for making a call.
    """
    def __init__(self, call):
        super(PreCallDependencies, self).__init__(call)
        self._calculate(get_contained_objects(call), side_effects_before(call))

class PostCallDependencies(Dependencies):
    """Dependencies regarding call's output value.
    """
    def __init__(self, call):
        super(PostCallDependencies, self).__init__(call)
        self._calculate(get_those_and_contained_objects([call.output]), call.side_effects)

# :: SerializedObject -> str
def get_name_base_for_object(obj):
    common_names = {'list': 'alist',
                    'dict': 'adict',
                    'array.array': 'array',
                    'types.FunctionType': 'function',
                    'types.GeneratorType': 'generator'}
    return common_names.get(obj.type_name, 'obj')

# :: [str], str -> str
def get_next_name(names, base):
    """Figure out a new name starting with base that doesn't appear in given
    list of names.

    >>> get_next_name(["alist", "adict1", "adict2"], "adict")
    'adict3'
    """
    base_length = len(base)
    def has_right_base(name):
        return name.startswith(base)
    def get_index(name):
        return int(name[base_length:])
    return base + str(max(map(get_index, filter(has_right_base, names))) + 1)

# :: SerializedObject, {SerializedObject: str} -> None
def assign_name_to_object(obj, assigned_names):
    """Assign a right name for given object.

    May reassign an existing name for an object as a side effect.
    """
    if assigned_names.has_key(obj):
        return
    base = get_name_base_for_object(obj)
    other_obj = key_for_value(assigned_names, base)

    if other_obj:
        # Avoid overlapping names by numbering objects with the same base.
        assigned_names[other_obj] = base+"1"
        assigned_names[obj] = base+"2"
    elif base+"1" in assigned_names.values():
        # We have some objects already numbered, insert a name with a new index.
        assigned_names[obj] = get_next_name(assigned_names.values(), base)
    else:
        # It's the first object with that base.
        assigned_names[obj] = base

# :: ([SerializedObject], {SerializedObject: str}) -> None
def assign_names_to_objects(objects, names):
    """Modifies names dictionary as a side effect.
    """
    for obj in sorted_by_timestamp(objects):
        assign_name_to_object(obj, names)

# :: (SerializedObject, str, {SerializedObject: str}) -> CodeString
def setup_for_named_object(obj, name, already_assigned_names):
    constructor = constructor_as_string(obj, already_assigned_names)
    setup = combine(name, constructor, "%s = %s\n")
    if constructor.uncomplete:
        setup = combine("# ", setup)
    return setup

# :: CodeString -> CodeString
def add_newline(code_string):
    return combine(code_string, "\n")

# :: (SideEffect, {SerializedObject: str}) -> CodeString
def setup_for_side_effect(side_effect, already_assigned_names):
    object_name = already_assigned_names[side_effect.obj]
    if isinstance(side_effect, BuiltinMethodWithPositionArgsSideEffect):
        return add_newline(call_as_string_for("%s.%s" % (object_name, side_effect.definition.name),
                                              side_effect.args_mapping(),
                                              side_effect.definition,
                                              already_assigned_names))
    else:
        raise TypeError("Unknown side effect type: %r" % side_effect)

# :: (Dependencies, {SerializedObject: str}) -> CodeString
def create_setup_for_dependencies(dependencies, names):
    """Returns a setup code string. Modifies names dictionary as a side effect.
    """
    already_assigned_names = names.copy()
    assign_names_to_objects(dependencies.unique_objects(), names)
    full_setup = CodeString("")
    for dependency in dependencies.sorted():
        if isinstance(dependency, SerializedObject):
            name = names[dependency]
            setup = setup_for_named_object(dependency, name, already_assigned_names)
            already_assigned_names[dependency] = name
        else:
            setup = setup_for_side_effect(dependency, already_assigned_names)
        full_setup = combine(full_setup, setup)
    return full_setup

# :: (Call, {SerializedObject : str}) -> CodeString
def assign_names_and_setup(call, names):
    """Returns a setup code string. Modifies names dictionary as a side effect.
    """
    if not isinstance(call, Call):
        raise TypeError("Tried to call assign_names_and_setup with %r instead of a call." % call)
    pre_setup = create_setup_for_input(call, names)
    post_setup = create_setup_for_output(call, names)
    return combine(pre_setup, post_setup)

# :: ([Call], {SerializedObject : str}) -> CodeString
def assign_names_and_setup_for_multiple_calls(calls, names):
    """Returns a setup code string. Modifies names dictionary as a side effect.
    """
    full_setup = CodeString("")
    for call in sorted_by_timestamp(calls):
        setup = assign_names_and_setup(call, names)
        full_setup = combine(full_setup, setup)
    return full_setup

# :: (Call, {SerializedObject : str}) -> CodeString
def create_setup_for_input(call, names):
    pre_dependencies = PreCallDependencies(call).fold().remove_objects_unworthty_of_naming()
    return create_setup_for_dependencies(pre_dependencies, names)

# :: (Call, {SerializedObject : str}) -> CodeString|str
def create_setup_for_output(call, names):
    """Sometimes the call's output doesn't depend on the input at all. If the
    object has been altered by the call we need to create and name a new copy of
    it to prepare a good equality assertion.
    """
    if call.output is not None and \
            can_be_constructed(call.output) and \
            call.output not in names.keys():
        post_dependencies = PostCallDependencies(call).fold().remove_objects_unworthty_of_naming()
        return create_setup_for_dependencies(post_dependencies, names)
    return ""
<commit_msg>  Fixed a typo.<commit_after>from pythoscope.compat import all, set, sorted
from pythoscope.generator.code_string import CodeString, combine
from pythoscope.generator.constructor import constructor_as_string, call_as_string_for
from pythoscope.serializer import BuiltinException, CompositeObject,\
    ImmutableObject, MapObject, UnknownObject, SequenceObject, SerializedObject
from pythoscope.side_effect import BuiltinMethodWithPositionArgsSideEffect
from pythoscope.store import Call, FunctionCall, UserObject, MethodCall,\
    GeneratorObject, GeneratorObjectInvocation
from pythoscope.util import counted, flatten, key_for_value

# :: SerializedObject | [SerializedObject] -> bool
def can_be_constructed(obj):
    if isinstance(obj, list):
        return all(map(can_be_constructed, obj))
    elif isinstance(obj, SequenceObject):
        return all(map(can_be_constructed, obj.contained_objects))
    elif isinstance(obj, GeneratorObject):
        return obj.is_activated()
    return not isinstance(obj, UnknownObject)

# :: Call -> Call
def top_caller(call):
    if call.caller is None:
        return call
    return top_caller(call.caller)

# :: (Call, int) -> [Call]
def subcalls_before_timestamp(call, reference_timestamp):
    for c in call.subcalls:
        if c.timestamp < reference_timestamp:
            yield c
            for sc in subcalls_before_timestamp(c, reference_timestamp):
                yield sc

# :: Call -> [Call]
def calls_before(call):
    """Go up the call graph and return all calls that happened before
    the given one.

    >>> class Call(object):
    ...     def __init__(self, caller, timestamp):
    ...         self.subcalls = []
    ...         self.caller = caller
    ...         self.timestamp = timestamp
    ...         if caller:
    ...             caller.subcalls.append(self)
    >>> top = Call(None, 1)
    >>> branch1 = Call(top, 2)
    >>> leaf1 = Call(branch1, 3)
    >>> branch2 = Call(top, 4)
    >>> leaf2 = Call(branch2, 5)
    >>> leaf3 = Call(branch2, 6)
    >>> leaf4 = Call(branch2, 7)
    >>> branch3 = Call(top, 8)
    >>> calls_before(branch3) == [top, branch1, leaf1, branch2, leaf2, leaf3, leaf4]
    True
    >>> calls_before(leaf3) == [top, branch1, leaf1, branch2, leaf2]
    True
    >>> calls_before(branch2) == [top, branch1, leaf1]
    True
    >>> calls_before(branch1) == [top]
    True
    """
    top = top_caller(call)
    return [top] + list(subcalls_before_timestamp(top, call.timestamp))

# :: Call -> [SideEffect]
def side_effects_before(call):
    return flatten(map(lambda c: c.side_effects, calls_before(call)))

# :: SerializedObject | Call | [SerializedObject] | [Call] -> [SerializedObject]
def get_contained_objects(obj):
    """Return a list of SerializedObjects this object requires during testing.

    This function will descend recursively if objects contained within given
    object are composite themselves.
    """
    if isinstance(obj, list):
        return flatten(map(get_contained_objects, obj))
    elif isinstance(obj, ImmutableObject):
        # ImmutableObjects are self-sufficient.
        return []
    elif isinstance(obj, UnknownObject):
        return []
    elif isinstance(obj, SequenceObject):
        return get_those_and_contained_objects(obj.contained_objects)
    elif isinstance(obj, MapObject):
        return get_those_and_contained_objects(flatten(obj.mapping))
    elif isinstance(obj, BuiltinException):
        return get_those_and_contained_objects(obj.args)
    elif isinstance(obj, UserObject):
        return get_contained_objects(obj.get_init_and_external_calls())
    elif isinstance(obj, (FunctionCall, MethodCall, GeneratorObjectInvocation)):
        if obj.raised_exception():
            output = obj.exception
        else:
            output = obj.output
        return get_those_and_contained_objects(obj.input.values() + [output])
    elif isinstance(obj, GeneratorObject):
        if obj.is_activated():
            return get_those_and_contained_objects(obj.args.values()) +\
                get_contained_objects(obj.calls)
        else:
            return []
    else:
        raise TypeError("Wrong argument to get_contained_objects: %r." % obj)

# :: [SerializedObject] -> [SerializedObject]
def get_those_and_contained_objects(objs):
    """Return a list containing given objects and all objects contained within
    them.
    """
    return objs + get_contained_objects(objs)

# :: [SerializedObject|Call] -> [SerializedObject|Call]
def sorted_by_timestamp(objects):
    return sorted(objects, key=lambda o: o.timestamp)

# :: [SideEffect] -> [SerializedObject]
def objects_referenced_by_side_effects(side_effects):
    return flatten(map(lambda se: se.referenced_objects, side_effects))

# :: [SideEffect] -> [SerializedObject]
def objects_affected_by_side_effects(side_effects):
    return flatten(map(lambda se: se.affected_objects, side_effects))

# :: ([SideEffect], set([SerializedObject])) -> [SideEffect]
def side_effects_that_affect_objects(side_effects, objects):
    "Filter out side effects that are irrelevant to given set of objects."
    for side_effect in side_effects:
        for obj in side_effect.affected_objects:
            if obj in objects:
                yield side_effect

class Dependencies(object):
    def __init__(self, call):
        self.objects = []
        self.side_effects = set()

    def _calculate(self, objects, relevant_side_effects):
        """
        First, we look at all objects required for the call's input/output
        (pre- and post- call dependencies each do one of those). Next, we look
        at all side effects that affected those objects before the call. Those
        side effects can reference more objects, which in turn can be affected
        by more side effects, so we do this back and forth until we have
        a complete set of objects and side effects that have direct or indirect
        relationship to the call.
        """
        def update(objects):
            self.objects.extend(objects)

            # We have some objects, let's see how many side effects affect them.
            new_side_effects = set(side_effects_that_affect_objects(relevant_side_effects, objects))
            previous_side_effects_count = len(self.side_effects)
            self.side_effects.update(new_side_effects)
            if len(self.side_effects) == previous_side_effects_count:
                return

            # Similarly, new side effects may yield some new objects, let's recur.
            update(get_those_and_contained_objects(objects_referenced_by_side_effects(new_side_effects)))
        # We start with objects required for the call itself.
        update(objects)

    def fold(self):
        return self # TODO

    def remove_objects_unworthty_of_naming(self):
        affected_objects = objects_affected_by_side_effects(self.side_effects)
        for obj, usage_count in counted(self.objects):
            # ImmutableObjects don't need to be named, as their identity is
            # always unambiguous.
            if not isinstance(obj, ImmutableObject):
                # Anything mentioned more than once have to be named.
                if usage_count > 1:
                    continue
                # Anything affected by side effects is also worth naming.
                if obj in affected_objects:
                    continue
            for i in range(usage_count):
                self.objects.remove(obj)
        return self

    def all(self):
        return list(set(self.objects).union(self.side_effects))

    def sorted(self):
        """
        Since data we have was gathered during real execution there is no way setup
        dependencies are cyclic, i.e. there is a strict order of object creation.
        We've chosen to sort objects by their creation timestamp.
        """
        return sorted_by_timestamp(self.all())

    def unique_objects(self):
        return set(self.objects)

class PreCallDependencies(Dependencies):
    """Dependencies for making a call.
    """
    def __init__(self, call):
        super(PreCallDependencies, self).__init__(call)
        self._calculate(get_contained_objects(call), side_effects_before(call))

class PostCallDependencies(Dependencies):
    """Dependencies regarding call's output value.
    """
    def __init__(self, call):
        super(PostCallDependencies, self).__init__(call)
        self._calculate(get_those_and_contained_objects([call.output]), call.side_effects)

# :: SerializedObject -> str
def get_name_base_for_object(obj):
    common_names = {'list': 'alist',
                    'dict': 'adict',
                    'array.array': 'array',
                    'types.FunctionType': 'function',
                    'types.GeneratorType': 'generator'}
    return common_names.get(obj.type_name, 'obj')

# :: [str], str -> str
def get_next_name(names, base):
    """Figure out a new name starting with base that doesn't appear in given
    list of names.

    >>> get_next_name(["alist", "adict1", "adict2"], "adict")
    'adict3'
    """
    base_length = len(base)
    def has_right_base(name):
        return name.startswith(base)
    def get_index(name):
        return int(name[base_length:])
    return base + str(max(map(get_index, filter(has_right_base, names))) + 1)

# :: SerializedObject, {SerializedObject: str} -> None
def assign_name_to_object(obj, assigned_names):
    """Assign a right name for given object.

    May reassign an existing name for an object as a side effect.
    """
    if assigned_names.has_key(obj):
        return
    base = get_name_base_for_object(obj)
    other_obj = key_for_value(assigned_names, base)

    if other_obj:
        # Avoid overlapping names by numbering objects with the same base.
        assigned_names[other_obj] = base+"1"
        assigned_names[obj] = base+"2"
    elif base+"1" in assigned_names.values():
        # We have some objects already numbered, insert a name with a new index.
        assigned_names[obj] = get_next_name(assigned_names.values(), base)
    else:
        # It's the first object with that base.
        assigned_names[obj] = base

# :: ([SerializedObject], {SerializedObject: str}) -> None
def assign_names_to_objects(objects, names):
    """Modifies names dictionary as a side effect.
    """
    for obj in sorted_by_timestamp(objects):
        assign_name_to_object(obj, names)

# :: (SerializedObject, str, {SerializedObject: str}) -> CodeString
def setup_for_named_object(obj, name, already_assigned_names):
    constructor = constructor_as_string(obj, already_assigned_names)
    setup = combine(name, constructor, "%s = %s\n")
    if constructor.uncomplete:
        setup = combine("# ", setup)
    return setup

# :: CodeString -> CodeString
def add_newline(code_string):
    return combine(code_string, "\n")

# :: (SideEffect, {SerializedObject: str}) -> CodeString
def setup_for_side_effect(side_effect, already_assigned_names):
    object_name = already_assigned_names[side_effect.obj]
    if isinstance(side_effect, BuiltinMethodWithPositionArgsSideEffect):
        return add_newline(call_as_string_for("%s.%s" % (object_name, side_effect.definition.name),
                                              side_effect.args_mapping(),
                                              side_effect.definition,
                                              already_assigned_names))
    else:
        raise TypeError("Unknown side effect type: %r" % side_effect)

# :: (Dependencies, {SerializedObject: str}) -> CodeString
def create_setup_for_dependencies(dependencies, names):
    """Returns a setup code string. Modifies names dictionary as a side effect.
    """
    already_assigned_names = names.copy()
    assign_names_to_objects(dependencies.unique_objects(), names)
    full_setup = CodeString("")
    for dependency in dependencies.sorted():
        if isinstance(dependency, SerializedObject):
            name = names[dependency]
            setup = setup_for_named_object(dependency, name, already_assigned_names)
            already_assigned_names[dependency] = name
        else:
            setup = setup_for_side_effect(dependency, already_assigned_names)
        full_setup = combine(full_setup, setup)
    return full_setup

# :: (Call, {SerializedObject : str}) -> CodeString
def assign_names_and_setup(call, names):
    """Returns a setup code string. Modifies names dictionary as a side effect.
    """
    if not isinstance(call, Call):
        raise TypeError("Tried to call assign_names_and_setup with %r instead of a call." % call)
    pre_setup = create_setup_for_input(call, names)
    post_setup = create_setup_for_output(call, names)
    return combine(pre_setup, post_setup)

# :: ([Call], {SerializedObject : str}) -> CodeString
def assign_names_and_setup_for_multiple_calls(calls, names):
    """Returns a setup code string. Modifies names dictionary as a side effect.
    """
    full_setup = CodeString("")
    for call in sorted_by_timestamp(calls):
        setup = assign_names_and_setup(call, names)
        full_setup = combine(full_setup, setup)
    return full_setup

# :: (Call, {SerializedObject : str}) -> CodeString
def create_setup_for_input(call, names):
    pre_dependencies = PreCallDependencies(call).fold().remove_objects_unworthty_of_naming()
    return create_setup_for_dependencies(pre_dependencies, names)

# :: (Call, {SerializedObject : str}) -> CodeString|str
def create_setup_for_output(call, names):
    """Sometimes the call's output doesn't depend on the input at all. If the
    object has been altered by the call we need to create and name a new copy of
    it to prepare a good equality assertion.
    """
    if call.output is not None and \
            can_be_constructed(call.output) and \
            call.output not in names.keys():
        post_dependencies = PostCallDependencies(call).fold().remove_objects_unworthty_of_naming()
        return create_setup_for_dependencies(post_dependencies, names)
    return ""
<|endoftext|>


<commit_before># ----------------------------------------------------------------------------
# Copyright (c) 2016-2019, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

import qiime2.plugin
import qiime2.plugin.model as model

from ._taxonomic_classifier import PickleFormat, JSONFormat
from .plugin_setup import plugin

# Semantic Types
ClassifierSpecification = qiime2.plugin.SemanticType('ClassifierSpecification')
KerasClassifier = qiime2.plugin.SemanticType('KerasClassifier')

class ClassifierSpecificationDirFmt(model.DirectoryFormat):
    classifier_specification = model.File(
        'classifier_specification.json', format=JSONFormat)

class KerasClassifierPickleDirFmt(model.DirectoryFormat):
    y_encoder = model.File('y_encoder.json', format=JSONFormat)
    x_encoder = model.File('x_encoder.json', format=JSONFormat)
    keras_pipeline = model.File('keras_model.tar', format=PickleFormat)

# Transformers
@plugin.regist_transformer
def _1(dirfmt: KerasClassifierPickleDirFmt) -> KerasPipeline:
    x_encoder_spec = dirfmt.x_encoder.view(object)
    y_encoder_spec = dirfmt.y_encoder.view(object)

<commit_msg>first draft<commit_after># ----------------------------------------------------------------------------
# Copyright (c) 2016-2019, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

import qiime2.plugin
import qiime2.plugin.model as model
import h5py
from keras.model import load_model

from ._taxonomic_classifier import PickleFormat, JSONFormat
from .plugin_setup import plugin

# Semantic Types
ClassifierSpecification = qiime2.plugin.SemanticType('ClassifierSpecification')
KerasClassifier = qiime2.plugin.SemanticType('KerasClassifier')

# Formats
ClassifierSpecificationDirectoryFormat = model.SingleFileDirectoryFormat(
    'ClassifierSpecificationDirectoryFormat', 'classifier-specification.json',
    JSONFormat)

class KerasModelFormat(model.BinaryFileFormat):
    def sniff(self):
        return h5py.is_hdf5(str(self))

class KerasClassifierDirectoryFormat(model.DirectoryFormat):
    y_encoder = model.File('y_encoder.json', format=JSONFormat)
    x_encoder = model.File('x_encoder.json', format=JSONFormat)
    keras_model = model.File('keras_model.hdf5', format=KerasModelFormat)

# Transformers
@plugin.register_transformer
def _1(dirfmt: KerasClassifierDirectoryFomat) -> tuple:
    x_encoder = dirfmt.x_encoder.view(object)
    y_encoder = dirfmt.y_encoder.view(object)
    keras_model_format = dirfmt.keras_model.view(KerasModelFormat)
    keras_model = load_model(str(keras_model_format))
    return (x_encoder, y_encoder, keras_model)

@plugin.register_transformer
def _2(data: tuple) -> KerasClassifierDirectoryFormat:
    x_encoder, y_encoder, keras_model = data
    dirfmt = KerasClassifierDirectoryFormat()
    dirfmt.x_encoder.write_data(x_encoder, object)
    dirfmt.y_encoder.write_data(y_encoder, object)

    keras_model_format = KerasModelFormat()
    keras_model.save(str(keras_model_format))
    dirfmt.keras_model.write_data(keras_model_format, KerasModelFormat)

    return dirfmt

# Registrations
plugin.register_semantic_types(ClassifierSpecification)
plugin.register_formats(ClassifierSpecificationDirectoryFormat)
plugin.register_semantic_type_to_format(
    ClassifierSpecification,
    artifact_format=ClassifierSpecificationDirectoryFormat
)

<|endoftext|>


<commit_before># ----------------------------------------------------------------------------
# Copyright (c) 2016-2017, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

import os
import pkg_resources
import shutil

import biom
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
from q2_types.feature_data import DNAIterator
import q2templates
import skbio
import qiime2

_blast_url_template = ("http://www.ncbi.nlm.nih.gov/BLAST/Blast.cgi?"
                       "ALIGNMENT_VIEW=Pairwise&PROGRAM=blastn&DATABASE"
                       "=nt&CMD=Put&QUERY=%s")

TEMPLATES = pkg_resources.resource_filename('q2_feature_table', '_summarize')

def tabulate_seqs(output_dir: str, data: DNAIterator) -> None:
    sequences = []
    with open(os.path.join(output_dir, 'sequences.fasta'), 'w') as fh:
        for sequence in data:
            skbio.io.write(sequence, format='fasta', into=fh)
            str_seq = str(sequence)
            sequences.append({'id': sequence.metadata['id'],
                              'url': _blast_url_template % str_seq,
                              'seq': str_seq})

    index = os.path.join(TEMPLATES, 'tabulate_seqs_assets', 'index.html')
    q2templates.render(index, output_dir, context={'data': sequences})

    js = os.path.join(
        TEMPLATES, 'tabulate_seqs_assets', 'js', 'tsorter.min.js')
    os.mkdir(os.path.join(output_dir, 'js'))
    shutil.copy(js, os.path.join(output_dir, 'js', 'tsorter.min.js'))

def summarize(output_dir: str, table: biom.Table,
              sample_metadata: qiime2.Metadata=None) -> None:
    number_of_features, number_of_samples = table.shape

    sample_summary, sample_frequencies = _frequency_summary(
        table, axis='sample')
    if number_of_samples > 1:
        # Calculate the bin count, with a minimum of 5 bins
        IQR = sample_summary['3rd quartile'] - sample_summary['1st quartile']
        if IQR == 0.0:
            bins = 5
        else:
            # Freedman–Diaconis rule
            bin_width = (2 * IQR) / (number_of_samples ** (1/3))

            bins = max((sample_summary['Maximum frequency'] -
                        sample_summary['Minimum frequency']) / bin_width, 5)

        sample_frequencies_ax = sns.distplot(sample_frequencies, kde=False,
                                             rug=True, bins=int(round(bins)))
        sample_frequencies_ax.get_xaxis().set_major_formatter(
            matplotlib.ticker.FuncFormatter(lambda x, p: format(int(x), ',')))
        sample_frequencies_ax.set_xlabel('Frequency per sample')
        sample_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'sample-frequencies.pdf'))
        sample_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'sample-frequencies.png'))
        plt.gcf().clear()

    feature_summary, feature_frequencies = _frequency_summary(
        table, axis='observation')
    if number_of_features > 1:
        feature_frequencies_ax = sns.distplot(feature_frequencies, kde=False,
                                              rug=True)
        feature_frequencies_ax.set_xlabel('Frequency per feature')
        feature_frequencies_ax.set_xscale('log')
        feature_frequencies_ax.set_yscale('log')
        feature_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'feature-frequencies.pdf'))
        feature_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'feature-frequencies.png'))

    sample_summary_table = q2templates.df_to_html(
        sample_summary.apply('{:,}'.format).to_frame('Frequency'))
    feature_summary_table = q2templates.df_to_html(
        feature_summary.apply('{:,}'.format).to_frame('Frequency'))

    index = os.path.join(TEMPLATES, 'summarize_assets', 'index.html')
    context = {
        'number_of_samples': number_of_samples,
        'number_of_features': number_of_features,
        'total_frequencies': int(np.sum(sample_frequencies)),
        'sample_summary_table': sample_summary_table,
        'feature_summary_table': feature_summary_table,
    }

    feature_qualitative_data = _compute_qualitative_summary(table)
    sample_frequencies.sort_values(inplace=True, ascending=False)
    feature_frequencies.sort_values(inplace=True, ascending=False)
    sample_frequencies.to_csv(
        os.path.join(output_dir, 'sample-frequency-detail.csv'))
    feature_frequencies.to_csv(
        os.path.join(output_dir, 'feature-frequency-detail.csv'))

    feature_frequencies = feature_frequencies.astype(int) \
        .apply('{:,}'.format).to_frame('Frequency')
    feature_frequencies['# of Samples Observed In'] = \
        pd.Series(feature_qualitative_data).astype(int).apply('{:,}'.format)
    feature_frequencies_table = q2templates.df_to_html(feature_frequencies)
    overview_template = os.path.join(
        TEMPLATES, 'summarize_assets', 'overview.html')
    sample_frequency_template = os.path.join(
        TEMPLATES, 'summarize_assets', 'sample-frequency-detail.html')
    feature_frequency_template = os.path.join(
        TEMPLATES, 'summarize_assets', 'feature-frequency-detail.html')

    context.update({'max_count': sample_frequencies.max(),
                    'feature_frequencies_table': feature_frequencies_table,
                    'feature_qualitative_data': feature_qualitative_data,
                    'tabs': [{'url': 'overview.html',
                              'title': 'Overview'},
                             {'url': 'sample-frequency-detail.html',
                              'title': 'Interactive Sample Detail'},
                             {'url': 'feature-frequency-detail.html',
                              'title': 'Feature Detail'}]})
    templates = [index, sample_frequency_template,
                 feature_frequency_template, overview_template]
    q2templates.render(templates, output_dir, context=context)

    shutil.copytree(os.path.join(TEMPLATES, 'summarize_assets', 'dist'),
                    os.path.join(output_dir, 'dist'))

    with open(os.path.join(output_dir, 'data.jsonp'), 'w') as fh:
        fh.write("app.init(")
        if sample_metadata:
            df = sample_metadata.to_dataframe()
            df.loc[sample_frequencies.index].to_json(fh)
        else:
            fh.write('{}')
        fh.write(', ')
        sample_frequencies.to_json(fh)
        fh.write(');')

def _compute_qualitative_summary(table):
    table = table.transpose()
    sample_count = {}
    for count_vector, feature_id, metadata in table.iter():
        sample_count[feature_id] = (count_vector != 0).sum()
    return sample_count

def _frequencies(table, axis):
    return pd.Series(data=table.sum(axis=axis), index=table.ids(axis=axis))

def _frequency_summary(table, axis='sample'):
    frequencies = _frequencies(table, axis=axis)

    summary = pd.Series([frequencies.min(), frequencies.quantile(0.25),
                         frequencies.median(), frequencies.quantile(0.75),
                         frequencies.max(), frequencies.mean()],
                        index=['Minimum frequency', '1st quartile',
                               'Median frequency', '3rd quartile',
                               'Maximum frequency', 'Mean frequency'])
    return summary, frequencies
<commit_msg>BUG: Add y-axis labels to `summarize` plots<commit_after># ----------------------------------------------------------------------------
# Copyright (c) 2016-2017, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

import os
import pkg_resources
import shutil

import biom
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
from q2_types.feature_data import DNAIterator
import q2templates
import skbio
import qiime2

_blast_url_template = ("http://www.ncbi.nlm.nih.gov/BLAST/Blast.cgi?"
                       "ALIGNMENT_VIEW=Pairwise&PROGRAM=blastn&DATABASE"
                       "=nt&CMD=Put&QUERY=%s")

TEMPLATES = pkg_resources.resource_filename('q2_feature_table', '_summarize')

def tabulate_seqs(output_dir: str, data: DNAIterator) -> None:
    sequences = []
    with open(os.path.join(output_dir, 'sequences.fasta'), 'w') as fh:
        for sequence in data:
            skbio.io.write(sequence, format='fasta', into=fh)
            str_seq = str(sequence)
            sequences.append({'id': sequence.metadata['id'],
                              'url': _blast_url_template % str_seq,
                              'seq': str_seq})

    index = os.path.join(TEMPLATES, 'tabulate_seqs_assets', 'index.html')
    q2templates.render(index, output_dir, context={'data': sequences})

    js = os.path.join(
        TEMPLATES, 'tabulate_seqs_assets', 'js', 'tsorter.min.js')
    os.mkdir(os.path.join(output_dir, 'js'))
    shutil.copy(js, os.path.join(output_dir, 'js', 'tsorter.min.js'))

def summarize(output_dir: str, table: biom.Table,
              sample_metadata: qiime2.Metadata=None) -> None:
    number_of_features, number_of_samples = table.shape

    sample_summary, sample_frequencies = _frequency_summary(
        table, axis='sample')
    if number_of_samples > 1:
        # Calculate the bin count, with a minimum of 5 bins
        IQR = sample_summary['3rd quartile'] - sample_summary['1st quartile']
        if IQR == 0.0:
            bins = 5
        else:
            # Freedman–Diaconis rule
            bin_width = (2 * IQR) / (number_of_samples ** (1/3))

            bins = max((sample_summary['Maximum frequency'] -
                        sample_summary['Minimum frequency']) / bin_width, 5)

        sample_frequencies_ax = sns.distplot(sample_frequencies, kde=False,
                                             rug=True, bins=int(round(bins)))
        sample_frequencies_ax.get_xaxis().set_major_formatter(
            matplotlib.ticker.FuncFormatter(lambda x, p: format(int(x), ',')))
        sample_frequencies_ax.set_xlabel('Frequency per sample')
        sample_frequencies_ax.set_ylabel('Number of samples')
        sample_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'sample-frequencies.pdf'))
        sample_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'sample-frequencies.png'))
        plt.gcf().clear()

    feature_summary, feature_frequencies = _frequency_summary(
        table, axis='observation')
    if number_of_features > 1:
        feature_frequencies_ax = sns.distplot(feature_frequencies, kde=False,
                                              rug=True)
        feature_frequencies_ax.set_xlabel('Frequency per feature')
        feature_frequencies_ax.set_ylabel('Number of features')
        feature_frequencies_ax.set_xscale('log')
        feature_frequencies_ax.set_yscale('log')
        feature_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'feature-frequencies.pdf'))
        feature_frequencies_ax.get_figure().savefig(
            os.path.join(output_dir, 'feature-frequencies.png'))

    sample_summary_table = q2templates.df_to_html(
        sample_summary.apply('{:,}'.format).to_frame('Frequency'))
    feature_summary_table = q2templates.df_to_html(
        feature_summary.apply('{:,}'.format).to_frame('Frequency'))

    index = os.path.join(TEMPLATES, 'summarize_assets', 'index.html')
    context = {
        'number_of_samples': number_of_samples,
        'number_of_features': number_of_features,
        'total_frequencies': int(np.sum(sample_frequencies)),
        'sample_summary_table': sample_summary_table,
        'feature_summary_table': feature_summary_table,
    }

    feature_qualitative_data = _compute_qualitative_summary(table)
    sample_frequencies.sort_values(inplace=True, ascending=False)
    feature_frequencies.sort_values(inplace=True, ascending=False)
    sample_frequencies.to_csv(
        os.path.join(output_dir, 'sample-frequency-detail.csv'))
    feature_frequencies.to_csv(
        os.path.join(output_dir, 'feature-frequency-detail.csv'))

    feature_frequencies = feature_frequencies.astype(int) \
        .apply('{:,}'.format).to_frame('Frequency')
    feature_frequencies['# of Samples Observed In'] = \
        pd.Series(feature_qualitative_data).astype(int).apply('{:,}'.format)
    feature_frequencies_table = q2templates.df_to_html(feature_frequencies)
    overview_template = os.path.join(
        TEMPLATES, 'summarize_assets', 'overview.html')
    sample_frequency_template = os.path.join(
        TEMPLATES, 'summarize_assets', 'sample-frequency-detail.html')
    feature_frequency_template = os.path.join(
        TEMPLATES, 'summarize_assets', 'feature-frequency-detail.html')

    context.update({'max_count': sample_frequencies.max(),
                    'feature_frequencies_table': feature_frequencies_table,
                    'feature_qualitative_data': feature_qualitative_data,
                    'tabs': [{'url': 'overview.html',
                              'title': 'Overview'},
                             {'url': 'sample-frequency-detail.html',
                              'title': 'Interactive Sample Detail'},
                             {'url': 'feature-frequency-detail.html',
                              'title': 'Feature Detail'}]})
    templates = [index, sample_frequency_template,
                 feature_frequency_template, overview_template]
    q2templates.render(templates, output_dir, context=context)

    shutil.copytree(os.path.join(TEMPLATES, 'summarize_assets', 'dist'),
                    os.path.join(output_dir, 'dist'))

    with open(os.path.join(output_dir, 'data.jsonp'), 'w') as fh:
        fh.write("app.init(")
        if sample_metadata:
            df = sample_metadata.to_dataframe()
            df.loc[sample_frequencies.index].to_json(fh)
        else:
            fh.write('{}')
        fh.write(', ')
        sample_frequencies.to_json(fh)
        fh.write(');')

def _compute_qualitative_summary(table):
    table = table.transpose()
    sample_count = {}
    for count_vector, feature_id, metadata in table.iter():
        sample_count[feature_id] = (count_vector != 0).sum()
    return sample_count

def _frequencies(table, axis):
    return pd.Series(data=table.sum(axis=axis), index=table.ids(axis=axis))

def _frequency_summary(table, axis='sample'):
    frequencies = _frequencies(table, axis=axis)

    summary = pd.Series([frequencies.min(), frequencies.quantile(0.25),
                         frequencies.median(), frequencies.quantile(0.75),
                         frequencies.max(), frequencies.mean()],
                        index=['Minimum frequency', '1st quartile',
                               'Median frequency', '3rd quartile',
                               'Maximum frequency', 'Mean frequency'])
    return summary, frequencies
<|endoftext|>


<commit_before>"""
This is the Scrapy engine which controls the Scheduler, Downloader and Spiders.

For more information see docs/topics/architecture.rst

"""
from time import time

from twisted.internet import reactor, defer
from twisted.python.failure import Failure

from scrapy import log, signals
from scrapy.stats import stats
from scrapy.core.downloader import Downloader
from scrapy.core.scraper import Scraper
from scrapy.exceptions import DontCloseSpider
from scrapy.http import Response, Request
from scrapy.utils.misc import load_object
from scrapy.utils.signal import send_catch_log, send_catch_log_deferred
from scrapy.utils.defer import mustbe_deferred

class Slot(object):

    def __init__(self):
        self.closing = False
        self.inprogress = set() # requests in progress

    def add_request(self, request):
        self.inprogress.add(request)

    def remove_request(self, request):
        self.inprogress.remove(request)
        self._maybe_fire_closing()

    def close(self):
        self.closing = defer.Deferred()
        self._maybe_fire_closing()
        return self.closing

    def _maybe_fire_closing(self):
        if self.closing and not self.inprogress:
            self.closing.callback(None)

class ExecutionEngine(object):

    def __init__(self, settings, spider_closed_callback):
        self.settings = settings
        self.slots = {}
        self.running = False
        self.paused = False
        self._next_request_calls = {}
        self.scheduler = load_object(settings['SCHEDULER'])()
        self.downloader = Downloader()
        self.scraper = Scraper(self, self.settings)
        self._spider_closed_callback = spider_closed_callback

    @defer.inlineCallbacks
    def start(self):
        """Start the execution engine"""
        assert not self.running, "Engine already running"
        self.start_time = time()
        yield send_catch_log_deferred(signal=signals.engine_started)
        self.running = True

    def stop(self):
        """Stop the execution engine gracefully"""
        assert self.running, "Engine not running"
        self.running = False
        dfd = self._close_all_spiders()
        return dfd.addBoth(lambda _: self._finish_stopping_engine())

    def pause(self):
        """Pause the execution engine"""
        self.paused = True

    def unpause(self):
        """Resume the execution engine"""
        self.paused = False

    def is_idle(self):
        return self.scheduler.is_idle() and self.downloader.is_idle() and \
            self.scraper.is_idle()

    def next_request(self, spider, now=False):
        """Scrape the next request for the spider passed.

        The next request to be scraped is retrieved from the scheduler and
        requested from the downloader.

        The spider is closed if there are no more pages to scrape.
        """
        if now:
            self._next_request_calls.pop(spider, None)
        elif spider not in self._next_request_calls:
            call = reactor.callLater(0, self.next_request, spider, now=True)
            self._next_request_calls[spider] = call
            return call
        else:
            return

        if self.paused:
            return reactor.callLater(5, self.next_request, spider)

        while not self._needs_backout(spider):
            if not self._next_request(spider):
                break

        if self.spider_is_idle(spider):
            self._spider_idle(spider)

    def _needs_backout(self, spider):
        slot = self.slots[spider]
        return not self.running \
            or slot.closing \
            or self.spider_is_closed(spider) \
            or self.downloader.sites[spider].needs_backout() \
            or self.scraper.sites[spider].needs_backout()

    def _next_request(self, spider):
        request = self.scheduler.next_request(spider)
        if not request:
            return
        d = self._download(request, spider)
        d.addBoth(self._handle_downloader_output, request, spider)
        d.addErrback(log.msg, spider=spider)
        slot = self.slots[spider]
        d.addBoth(lambda _: slot.remove_request(request))
        d.addErrback(log.msg, spider=spider)
        d.addBoth(lambda _: self.next_request(spider))
        return d

    def _handle_downloader_output(self, response, request, spider):
        assert isinstance(response, (Request, Response, Failure)), response
        # downloader middleware can return requests (for example, redirects)
        if isinstance(response, Request):
            self.crawl(response, spider)
            return
        # response is a Response or Failure
        d = defer.Deferred()
        d.addBoth(self.scraper.enqueue_scrape, request, spider)
        d.addErrback(log.err, spider=spider)
        if isinstance(response, Failure):
            d.errback(response)
        else:
            d.callback(response)
        return d

    def spider_is_idle(self, spider):
        scraper_idle = spider in self.scraper.sites \
            and self.scraper.sites[spider].is_idle()
        pending = self.scheduler.spider_has_pending_requests(spider)
        downloading = spider in self.downloader.sites \
            and self.downloader.sites[spider].active
        return scraper_idle and not (pending or downloading)

    def spider_is_closed(self, spider):
        """Return True if the spider is fully closed (ie. not even in the
        closing stage)"""
        return spider not in self.downloader.sites

    @property
    def open_spiders(self):
        return self.downloader.sites.keys()

    def has_capacity(self):
        """Does the engine have capacity to handle more spiders"""
        return len(self.downloader.sites) < self.downloader.concurrent_spiders

    def crawl(self, request, spider):
        assert spider in self.open_spiders, \
            "Spider %r not opened when crawling: %s" % (spider.name, request)
        self.schedule(request, spider)
        self.next_request(spider)

    def schedule(self, request, spider):
        return self.scheduler.enqueue_request(spider, request)

    def download(self, request, spider):
        slot = self.slots[spider]
        slot.add_request(request)
        if isinstance(request, Response):
            return request
        d = self._download(request, spider)
        d.addCallback(self.download, spider)
        d.addBoth(self._remove_request, slot, request)
        return d

    def _remove_request(self, _, slot, request):
        slot.remove_request(request)
        return _

    def _download(self, request, spider):
        slot = self.slots[spider]
        slot.add_request(request)
        def _on_success(response):
            """handle the result of a page download"""
            assert isinstance(response, (Response, Request))
            if isinstance(response, Response):
                response.request = request # tie request to response received
                log.msg(log.formatter.crawled(request, response, spider), \
                    level=log.DEBUG, spider=spider)
                send_catch_log(signal=signals.response_received, \
                    response=response, request=request, spider=spider)
            return response

        def _on_complete(_):
            self.next_request(spider)
            return _

        dwld = mustbe_deferred(self.downloader.fetch, request, spider)
        dwld.addCallback(_on_success)
        dwld.addBoth(_on_complete)
        return dwld

    @defer.inlineCallbacks
    def open_spider(self, spider):
        assert self.has_capacity(), "No free spider slots when opening %r" % \
            spider.name
        log.msg("Spider opened", spider=spider)
        self.slots[spider] = Slot()
        yield self.scheduler.open_spider(spider)
        self.downloader.open_spider(spider)
        yield self.scraper.open_spider(spider)
        stats.open_spider(spider)
        yield send_catch_log_deferred(signals.spider_opened, spider=spider)
        self.next_request(spider)

    def _spider_idle(self, spider):
        """Called when a spider gets idle. This function is called when there
        are no remaining pages to download or schedule. It can be called
        multiple times. If some extension raises a DontCloseSpider exception
        (in the spider_idle signal handler) the spider is not closed until the
        next loop and this function is guaranteed to be called (at least) once
        again for this spider.
        """
        res = send_catch_log(signal=signals.spider_idle, \
            spider=spider, dont_log=DontCloseSpider)
        if any(isinstance(x, Failure) and isinstance(x.value, DontCloseSpider) \
                for _, x in res):
            reactor.callLater(5, self.next_request, spider)
            return

        if self.spider_is_idle(spider):
            self.close_spider(spider, reason='finished')

    def close_spider(self, spider, reason='cancelled'):
        """Close (cancel) spider and clear all its outstanding requests"""

        slot = self.slots[spider]
        if slot.closing:
            return slot.closing
        log.msg("Closing spider (%s)" % reason, spider=spider)

        self.scheduler.clear_pending_requests(spider)

        dfd = slot.close()

        dfd.addBoth(lambda _: self.downloader.close_spider(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self.scraper.close_spider(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self.scheduler.close_spider(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self._cancel_next_call(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: send_catch_log_deferred(signal=signals.spider_closed, \
            spider=spider, reason=reason))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: stats.close_spider(spider, reason=reason))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: log.msg("Spider closed (%s)" % reason, spider=spider))

        dfd.addBoth(lambda _: self.slots.pop(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self._spider_closed_callback(spider))

        return dfd

    def _cancel_next_call(self, spider):
        call = self._next_request_calls.pop(spider, None)
        if call and call.active:
            call.cancel()

    def _close_all_spiders(self):
        dfds = [self.close_spider(s, reason='shutdown') for s in self.open_spiders]
        dlist = defer.DeferredList(dfds)
        return dlist

    @defer.inlineCallbacks
    def _finish_stopping_engine(self):
        yield send_catch_log_deferred(signal=signals.engine_stopped)
        yield stats.engine_stopped()
<commit_msg>fix issue with responses preventing spiders to be idle in engine counts<commit_after>"""
This is the Scrapy engine which controls the Scheduler, Downloader and Spiders.

For more information see docs/topics/architecture.rst

"""
from time import time

from twisted.internet import reactor, defer
from twisted.python.failure import Failure

from scrapy import log, signals
from scrapy.stats import stats
from scrapy.core.downloader import Downloader
from scrapy.core.scraper import Scraper
from scrapy.exceptions import DontCloseSpider
from scrapy.http import Response, Request
from scrapy.utils.misc import load_object
from scrapy.utils.signal import send_catch_log, send_catch_log_deferred
from scrapy.utils.defer import mustbe_deferred

class Slot(object):

    def __init__(self):
        self.closing = False
        self.inprogress = set() # requests in progress

    def add_request(self, request):
        self.inprogress.add(request)

    def remove_request(self, request):
        self.inprogress.remove(request)
        self._maybe_fire_closing()

    def close(self):
        self.closing = defer.Deferred()
        self._maybe_fire_closing()
        return self.closing

    def _maybe_fire_closing(self):
        if self.closing and not self.inprogress:
            self.closing.callback(None)

class ExecutionEngine(object):

    def __init__(self, settings, spider_closed_callback):
        self.settings = settings
        self.slots = {}
        self.running = False
        self.paused = False
        self._next_request_calls = {}
        self.scheduler = load_object(settings['SCHEDULER'])()
        self.downloader = Downloader()
        self.scraper = Scraper(self, self.settings)
        self._spider_closed_callback = spider_closed_callback

    @defer.inlineCallbacks
    def start(self):
        """Start the execution engine"""
        assert not self.running, "Engine already running"
        self.start_time = time()
        yield send_catch_log_deferred(signal=signals.engine_started)
        self.running = True

    def stop(self):
        """Stop the execution engine gracefully"""
        assert self.running, "Engine not running"
        self.running = False
        dfd = self._close_all_spiders()
        return dfd.addBoth(lambda _: self._finish_stopping_engine())

    def pause(self):
        """Pause the execution engine"""
        self.paused = True

    def unpause(self):
        """Resume the execution engine"""
        self.paused = False

    def is_idle(self):
        return self.scheduler.is_idle() and self.downloader.is_idle() and \
            self.scraper.is_idle()

    def next_request(self, spider, now=False):
        """Scrape the next request for the spider passed.

        The next request to be scraped is retrieved from the scheduler and
        requested from the downloader.

        The spider is closed if there are no more pages to scrape.
        """
        if now:
            self._next_request_calls.pop(spider, None)
        elif spider not in self._next_request_calls:
            call = reactor.callLater(0, self.next_request, spider, now=True)
            self._next_request_calls[spider] = call
            return call
        else:
            return

        if self.paused:
            return reactor.callLater(5, self.next_request, spider)

        while not self._needs_backout(spider):
            if not self._next_request(spider):
                break

        if self.spider_is_idle(spider):
            self._spider_idle(spider)

    def _needs_backout(self, spider):
        slot = self.slots[spider]
        return not self.running \
            or slot.closing \
            or self.spider_is_closed(spider) \
            or self.downloader.sites[spider].needs_backout() \
            or self.scraper.sites[spider].needs_backout()

    def _next_request(self, spider):
        request = self.scheduler.next_request(spider)
        if not request:
            return
        d = self._download(request, spider)
        d.addBoth(self._handle_downloader_output, request, spider)
        d.addErrback(log.msg, spider=spider)
        slot = self.slots[spider]
        d.addBoth(lambda _: slot.remove_request(request))
        d.addErrback(log.msg, spider=spider)
        d.addBoth(lambda _: self.next_request(spider))
        return d

    def _handle_downloader_output(self, response, request, spider):
        assert isinstance(response, (Request, Response, Failure)), response
        # downloader middleware can return requests (for example, redirects)
        if isinstance(response, Request):
            self.crawl(response, spider)
            return
        # response is a Response or Failure
        d = defer.Deferred()
        d.addBoth(self.scraper.enqueue_scrape, request, spider)
        d.addErrback(log.err, spider=spider)
        if isinstance(response, Failure):
            d.errback(response)
        else:
            d.callback(response)
        return d

    def spider_is_idle(self, spider):
        scraper_idle = spider in self.scraper.sites \
            and self.scraper.sites[spider].is_idle()
        pending = self.scheduler.spider_has_pending_requests(spider)
        downloading = spider in self.downloader.sites \
            and self.downloader.sites[spider].active
        return scraper_idle and not (pending or downloading)

    def spider_is_closed(self, spider):
        """Return True if the spider is fully closed (ie. not even in the
        closing stage)"""
        return spider not in self.downloader.sites

    @property
    def open_spiders(self):
        return self.downloader.sites.keys()

    def has_capacity(self):
        """Does the engine have capacity to handle more spiders"""
        return len(self.downloader.sites) < self.downloader.concurrent_spiders

    def crawl(self, request, spider):
        assert spider in self.open_spiders, \
            "Spider %r not opened when crawling: %s" % (spider.name, request)
        self.schedule(request, spider)
        self.next_request(spider)

    def schedule(self, request, spider):
        return self.scheduler.enqueue_request(spider, request)

    def download(self, request, spider):
        slot = self.slots[spider]
        slot.add_request(request)
        d = self._download(request, spider)
        d.addBoth(self._downloaded, slot, request, spider)
        return d

    def _downloaded(self, response, slot, request, spider):
        slot.remove_request(request)
        return self.download(response, spider) \
                if isinstance(response, Request) else response

    def _download(self, request, spider):
        slot = self.slots[spider]
        slot.add_request(request)
        def _on_success(response):
            """handle the result of a page download"""
            assert isinstance(response, (Response, Request))
            if isinstance(response, Response):
                response.request = request # tie request to response received
                log.msg(log.formatter.crawled(request, response, spider), \
                    level=log.DEBUG, spider=spider)
                send_catch_log(signal=signals.response_received, \
                    response=response, request=request, spider=spider)
            return response

        def _on_complete(_):
            self.next_request(spider)
            return _

        dwld = mustbe_deferred(self.downloader.fetch, request, spider)
        dwld.addCallback(_on_success)
        dwld.addBoth(_on_complete)
        return dwld

    @defer.inlineCallbacks
    def open_spider(self, spider):
        assert self.has_capacity(), "No free spider slots when opening %r" % \
            spider.name
        log.msg("Spider opened", spider=spider)
        self.slots[spider] = Slot()
        yield self.scheduler.open_spider(spider)
        self.downloader.open_spider(spider)
        yield self.scraper.open_spider(spider)
        stats.open_spider(spider)
        yield send_catch_log_deferred(signals.spider_opened, spider=spider)
        self.next_request(spider)

    def _spider_idle(self, spider):
        """Called when a spider gets idle. This function is called when there
        are no remaining pages to download or schedule. It can be called
        multiple times. If some extension raises a DontCloseSpider exception
        (in the spider_idle signal handler) the spider is not closed until the
        next loop and this function is guaranteed to be called (at least) once
        again for this spider.
        """
        res = send_catch_log(signal=signals.spider_idle, \
            spider=spider, dont_log=DontCloseSpider)
        if any(isinstance(x, Failure) and isinstance(x.value, DontCloseSpider) \
                for _, x in res):
            reactor.callLater(5, self.next_request, spider)
            return

        if self.spider_is_idle(spider):
            self.close_spider(spider, reason='finished')

    def close_spider(self, spider, reason='cancelled'):
        """Close (cancel) spider and clear all its outstanding requests"""

        slot = self.slots[spider]
        if slot.closing:
            return slot.closing
        log.msg("Closing spider (%s)" % reason, spider=spider)

        self.scheduler.clear_pending_requests(spider)

        dfd = slot.close()

        dfd.addBoth(lambda _: self.downloader.close_spider(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self.scraper.close_spider(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self.scheduler.close_spider(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self._cancel_next_call(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: send_catch_log_deferred(signal=signals.spider_closed, \
            spider=spider, reason=reason))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: stats.close_spider(spider, reason=reason))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: log.msg("Spider closed (%s)" % reason, spider=spider))

        dfd.addBoth(lambda _: self.slots.pop(spider))
        dfd.addErrback(log.err, spider=spider)

        dfd.addBoth(lambda _: self._spider_closed_callback(spider))

        return dfd

    def _cancel_next_call(self, spider):
        call = self._next_request_calls.pop(spider, None)
        if call and call.active:
            call.cancel()

    def _close_all_spiders(self):
        dfds = [self.close_spider(s, reason='shutdown') for s in self.open_spiders]
        dlist = defer.DeferredList(dfds)
        return dlist

    @defer.inlineCallbacks
    def _finish_stopping_engine(self):
        yield send_catch_log_deferred(signal=signals.engine_stopped)
        yield stats.engine_stopped()
<|endoftext|>


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

import re
import time
import logging
from copy import copy

from wlauto.exceptions import ConfigError
from wlauto.utils.serial_port import TIMEOUT
from wlauto.utils.types import boolean

logger = logging.getLogger('UEFI')

class UefiConfig(object):

    def __init__(self, config_dict):
        if isinstance(config_dict, UefiConfig):
            self.__dict__ = copy(config_dict.__dict__)
        else:
            try:
                self.image_name = config_dict['image_name']
                self.image_args = config_dict['image_args']
                self.fdt_support = boolean(config_dict['fdt_support'])
            except KeyError as e:
                raise ConfigError('Missing mandatory parameter for UEFI entry config: "{}"'.format(e))
            self.initrd = config_dict.get('initrd')
            self.fdt_path = config_dict.get('fdt_path')
            if self.fdt_path and not self.fdt_support:
                raise ConfigError('FDT path has been specfied for UEFI entry, when FDT support is "False"')

class UefiMenu(object):
    """
    Allows navigating UEFI menu over serial (it relies on a pexpect connection).

    """

    option_regex = re.compile(r'^\[(\d+)\]\s+([^\r]+)\r\n', re.M)
    prompt_regex = re.compile(r'^([^\r\n]+):\s*', re.M)
    invalid_regex = re.compile(r'Invalid input \(max (\d+)\)', re.M)

    load_delay = 1  # seconds
    default_timeout = 60  # seconds

    def __init__(self, conn, prompt='The default boot selection will start in'):
        """
        :param conn: A serial connection as returned by ``pexect.spawn()``.
        :param prompt: The starting prompt to wait for during ``open()``.

        """
        self.conn = conn
        self.start_prompt = prompt
        self.options = {}
        self.prompt = None

    def open(self, timeout=default_timeout):
        """
        "Open" the UEFI menu by sending an interrupt on STDIN after seeing the
        starting prompt (configurable upon creation of the ``UefiMenu`` object.

        """
        self.conn.expect(self.start_prompt, timeout)
        self.conn.sendline('')
        time.sleep(self.load_delay)

    def create_entry(self, name, config):
        """Create a new UEFI entry using the parameters. The menu is assumed
        to be at the top level. Upon return, the menu will be at the top level."""
        logger.debug('Creating UEFI entry {}'.format(name))
        self.nudge()
        self.select('Boot Manager')
        self.select('Add Boot Device Entry')
        self.select('NOR Flash')
        self.enter(config.image_name)
        self.enter('y' if config.fdt_support else 'n')
        if config.initrd:
            self.enter('y')
            self.enter(config.initrd)
        else:
            self.enter('n')
        self.enter(config.image_args)
        self.enter(name)

        if config.fdt_path:
            self.select('Update FDT path')
            self.enter(config.fdt_path)

        self.select('Return to main menu')

    def delete_entry(self, name):
        """Delete the specified UEFI entry. The menu is assumed
        to be at the top level. Upon return, the menu will be at the top level."""
        logger.debug('Removing UEFI entry {}'.format(name))
        self.nudge()
        self.select('Boot Manager')
        self.select('Remove Boot Device Entry')
        self.select(name)
        self.select('Return to main menu')

    def select(self, option, timeout=default_timeout):
        """
        Select the specified option from the current menu.

        :param option: Could be an ``int`` index of the option, or a string/regex to
                       match option text against.
        :param timeout: If a non-``int`` option is specified, the option list may need
                        need to be parsed (if it hasn't been already), this may block
                        and the timeout is used to cap that , resulting in a ``TIMEOUT``
                        exception.
        :param delay: A fixed delay to wait after sending the input to the serial connection.
                      This should be set if input this action is known to result in a
                      long-running operation.

        """
        if isinstance(option, basestring):
            option = self.get_option_index(option, timeout)
        self.enter(option)

    def enter(self, value, delay=load_delay):
        """Like ``select()`` except no resolution is performed -- the value is sent directly
        to the serial connection."""
        # Empty the buffer first, so that only response to the input about to
        # be sent will be processed by subsequent commands.
        value = str(value)
        self._reset()
        self.write_characters(value)
        # TODO: in case the value is long an complicated, things may get
        # screwed up (e.g. there may be line breaks injected), additionally,
        # special chars might cause regex to fail. To avoid these issues i'm
        # only matching against the first 5 chars of the value. This is
        # entirely arbitrary and I'll probably have to find a better way of
        # doing this at some point.
        self.conn.expect(value[:5], timeout=delay)
        time.sleep(self.load_delay)

    def read_menu(self, timeout=default_timeout):
        """Parse serial output to get the menu options and the following prompt."""
        attempting_timeout_retry = False
        attempting_invalid_retry = False
        while True:
            index = self.conn.expect([self.option_regex, self.prompt_regex, self.invalid_regex, TIMEOUT],
                                     timeout=timeout)
            match = self.conn.match
            if index == 0:  # matched menu option
                self.options[match.group(1)] = match.group(2)
            elif index == 1:  # matched prompt
                self.prompt = match.group(1)
                break
            elif index == 2:  # matched invalid selection
                # We've sent an invalid input (which includes an empty line) at
                # the top-level menu. To get back the menu options, it seems we
                # need to enter what the error reports as the max + 1, so...
                if not attempting_invalid_retry:
                    attempting_invalid_retry = True
                    val = int(match.group(1)) + 1
                    self.empty_buffer()
                    self.enter(val)
                else:   # OK, that didn't work; panic!
                    raise RuntimeError('Could not read menu entries stuck on "{}" prompt'.format(self.prompt))
            elif index == 3:  # timed out
                if not attempting_timeout_retry:
                    attempting_timeout_retry = True
                    self.nudge()
                else:  # Didn't help. Run away!
                    raise RuntimeError('Did not see a valid UEFI menu.')
            else:
                raise AssertionError('Unexpected response waiting for UEFI menu')  # should never get here

    def get_option_index(self, text, timeout=default_timeout):
        """Returns the menu index of the specified option text (uses regex matching). If the option
        is not in the current menu, ``LookupError`` will be raised."""
        if not self.prompt:
            self.read_menu(timeout)
        for k, v in self.options.iteritems():
            if re.search(text, v):
                return k
        raise LookupError(text)

    def has_option(self, text, timeout=default_timeout):
        """Returns ``True`` if at least one of the options in the current menu has
        matched (using regex) the specified text."""
        try:
            self.get_option_index(text, timeout)
            return True
        except LookupError:
            return False

    def nudge(self):
        """Send a little nudge to ensure there is something to read. This is useful when you're not
        sure if all out put from the serial has been read already."""
        self.enter('')

    def empty_buffer(self):
        """Read everything from the serial and clear the internal pexpect buffer. This ensures
        that the next ``expect()`` call will time out (unless further input will be sent to the
        serial beforehand. This is used to create a "known" state and avoid unexpected matches."""
        try:
            while True:
                time.sleep(0.1)
                self.conn.read_nonblocking(size=1024, timeout=0.1)
        except TIMEOUT:
            pass
        self.conn.buffer = ''

    def write_characters(self, line):
        """Write a single line out to serial charcter-by-character. This will ensure that nothing will
        be dropped for longer lines."""
        line = line.rstrip('\r\n')
        for c in line:
            self.conn.send(c)
            time.sleep(0.05)
        self.conn.sendline('')

    def _reset(self):
        self.options = {}
        self.prompt = None
        self.empty_buffer()

<commit_msg>uefi menu: updated prompt_regex to handle debug UEFI builds.<commit_after>#    Copyright 2014-2015 ARM Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import re
import time
import logging
from copy import copy

from wlauto.exceptions import ConfigError
from wlauto.utils.serial_port import TIMEOUT
from wlauto.utils.types import boolean

logger = logging.getLogger('UEFI')

class UefiConfig(object):

    def __init__(self, config_dict):
        if isinstance(config_dict, UefiConfig):
            self.__dict__ = copy(config_dict.__dict__)
        else:
            try:
                self.image_name = config_dict['image_name']
                self.image_args = config_dict['image_args']
                self.fdt_support = boolean(config_dict['fdt_support'])
            except KeyError as e:
                raise ConfigError('Missing mandatory parameter for UEFI entry config: "{}"'.format(e))
            self.initrd = config_dict.get('initrd')
            self.fdt_path = config_dict.get('fdt_path')
            if self.fdt_path and not self.fdt_support:
                raise ConfigError('FDT path has been specfied for UEFI entry, when FDT support is "False"')

class UefiMenu(object):
    """
    Allows navigating UEFI menu over serial (it relies on a pexpect connection).

    """

    option_regex = re.compile(r'^\[(\d+)\]\s+([^\r]+)\r\n', re.M)
    prompt_regex = re.compile(r'^(\S[^\r\n]+):\s*', re.M)
    invalid_regex = re.compile(r'Invalid input \(max (\d+)\)', re.M)

    load_delay = 1  # seconds
    default_timeout = 60  # seconds

    def __init__(self, conn, prompt='The default boot selection will start in'):
        """
        :param conn: A serial connection as returned by ``pexect.spawn()``.
        :param prompt: The starting prompt to wait for during ``open()``.

        """
        self.conn = conn
        self.start_prompt = prompt
        self.options = {}
        self.prompt = None

    def open(self, timeout=default_timeout):
        """
        "Open" the UEFI menu by sending an interrupt on STDIN after seeing the
        starting prompt (configurable upon creation of the ``UefiMenu`` object.

        """
        self.conn.expect(self.start_prompt, timeout)
        self.conn.sendline('')
        time.sleep(self.load_delay)

    def create_entry(self, name, config):
        """Create a new UEFI entry using the parameters. The menu is assumed
        to be at the top level. Upon return, the menu will be at the top level."""
        logger.debug('Creating UEFI entry {}'.format(name))
        self.nudge()
        self.select('Boot Manager')
        self.select('Add Boot Device Entry')
        self.select('NOR Flash')
        self.enter(config.image_name)
        self.enter('y' if config.fdt_support else 'n')
        if config.initrd:
            self.enter('y')
            self.enter(config.initrd)
        else:
            self.enter('n')
        self.enter(config.image_args)
        self.enter(name)

        if config.fdt_path:
            self.select('Update FDT path')
            self.enter(config.fdt_path)

        self.select('Return to main menu')

    def delete_entry(self, name):
        """Delete the specified UEFI entry. The menu is assumed
        to be at the top level. Upon return, the menu will be at the top level."""
        logger.debug('Removing UEFI entry {}'.format(name))
        self.nudge()
        self.select('Boot Manager')
        self.select('Remove Boot Device Entry')
        self.select(name)
        self.select('Return to main menu')

    def select(self, option, timeout=default_timeout):
        """
        Select the specified option from the current menu.

        :param option: Could be an ``int`` index of the option, or a string/regex to
                       match option text against.
        :param timeout: If a non-``int`` option is specified, the option list may need
                        need to be parsed (if it hasn't been already), this may block
                        and the timeout is used to cap that , resulting in a ``TIMEOUT``
                        exception.
        :param delay: A fixed delay to wait after sending the input to the serial connection.
                      This should be set if input this action is known to result in a
                      long-running operation.

        """
        if isinstance(option, basestring):
            option = self.get_option_index(option, timeout)
        self.enter(option)

    def enter(self, value, delay=load_delay):
        """Like ``select()`` except no resolution is performed -- the value is sent directly
        to the serial connection."""
        # Empty the buffer first, so that only response to the input about to
        # be sent will be processed by subsequent commands.
        value = str(value)
        self._reset()
        self.write_characters(value)
        # TODO: in case the value is long an complicated, things may get
        # screwed up (e.g. there may be line breaks injected), additionally,
        # special chars might cause regex to fail. To avoid these issues i'm
        # only matching against the first 5 chars of the value. This is
        # entirely arbitrary and I'll probably have to find a better way of
        # doing this at some point.
        self.conn.expect(value[:5], timeout=delay)
        time.sleep(self.load_delay)

    def read_menu(self, timeout=default_timeout):
        """Parse serial output to get the menu options and the following prompt."""
        attempting_timeout_retry = False
        attempting_invalid_retry = False
        while True:
            index = self.conn.expect([self.option_regex, self.prompt_regex, self.invalid_regex, TIMEOUT],
                                     timeout=timeout)
            match = self.conn.match
            if index == 0:  # matched menu option
                self.options[match.group(1)] = match.group(2)
            elif index == 1:  # matched prompt
                self.prompt = match.group(1)
                break
            elif index == 2:  # matched invalid selection
                # We've sent an invalid input (which includes an empty line) at
                # the top-level menu. To get back the menu options, it seems we
                # need to enter what the error reports as the max + 1, so...
                if not attempting_invalid_retry:
                    attempting_invalid_retry = True
                    val = int(match.group(1)) + 1
                    self.empty_buffer()
                    self.enter(val)
                else:   # OK, that didn't work; panic!
                    raise RuntimeError('Could not read menu entries stuck on "{}" prompt'.format(self.prompt))
            elif index == 3:  # timed out
                if not attempting_timeout_retry:
                    attempting_timeout_retry = True
                    self.nudge()
                else:  # Didn't help. Run away!
                    raise RuntimeError('Did not see a valid UEFI menu.')
            else:
                raise AssertionError('Unexpected response waiting for UEFI menu')  # should never get here

    def get_option_index(self, text, timeout=default_timeout):
        """Returns the menu index of the specified option text (uses regex matching). If the option
        is not in the current menu, ``LookupError`` will be raised."""
        if not self.prompt:
            self.read_menu(timeout)
        for k, v in self.options.iteritems():
            if re.search(text, v):
                return k
        raise LookupError(text)

    def has_option(self, text, timeout=default_timeout):
        """Returns ``True`` if at least one of the options in the current menu has
        matched (using regex) the specified text."""
        try:
            self.get_option_index(text, timeout)
            return True
        except LookupError:
            return False

    def nudge(self):
        """Send a little nudge to ensure there is something to read. This is useful when you're not
        sure if all out put from the serial has been read already."""
        self.enter('')

    def empty_buffer(self):
        """Read everything from the serial and clear the internal pexpect buffer. This ensures
        that the next ``expect()`` call will time out (unless further input will be sent to the
        serial beforehand. This is used to create a "known" state and avoid unexpected matches."""
        try:
            while True:
                time.sleep(0.1)
                self.conn.read_nonblocking(size=1024, timeout=0.1)
        except TIMEOUT:
            pass
        self.conn.buffer = ''

    def write_characters(self, line):
        """Write a single line out to serial charcter-by-character. This will ensure that nothing will
        be dropped for longer lines."""
        line = line.rstrip('\r\n')
        for c in line:
            self.conn.send(c)
            time.sleep(0.05)
        self.conn.sendline('')

    def _reset(self):
        self.options = {}
        self.prompt = None
        self.empty_buffer()

<|endoftext|>