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

"""

This module assists in running complex workflows on BEL graphs.

Example Pipeline #1
~~~~~~~~~~~~~~~~~~~
This example shows a pipeline that acquires a subgraph and finds possible additions to the subgraph.

>>> network = ...
>>> example = Pipeline()
>>> example.append('get_subgraph_by_annotation_value', 'Subgraph', 'Blood vessel dilation subgraph')
>>> example.append('enrich_unqualified')
>>> example.append('infer_central_dogma')
>>> example.append('expand_periphery')
>>> result = example.run(network)

Example Pipeline #2
~~~~~~~~~~~~~~~~~~~
This example shows how additional data can be integrated into a graph.

>>> from pybel.constants import PROTEIN
>>> network = ...
>>> example = Pipeline()
>>> example.append('infer_central_dogma')
>>> example.append('enrich_unqualified')
>>> example.append('infer_central_dogma')
>>> example.append('expand_periphery')
>>> example.append('expand_nodes_neighborhoods', [(PROTEIN, 'HGNC', 'AKT1'), (PROTEIN, 'HGNC', 'AKT2')])
>>> result = example.run(network)

Example Pipeline #3
~~~~~~~~~~~~~~~~~~~
This example shows how the results from multiple pipelines can be combine.

>>> network = ...
>>> pipeline_a = Pipeline()
>>> pipeline_a.append('get_subgraph_by_annotation_value', 'Subgraph', 'Blood vessel dilation subgraph')
>>> pipeline_b = Pipeline()
>>> pipeline_b.append('get_subgraph_by_annotation_value', 'Subgraph', 'Tau protein subgraph')
>>> pipeline_c = Pipeline.union(pipeline_a, pipeline_b)
>>> result = pipeline_c.run(network)

"""

from __future__ import print_function

import json
import logging
import types
from functools import wraps
from inspect import signature

from pybel.struct.operations import node_intersection, union

__all__ = [
    'Pipeline',
    'in_place_mutator',
    'uni_in_place_mutator',
    'uni_mutator',
    'mutator',
    'splitter'
]

log = logging.getLogger(__name__)

META_UNION = 'union'
META_INTERSECTION = 'intersection'

mapped = {}
universe_map = {}
no_universe_map = {}
in_place_map = {}
out_place_map = {}
has_arguments_map = {}
no_arguments_map = {}

#: A map of function names to functions that split graphs into dictionaries of graphs
splitter_map = {}


class MissingPipelineFunctionError(KeyError):
    """Raised when trying to run the pipeline with a function that isn't registered"""


def _register(universe, in_place, **kwargs):
    """Builds a decorator function to tag mutator functions

    :param bool universe: Does the first positional argument of this function correspond to a universe graph?
    :param bool in_place: Does this function return a new graph, or just modify it in-place?
    :param dict kwargs: Other parameters to annotate to the function
    """

    def decorator(f):
        """A decorator that tags mutator functions

        :param f: A function
        :return: The same function, with additional properties added
        """
        f.__dict__.update(kwargs)

        if universe:
            universe_map[f.__name__] = f
        else:
            no_universe_map[f.__name__] = f

        if in_place:
            in_place_map[f.__name__] = f
        else:
            out_place_map[f.__name__] = f

        z = signature(f)
        if (universe and 3 <= len(z.parameters)) or (not universe and 2 <= len(z.parameters)):
            has_arguments_map[f.__name__] = f
        else:
            no_arguments_map[f.__name__] = f

        mapped[f.__name__] = f

        return f

    return decorator


#: A function decorator to inform the Pipeline how to handle a function
in_place_mutator = _register(universe=False, in_place=True)
#: A function decorator to inform the Pipeline how to handle a function
uni_in_place_mutator = _register(universe=True, in_place=True)
#: A function decorator to inform the Pipeline how to handle a function
uni_mutator = _register(universe=True, in_place=False)
#: A function decorator to inform the Pipeline how to handle a function
mutator = _register(universe=False, in_place=False)

SET_UNIVERSE = 'UNIVERSE'


def splitter(f):
    """A function decorator that signifies a function that takes in a graph and returns a dictionary of keys to graphs

    :param types.FunctionType f: A function
    :rtype: types.FunctionType
    """
    splitter_map[f.__name__] = f
    return f


def function_is_registered(name):
    """Checks if a function is a valid pipeline function

    :param str or types.FunctionType name: The name of the function
    :rtype: bool
    """
    if not isinstance(name, str):
        return name.__name__ in mapped

    return name in mapped


class Pipeline:
    """Builds and runs analytical pipelines on BEL graphs"""

    def __init__(self, protocol=None, universe=None):
        """
        :param list[dict] protocol: The list of dictionaries describing how to mutate/filter a network
        :param pybel.BELGraph universe: The entire set of known knowledge to draw from
        """
        self.universe = universe
        self.protocol = [] if protocol is None else protocol

    @staticmethod
    def from_functions(functions):
        """Builds a pipeline from a list of functions

        :param iter[(pybel.BELGraph) -> pybel.BELGraph or (pybel.BELGraph) -> None] functions: A list of functions
        :rtype: Pipeline
        """
        result = Pipeline()
        for func in functions:
            result.append(func)
        return result

    def get_function(self, name):
        """Wraps a function with the universe and in-place

        :param str name: The name of the function
        :rtype: types.FunctionType
        """
        f = mapped.get(name)
        if f is None:
            raise MissingPipelineFunctionError('{} is not registered as a pipeline function'.format(name))

        if name in universe_map and name in in_place_map:
            return self.wrap_in_place(self.wrap_universe(f))

        if name in universe_map:
            return self.wrap_universe(f)

        if name in in_place_map:
            return self.wrap_in_place(f)

        return f

    def append(self, name, *args, **kwargs):
        """Adds a function (either as a reference, or by name) and arguments to the pipeline

        :param name: The name of the function
        :type name: str or (pybel.BELGraph -> pybel.BELGraph)
        :param args: The positional arguments to call in the function
        :param kwargs: The keyword arguments to call in the function
        :return: This pipeline for fluid query building
        :rtype: Pipeline
        """
        if isinstance(name, types.FunctionType):
            return self.append(name.__name__, *args, **kwargs)
        elif not isinstance(name, str):
            raise TypeError('invalid function argument: {}'.format(name))

        if not function_is_registered(name):
            raise KeyError(name)

        av = {
            'function': name,
        }

        if args:
            av['args'] = args

        if kwargs:
            av['kwargs'] = kwargs

        self.protocol.append(av)
        return self
<commit_msg>Better error handling for missing pipeline functions<commit_after># -*- coding: utf-8 -*-

"""

This module assists in running complex workflows on BEL graphs.

Example Pipeline #1
~~~~~~~~~~~~~~~~~~~
This example shows a pipeline that acquires a subgraph and finds possible additions to the subgraph.

>>> network = ...
>>> example = Pipeline()
>>> example.append('get_subgraph_by_annotation_value', 'Subgraph', 'Blood vessel dilation subgraph')
>>> example.append('enrich_unqualified')
>>> example.append('infer_central_dogma')
>>> example.append('expand_periphery')
>>> result = example.run(network)

Example Pipeline #2
~~~~~~~~~~~~~~~~~~~
This example shows how additional data can be integrated into a graph.

>>> from pybel.constants import PROTEIN
>>> network = ...
>>> example = Pipeline()
>>> example.append('infer_central_dogma')
>>> example.append('enrich_unqualified')
>>> example.append('infer_central_dogma')
>>> example.append('expand_periphery')
>>> example.append('expand_nodes_neighborhoods', [(PROTEIN, 'HGNC', 'AKT1'), (PROTEIN, 'HGNC', 'AKT2')])
>>> result = example.run(network)

Example Pipeline #3
~~~~~~~~~~~~~~~~~~~
This example shows how the results from multiple pipelines can be combine.

>>> network = ...
>>> pipeline_a = Pipeline()
>>> pipeline_a.append('get_subgraph_by_annotation_value', 'Subgraph', 'Blood vessel dilation subgraph')
>>> pipeline_b = Pipeline()
>>> pipeline_b.append('get_subgraph_by_annotation_value', 'Subgraph', 'Tau protein subgraph')
>>> pipeline_c = Pipeline.union(pipeline_a, pipeline_b)
>>> result = pipeline_c.run(network)

"""

from __future__ import print_function

import json
import logging
import types
from functools import wraps
from inspect import signature

from pybel.struct.operations import node_intersection, union

__all__ = [
    'Pipeline',
    'in_place_mutator',
    'uni_in_place_mutator',
    'uni_mutator',
    'mutator',
    'splitter'
]

log = logging.getLogger(__name__)

META_UNION = 'union'
META_INTERSECTION = 'intersection'

mapped = {}
universe_map = {}
no_universe_map = {}
in_place_map = {}
out_place_map = {}
has_arguments_map = {}
no_arguments_map = {}

#: A map of function names to functions that split graphs into dictionaries of graphs
splitter_map = {}


class MissingPipelineFunctionError(KeyError):
    """Raised when trying to run the pipeline with a function that isn't registered"""


def _register(universe, in_place, **kwargs):
    """Builds a decorator function to tag mutator functions

    :param bool universe: Does the first positional argument of this function correspond to a universe graph?
    :param bool in_place: Does this function return a new graph, or just modify it in-place?
    :param dict kwargs: Other parameters to annotate to the function
    """

    def decorator(f):
        """A decorator that tags mutator functions

        :param f: A function
        :return: The same function, with additional properties added
        """
        f.__dict__.update(kwargs)

        if universe:
            universe_map[f.__name__] = f
        else:
            no_universe_map[f.__name__] = f

        if in_place:
            in_place_map[f.__name__] = f
        else:
            out_place_map[f.__name__] = f

        z = signature(f)
        if (universe and 3 <= len(z.parameters)) or (not universe and 2 <= len(z.parameters)):
            has_arguments_map[f.__name__] = f
        else:
            no_arguments_map[f.__name__] = f

        mapped[f.__name__] = f

        return f

    return decorator


#: A function decorator to inform the Pipeline how to handle a function
in_place_mutator = _register(universe=False, in_place=True)
#: A function decorator to inform the Pipeline how to handle a function
uni_in_place_mutator = _register(universe=True, in_place=True)
#: A function decorator to inform the Pipeline how to handle a function
uni_mutator = _register(universe=True, in_place=False)
#: A function decorator to inform the Pipeline how to handle a function
mutator = _register(universe=False, in_place=False)

SET_UNIVERSE = 'UNIVERSE'


def assert_is_mapped_to_pipeline(name):
    """
    :param str name:
    :raises: MissingPipelineFunctionError
    """
    if name not in mapped:
        raise MissingPipelineFunctionError('{} is not registered as a pipeline function'.format(name))


def splitter(f):
    """A function decorator that signifies a function that takes in a graph and returns a dictionary of keys to graphs

    :param types.FunctionType f: A function
    :rtype: types.FunctionType
    """
    splitter_map[f.__name__] = f
    return f


def function_is_registered(name):
    """Checks if a function is a valid pipeline function

    :param str or types.FunctionType name: The name of the function
    :rtype: bool
    """
    if not isinstance(name, str):
        return name.__name__ in mapped

    return name in mapped


class Pipeline:
    """Builds and runs analytical pipelines on BEL graphs"""

    def __init__(self, protocol=None, universe=None):
        """
        :param list[dict] protocol: The list of dictionaries describing how to mutate/filter a network
        :param pybel.BELGraph universe: The entire set of known knowledge to draw from
        """
        self.universe = universe
        self.protocol = [] if protocol is None else protocol

    @staticmethod
    def from_functions(functions):
        """Builds a pipeline from a list of functions

        :param iter[(pybel.BELGraph) -> pybel.BELGraph or (pybel.BELGraph) -> None] functions: A list of functions
        :rtype: Pipeline
        """
        result = Pipeline()
        for func in functions:
            result.append(func)
        return result

    def get_function(self, name):
        """Wraps a function with the universe and in-place

        :param str name: The name of the function
        :rtype: types.FunctionType
        :raises: MissingPipelineFunctionError
        """
        assert_is_mapped_to_pipeline(name)

        f = mapped[name]

        if name in universe_map and name in in_place_map:
            return self.wrap_in_place(self.wrap_universe(f))

        if name in universe_map:
            return self.wrap_universe(f)

        if name in in_place_map:
            return self.wrap_in_place(f)

        return f

    def append(self, name, *args, **kwargs):
        """Adds a function (either as a reference, or by name) and arguments to the pipeline

        :param name: The name of the function
        :type name: str or (pybel.BELGraph -> pybel.BELGraph)
        :param args: The positional arguments to call in the function
        :param kwargs: The keyword arguments to call in the function
        :return: This pipeline for fluid query building
        :rtype: Pipeline
        :raises: MissingPipelineFunctionError
        """
        if isinstance(name, types.FunctionType):
            return self.append(name.__name__, *args, **kwargs)
        elif isinstance(name, str):
            assert_is_mapped_to_pipeline(name)
        else:
            raise TypeError('invalid function argument: {}'.format(name))

        if not function_is_registered(name):
            raise KeyError(name)

        av = {
            'function': name,
        }

        if args:
            av['args'] = args

        if kwargs:
            av['kwargs'] = kwargs

        self.protocol.append(av)
        return self

{'fraction_non_alphanumeric': 0.10185569865669575, 'fraction_numerical': 0.001107879795042238, 'mean_word_length': 3.1778999132195547, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 147, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6858830', 'n_tokens_mistral': 4184, 'n_tokens_neox': 4026, 'n_words': 1360}