def country_view(multi_graph: nx.MultiDiGraph, client_country: str) -> nx.MultiDiGraph: """Returns a view of the edges restricted to the given client country.""" country_edges = [(u, v, k) for u, v, k, measurement in multi_graph.edges( keys=True, data="measurement") if measurement.client_country == client_country] return multi_graph.edge_subgraph(country_edges)
def remove_edges(source: nx.MultiDiGraph, predicate: Callable[[Any, Any, Dict[str, Any]], bool]): """ Returns a subgraph of source that does not contain the edges for which the predicate returns true. Args: source: source graph predicate: a function(u, v, attr) that returns true if the edge from node u to node v with the attributes attr should be removed. Returns: the subgraph of source induced by the edges that are not selected by the predicate. This is a read-only view, you may want to use copy() on the result. """ to_keep = [(u, v, k) for u, v, k, attr in source.edges(data=True, keys=True) if not predicate(u, v, attr)] return source.edge_subgraph(to_keep)