Source code for pybel.struct.mutation.induction.neighborhood

# -*- coding: utf-8 -*-

"""Functions for selecting by the neighborhoods of nodes."""

import itertools as itt
from typing import Iterable, Optional

from ...graph import BELGraph
from ...pipeline import transformation
from ...utils import update_metadata
from ....dsl import BaseEntity

__all__ = [
    "get_subgraph_by_neighborhood",
]


[docs]@transformation def get_subgraph_by_neighborhood(graph: BELGraph, nodes: Iterable[BaseEntity]) -> Optional[BELGraph]: """Get a BEL graph around the neighborhoods of the given nodes. Returns none if no nodes are in the graph. :param graph: A BEL graph :param nodes: An iterable of BEL nodes :return: A BEL graph induced around the neighborhoods of the given nodes """ node_set = set(nodes) if not any(node in graph for node in node_set): return rv = graph.child() rv.add_edges_from( itt.chain( graph.in_edges(nodes, keys=True, data=True), graph.out_edges(nodes, keys=True, data=True), ), ) return rv