restricted paths

restricted_paths.py

import networkx as nx
from contextlib import contextmanager

@contextmanager
def hidden(G, nodes=None, edges=None):
    if nodes is None:
        nodes = []
    if edges is None:
        edges = []
    hidden_nodes = [(u,G.node[u]) for u in nodes]
    # add edges adjacent to any hidden (to be removed) nodes
    extra_edges = G.edges(nodes)
    hidden_edges = [(u,v,G[u][v]) for u,v in edges+extra_edges]
    try:
        G.remove_edges_from(hidden_edges)
        # Should remove_nodes_from be modified so we can just use
        # G.remove_nodes_from(nodes) ?
        G.remove_nodes_from(n for (n,d) in hidden_nodes)
        yield G
    finally:
        G.add_edges_from(hidden_edges)
        G.add_nodes_from(hidden_nodes)

G = nx.Graph()
G.add_path([1,2,3,4])
G.add_path([1,10,4])

print nx.shortest_path(G, 1,4)

with hidden(G, edges=[(1,10)]) as H:
    print nx.shortest_path(H, 1, 4)

print nx.shortest_path(G, 1, 4)

with hidden(G, nodes=[10]) as H:
    print nx.shortest_path(H, 1, 4)

print nx.shortest_path(G, 1, 4)